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
  
42e3aea6   tangwang   tidy
21
22
  from config.services_config import get_embedding_base_url
  
453992a8   tangwang   需求:
23
24
25
26
27
28
  # 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
29
30
31
  
  class BgeEncoder:
      """
325eec03   tangwang   1. 日志、配置基础设施,使用优化
32
      Singleton text encoder using network service.
be52af70   tangwang   first commit
33
  
325eec03   tangwang   1. 日志、配置基础设施,使用优化
34
      Thread-safe singleton pattern ensures only one instance exists.
be52af70   tangwang   first commit
35
36
37
38
      """
      _instance = None
      _lock = threading.Lock()
  
7bfb9946   tangwang   向量化模块
39
      def __new__(cls, service_url: Optional[str] = None):
be52af70   tangwang   first commit
40
41
42
          with cls._lock:
              if cls._instance is None:
                  cls._instance = super(BgeEncoder, cls).__new__(cls)
42e3aea6   tangwang   tidy
43
                  resolved_url = service_url or os.getenv("EMBEDDING_SERVICE_URL") or get_embedding_base_url()
7bfb9946   tangwang   向量化模块
44
45
46
                  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   需求:
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
                  
                  # 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
67
68
          return cls._instance
  
7bfb9946   tangwang   向量化模块
69
      def _call_service(self, request_data: List[str]) -> List[Any]:
325eec03   tangwang   1. 日志、配置基础设施,使用优化
70
71
72
73
          """
          Call the embedding service API.
  
          Args:
7bfb9946   tangwang   向量化模块
74
              request_data: List of texts
325eec03   tangwang   1. 日志、配置基础设施,使用优化
75
76
  
          Returns:
7bfb9946   tangwang   向量化模块
77
              List of embeddings (list[float]) or nulls (None), aligned to input order
325eec03   tangwang   1. 日志、配置基础设施,使用优化
78
79
80
81
82
83
84
85
86
87
88
89
90
          """
          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
91
92
93
94
      def encode(
          self,
          sentences: Union[str, List[str]],
          normalize_embeddings: bool = True,
325eec03   tangwang   1. 日志、配置基础设施,使用优化
95
          device: str = 'cpu',
be52af70   tangwang   first commit
96
97
98
          batch_size: int = 32
      ) -> np.ndarray:
          """
453992a8   tangwang   需求:
99
          Encode text into embeddings via network service with Redis caching.
be52af70   tangwang   first commit
100
101
102
  
          Args:
              sentences: Single string or list of strings to encode
325eec03   tangwang   1. 日志、配置基础设施,使用优化
103
104
105
              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
106
107
  
          Returns:
b2e50710   tangwang   BgeEncoder.encode...
108
109
110
              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
111
          """
325eec03   tangwang   1. 日志、配置基础设施,使用优化
112
113
114
          # Convert single string to list
          if isinstance(sentences, str):
              sentences = [sentences]
be52af70   tangwang   first commit
115
  
453992a8   tangwang   需求:
116
          # Check cache first
b2e50710   tangwang   BgeEncoder.encode...
117
118
          uncached_indices: List[int] = []
          uncached_texts: List[str] = []
453992a8   tangwang   需求:
119
          
70a318c6   tangwang   fix bug
120
121
122
123
124
125
126
127
128
129
130
131
132
          # 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   向量化模块
133
          request_data = list(uncached_texts)
453992a8   tangwang   需求:
134
135
136
137
138
139
140
141
142
143
          
          # 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   向量化模块
144
145
146
                      if response_data and i < len(response_data):
                          embedding = response_data[i]
                      else:
453992a8   tangwang   需求:
147
                          embedding = None
7bfb9946   tangwang   向量化模块
148
149
150
151
152
153
154
155
  
                      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   需求:
156
                          else:
7bfb9946   tangwang   向量化模块
157
158
159
160
161
                              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...
162
                              embeddings[original_idx] = None
325eec03   tangwang   1. 日志、配置基础设施,使用优化
163
                      else:
7bfb9946   tangwang   向量化模块
164
                          logger.warning(f"No embedding found for text {original_idx}: {text[:50]}...")
b2e50710   tangwang   BgeEncoder.encode...
165
                          embeddings[original_idx] = None
453992a8   tangwang   需求:
166
167
168
  
              except Exception as e:
                  logger.error(f"Failed to encode texts: {e}", exc_info=True)
b2e50710   tangwang   BgeEncoder.encode...
169
170
                  # 出错时不要生成兜底全零向量,保持为 None
                  pass
453992a8   tangwang   需求:
171
          
b2e50710   tangwang   BgeEncoder.encode...
172
173
          # 返回 numpy 数组(dtype=object),元素为 np.ndarray 或 None
          return np.array(embeddings, dtype=object)
be52af70   tangwang   first commit
174
175
176
177
178
  
      def encode_batch(
          self,
          texts: List[str],
          batch_size: int = 32,
325eec03   tangwang   1. 日志、配置基础设施,使用优化
179
          device: str = 'cpu'
be52af70   tangwang   first commit
180
181
      ) -> np.ndarray:
          """
325eec03   tangwang   1. 日志、配置基础设施,使用优化
182
          Encode a batch of texts efficiently via network service.
be52af70   tangwang   first commit
183
184
185
186
  
          Args:
              texts: List of texts to encode
              batch_size: Batch size for processing
325eec03   tangwang   1. 日志、配置基础设施,使用优化
187
              device: Device parameter ignored for service compatibility
be52af70   tangwang   first commit
188
189
190
191
192
  
          Returns:
              numpy array of embeddings
          """
          return self.encode(texts, batch_size=batch_size, device=device)
453992a8   tangwang   需求:
193
194
195
196
197
      
      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...
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
      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   需求:
219
220
221
222
223
224
225
226
227
      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...
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
                  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   需求:
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
271
272
              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