Blame view

embeddings/text_encoder.py 10.5 KB
be52af70   tangwang   first commit
1
  """
325eec03   tangwang   1. 日志、配置基础设施,使用优化
2
  Text embedding encoder using network service.
be52af70   tangwang   first commit
3
  
7bfb9946   tangwang   向量化模块
4
  Generates embeddings via HTTP API service (default localhost:6005).
be52af70   tangwang   first commit
5
6
7
  """
  
  import sys
325eec03   tangwang   1. 日志、配置基础设施,使用优化
8
  import requests
be52af70   tangwang   first commit
9
10
  import time
  import threading
be52af70   tangwang   first commit
11
  import numpy as np
453992a8   tangwang   需求:
12
13
  import pickle
  import redis
7bfb9946   tangwang   向量化模块
14
  import os
453992a8   tangwang   需求:
15
16
  from datetime import timedelta
  from typing import List, Union, Dict, Any, Optional
325eec03   tangwang   1. 日志、配置基础设施,使用优化
17
  import logging
325eec03   tangwang   1. 日志、配置基础设施,使用优化
18
19
  
  logger = logging.getLogger(__name__)
be52af70   tangwang   first commit
20
  
453992a8   tangwang   需求:
21
22
23
24
25
26
  # Try to import REDIS_CONFIG, but allow import to fail
  try:
      from config.env_config import REDIS_CONFIG
  except ImportError:
      REDIS_CONFIG = {}
  
be52af70   tangwang   first commit
27
28
29
  
  class BgeEncoder:
      """
325eec03   tangwang   1. 日志、配置基础设施,使用优化
30
      Singleton text encoder using network service.
be52af70   tangwang   first commit
31
  
325eec03   tangwang   1. 日志、配置基础设施,使用优化
32
      Thread-safe singleton pattern ensures only one instance exists.
be52af70   tangwang   first commit
33
34
35
36
      """
      _instance = None
      _lock = threading.Lock()
  
7bfb9946   tangwang   向量化模块
37
      def __new__(cls, service_url: Optional[str] = None):
be52af70   tangwang   first commit
38
39
40
          with cls._lock:
              if cls._instance is None:
                  cls._instance = super(BgeEncoder, cls).__new__(cls)
7bfb9946   tangwang   向量化模块
41
42
43
44
                  resolved_url = service_url or os.getenv("EMBEDDING_SERVICE_URL", "http://localhost:6005")
                  logger.info(f"Creating BgeEncoder instance with service URL: {resolved_url}")
                  cls._instance.service_url = resolved_url
                  cls._instance.endpoint = f"{resolved_url}/embed/text"
453992a8   tangwang   需求:
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
                  
                  # Initialize Redis cache
                  try:
                      cls._instance.redis_client = redis.Redis(
                          host=REDIS_CONFIG.get('host', 'localhost'),
                          port=REDIS_CONFIG.get('port', 6479),
                          password=REDIS_CONFIG.get('password'),
                          decode_responses=False,  # Keep binary data as is
                          socket_timeout=REDIS_CONFIG.get('socket_timeout', 1),
                          socket_connect_timeout=REDIS_CONFIG.get('socket_connect_timeout', 1),
                          retry_on_timeout=REDIS_CONFIG.get('retry_on_timeout', False),
                          health_check_interval=10  # 避免复用坏连接
                      )
                      # Test connection
                      cls._instance.redis_client.ping()
                      cls._instance.expire_time = timedelta(days=REDIS_CONFIG.get('cache_expire_days', 180))
                      logger.info("Redis cache initialized for embeddings")
                  except Exception as e:
                      logger.warning(f"Failed to initialize Redis cache for embeddings: {e}, continuing without cache")
                      cls._instance.redis_client = None
be52af70   tangwang   first commit
65
66
          return cls._instance
  
7bfb9946   tangwang   向量化模块
67
      def _call_service(self, request_data: List[str]) -> List[Any]:
325eec03   tangwang   1. 日志、配置基础设施,使用优化
68
69
70
71
          """
          Call the embedding service API.
  
          Args:
7bfb9946   tangwang   向量化模块
72
              request_data: List of texts
325eec03   tangwang   1. 日志、配置基础设施,使用优化
73
74
  
          Returns:
7bfb9946   tangwang   向量化模块
75
              List of embeddings (list[float]) or nulls (None), aligned to input order
325eec03   tangwang   1. 日志、配置基础设施,使用优化
76
77
78
79
80
81
82
83
84
85
86
87
88
          """
          try:
              response = requests.post(
                  self.endpoint,
                  json=request_data,
                  timeout=60
              )
              response.raise_for_status()
              return response.json()
          except requests.exceptions.RequestException as e:
              logger.error(f"BgeEncoder service request failed: {e}", exc_info=True)
              raise
  
