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
|