be52af70   tangwang   first commit
89
90
91
92
      def encode(
          self,
          sentences: Union[str, List[str]],
          normalize_embeddings: bool = True,
325eec03   tangwang   1. 日志、配置基础设施,使用优化
93
          device: str = 'cpu',
be52af70   tangwang   first commit
94
95
96
          batch_size: int = 32
      ) -> np.ndarray:
          """
453992a8   tangwang   需求:
97
          Encode text into embeddings via network service with Redis caching.
be52af70   tangwang   first commit
98
99
100
  
          Args:
              sentences: Single string or list of strings to encode
325eec03   tangwang   1. 日志、配置基础设施,使用优化
101
102
103
              normalize_embeddings: Whether to normalize embeddings (ignored for service)
              device: Device parameter ignored for service compatibility
              batch_size: Batch size for processing (used for service requests)
be52af70   tangwang   first commit
104
105
  
          Returns:
b2e50710   tangwang   BgeEncoder.encode...
106
107
108
              numpy array of dtype=object, where each element is either:
              - np.ndarray (valid embedding vector) or
              - None (no embedding available for that text)
be52af70   tangwang   first commit
109
          """
325eec03   tangwang   1. 日志、配置基础设施,使用优化
110
111
112
          # Convert single string to list
          if isinstance(sentences, str):
              sentences = [sentences]
be52af70   tangwang   first commit
113
  
453992a8   tangwang   需求:
114
          # Check cache first
b2e50710   tangwang   BgeEncoder.encode...
115
116
          uncached_indices: List[int] = []
          uncached_texts: List[str] = []
453992a8   tangwang   需求:
117
          
70a318c6   tangwang   fix bug
118
119
120
121
122
123
124
125
126
127
128
129
130
          # Process response
          # Each element can be np.ndarray or None (表示该文本没有可用的向量)
          embeddings: List[Optional[np.ndarray]] = [None] * len(sentences)
  
          for i, text in enumerate(sentences):
              cached = self._get_cached_embedding(text, 'en')  # Use 'en' as default language for title embedding
              if cached is not None:
                  embeddings[i] = cached
              else:
                  uncached_indices.append(i)
                  uncached_texts.append(text)
          
          # Prepare request data for uncached texts (after cache check)
7bfb9946   tangwang   向量化模块
131
          request_data = list(uncached_texts)
453992a8   tangwang   需求:
132
133
134
135
136
137
138
139
140
141
          
          # If there are uncached texts, call service
          if uncached_texts:
              try:
                  # Call service
                  response_data = self._call_service(request_data)
  
                  # Process response
                  for i, text in enumerate(uncached_texts):
                      original_idx = uncached_indices[i]
7bfb9946   tangwang   向量化模块
142
143
144
                      if response_data and i < len(response_data):
                          embedding = response_data[i]
                      else:
453992a8   tangwang   需求:
145
                          embedding = None
7bfb9946   tangwang   向量化模块
146
147
148
149
150
151
152
153
  
                      if embedding is not None:
                          embedding_array = np.array(embedding, dtype=np.float32)
                          # Validate embedding from service - if invalid, treat as no result
                          if self._is_valid_embedding(embedding_array):
                              embeddings[original_idx] = embedding_array
                              # Cache the embedding
                              self._set_cached_embedding(text, 'en', embedding_array)
453992a8   tangwang   需求:
154
                          else:
7bfb9946   tangwang   向量化模块
155
156
157
158
159
                              logger.warning(
                                  f"Invalid embedding returned from service for text {original_idx} "
                                  f"(contains NaN/Inf or invalid shape), treating as no result. "
                                  f"Text preview: {text[:50]}..."
                              )
b2e50710   tangwang   BgeEncoder.encode...
160
                              embeddings[original_idx] = None
325eec03   tangwang   1. 日志、配置基础设施,使用优化
161
                      else:
7bfb9946   tangwang   向量化模块
162
                          logger.warning(f"No embedding found for text {original_idx}: {text[:50]}...")
b2e50710   tangwang   BgeEncoder.encode...
163
                          embeddings[original_idx] = None
453992a8   tangwang   需求:
164
165
166
  
              except Exception as e:
                  logger.error(f"Failed to encode texts: {e}", exc_info=True)
b2e50710   tangwang   BgeEncoder.encode...
167
168
                  # 出错时不要生成兜底全零向量,保持为 None
                  pass
453992a8   tangwang   需求:
169
          
b2e50710   tangwang   BgeEncoder.encode...
170
171
          # 返回 numpy 数组(dtype=object),元素为 np.ndarray 或 None
          return np.array(embeddings, dtype=object)
be52af70   tangwang   first commit
172
173
174
175
176
  
      def encode_batch(
          self,
          texts: List[str],
          batch_size: int = 32,
325eec03   tangwang   1. 日志、配置基础设施,使用优化
177
          device: str = 'cpu'
be52af70   tangwang   first commit
178
179
      ) -> np.ndarray:
          """
325eec03   tangwang   1. 日志、配置基础设施,使用优化
180
          Encode a batch of texts efficiently via network service.
be52af70   tangwang   first commit
181
182
183
184
  
          Args:
              texts: List of texts to encode
              batch_size: Batch size for processing
325eec03   tangwang   1. 日志、配置基础设施,使用优化
185
              device: Device parameter ignored for service compatibility
be52af70   tangwang   first commit
186
187
188
189
190
  
          Returns:
              numpy array of embeddings
          """
          return self.encode(texts, batch_size=batch_size, device=device)
453992a8   tangwang   需求:
191
192
193
194
195
      
      def _get_cache_key(self, query: str, language: str) -> str:
          """Generate a cache key for the query"""
          return f"embedding:{language}:{query}"
      
b2e50710   tangwang   BgeEncoder.encode...
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
      def _is_valid_embedding(self, embedding: np.ndarray) -> bool:
          """
          Check if embedding is valid (not None, correct shape, no NaN/Inf).
          
          Args:
              embedding: Embedding array to validate
              
          Returns:
              True if valid, False otherwise
          """
          if embedding is None:
              return False
          if not isinstance(embedding, np.ndarray):
              return False
          if embedding.size == 0:
              return False
          # Check for NaN or Inf values
          if not np.isfinite(embedding).all():
              return False
          return True
      
453992a8   tangwang   需求:
217
218
219
220
221
222
223
224
225
      def _get_cached_embedding(self, query: str, language: str) -> Optional[np.ndarray]:
          """Get embedding from cache if exists (with sliding expiration)"""
          if not self.redis_client:
              return None
              
          try:
              cache_key = self._get_cache_key(query, language)
              cached_data = self.redis_client.get(cache_key)
              if cached_data:
b2e50710   tangwang   BgeEncoder.encode...
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
                  embedding = pickle.loads(cached_data)
                  # Validate cached embedding - if invalid, ignore cache and return None
                  if self._is_valid_embedding(embedding):
                      logger.debug(f"Cache hit for embedding: {query}")
                      # Update expiration time on access (sliding expiration)
                      self.redis_client.expire(cache_key, self.expire_time)
                      return embedding
                  else:
                      logger.warning(
                          f"Invalid embedding found in cache (contains NaN/Inf or invalid shape), "
                          f"ignoring cache for query: {query[:50]}..."
                      )
                      # Delete invalid cache entry
                      try:
                          self.redis_client.delete(cache_key)
                      except Exception as e:
                          logger.debug(f"Failed to delete invalid cache entry: {e}")
                      return None
453992a8   tangwang   需求:
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
              return None
          except Exception as e:
              logger.error(f"Error retrieving embedding from cache: {e}")
              return None
      
      def _set_cached_embedding(self, query: str, language: str, embedding: np.ndarray) -> bool:
          """Store embedding in cache"""
          if not self.redis_client:
              return False
              
          try:
              cache_key = self._get_cache_key(query, language)
              serialized_data = pickle.dumps(embedding)
              self.redis_client.setex(
                  cache_key,
                  self.expire_time,
                  serialized_data
              )
              logger.debug(f"Successfully cached embedding for query: {query}")
              return True
          except (redis.exceptions.BusyLoadingError, redis.exceptions.ConnectionError, 
                  redis.exceptions.TimeoutError, redis.exceptions.RedisError) as e:
              logger.warning(f"Redis error storing embedding in cache: {e}")
              return False
          except Exception as e:
              logger.error(f"Error storing embedding in cache: {e}")
              return False