tidy embeddings

tangwang
1 parent 3d588bef
Showing 9 changed files with 49 additions and 65 deletions Show diff stats
docs/DEVELOPER_GUIDE.md
docs/MySQL到ES文档映射说明.md
embeddings/README.md
embeddings/server.py
embeddings/qwen3_model.py -> embeddings/text_embedding_sentence_transformers.py
embeddings/tei_model.py -> embeddings/text_embedding_tei.py
embeddings/text_encoder.py
indexer/incremental_service.py
indexer/product_enrich.py
@@ -315,7 +315,7 @@ services:
 **重排后端协议（服务内）**：所有在 reranker 服务内加载的后端须实现 `score_with_meta(query, docs, normalize=True) -> (scores: List[float], meta: dict)`。返回的 `scores[i]` 与 `docs[i]` 一一对应；meta 至少含 `input_docs`、`usable_docs`、`elapsed_ms` 等。对外 HTTP 契约固定：`POST /rerank` 请求体 `{ "query": str, "docs": [str] }`，响应体 `{ "scores": [float], "meta": object }`；`GET /health` 返回 `status`、`model`、`backend` 等。
-**向量化后端协议（服务内）**：文本后端需支持 `encode_batch(texts, batch_size, device) -> List[ndarray]`，与 texts 一一对应；图片后端实现 `embeddings/protocols.ImageEncoderProtocol`：`encode_image_urls(urls, batch_size) -> List[Optional[ndarray]]`，与 urls 等长。
+**向量化后端协议（服务内）**：文本后端需支持 `encode(sentences: Union[str, List[str]], batch_size, device) -> ndarray | List[ndarray]`，单条与批量输入统一通过一个接口处理；图片后端实现 `embeddings/protocols.ImageEncoderProtocol`：`encode_image_urls(urls, batch_size) -> List[Optional[ndarray]]`，与 urls 等长。
 **配置速查**：
@@ -679,7 +679,7 @@ if enable_embedding and encoder and documents:
             title_doc_indices.append(i)
     if title_texts:
-        embeddings = encoder.encode_batch(title_texts, batch_size=32)
+        embeddings = encoder.encode(title_texts, batch_size=32)
         for j, emb in enumerate(embeddings):
             doc_idx = title_doc_indices[j]
             if isinstance(emb, np.ndarray):
@@ -731,7 +731,7 @@ if enable_embedding and encoder and documents:
 7. **批量生成 Embedding**（如果启用）
    - 收集所有文档的标题文本
-   - 批量调用 `encoder.encode_batch()` 生成 embedding
+   - 批量调用 `encoder.encode()`（传入 list[str]）生成 embedding
    - 填充到对应文档
 8. **批量写入 ES**
@@ -10,7 +10,7 @@
 这个目录是一个完整的“向量化模块”，包含：
 - **HTTP 客户端**：`text_encoder.py` / `image_encoder.py`（供搜索/索引模块调用）
-- **本地模型实现**：`qwen3_model.py` / `clip_model.py`
+- **本地模型实现**：`text_embedding_sentence_transformers.py` / `clip_model.py`
 - **clip-as-service 客户端**：`clip_as_service_encoder.py`（图片向量，推荐）
 - **向量化服务（FastAPI）**：`server.py`
 - **统一配置**：`config.py`
@@ -83,7 +83,7 @@ def _preview_inputs(items: List[str], max_items: int, max_chars: int) -&gt; List[Di
 def _encode_local_st(texts: List[str], normalize_embeddings: bool) -> Any:
     with _text_encode_lock:
-        return _text_model.encode_batch(
+        return _text_model.encode(
             texts,
             batch_size=int(CONFIG.TEXT_BATCH_SIZE),
             device=CONFIG.TEXT_DEVICE,
@@ -198,7 +198,7 @@ def load_models():
             backend_name, backend_cfg = get_embedding_backend_config()
             _text_backend_name = backend_name
             if backend_name == "tei":
-                from embeddings.tei_model import TEITextModel
+                from embeddings.text_embedding_tei import TEITextModel
                 base_url = (
                     os.getenv("TEI_BASE_URL")
@@ -216,7 +216,7 @@ def load_models():
                     timeout_sec=timeout_sec,
                 )
             elif backend_name == "local_st":
-                from embeddings.qwen3_model import Qwen3TextModel
+                from embeddings.text_embedding_sentence_transformers import Qwen3TextModel
                 model_id = (
                     os.getenv("TEXT_MODEL_ID")
@@ -342,7 +342,7 @@ def embed_text(texts: List[str], normalize: Optional[bool] = None) -&gt; List[Optio
                 return out
             embs = _encode_local_st(normalized, normalize_embeddings=False)
         else:
-            embs = _text_model.encode_batch(
+            embs = _text_model.encode(
                 normalized,
                 batch_size=int(CONFIG.TEXT_BATCH_SIZE),
                 device=CONFIG.TEXT_DEVICE,
@@ -47,6 +47,7 @@ class Qwen3TextModel(object):
         device: str = "cuda",
         batch_size: int = 32,
     ) -> np.ndarray:
+
         # SentenceTransformer + CUDA inference is not thread-safe in our usage;
         # keep one in-flight encode call while avoiding repeated .to(device) hops.
         with self._encode_lock:
@@ -60,16 +61,3 @@ class Qwen3TextModel(object):
             )
             return embeddings
-    def encode_batch(
-        self,
-        texts: List[str],
-        batch_size: int = 32,
-        device: str = "cuda",
-        normalize_embeddings: bool = True,
-    ) -> np.ndarray:
-        return self.encode(
-            texts,
-            batch_size=batch_size,
-            device=device,
-            normalize_embeddings=normalize_embeddings,
-        )
@@ -54,24 +54,17 @@ class TEITextModel:
         device: str = "cuda",
         batch_size: int = 32,
     ) -> np.ndarray:
-        if isinstance(sentences, str):
-            sentences = [sentences]
-        return self.encode_batch(
-            texts=sentences,
-            batch_size=batch_size,
-            device=device,
-            normalize_embeddings=normalize_embeddings,
-        )
+        """
+        Encode a single sentence or a list of sentences.
-    def encode_batch(
-        self,
-        texts: List[str],
-        batch_size: int = 32,
-        device: str = "cuda",
-        normalize_embeddings: bool = True,
-    ) -> np.ndarray:
-        del batch_size  # TEI performs its own batching.
-        del device      # Not used by HTTP backend.
+        TEI HTTP 后端天然是批量接口，这里统一通过 encode 处理单条和批量输入，
+        不再额外暴露 encode_batch。
+        """
+
+        if isinstance(sentences, str):
+            texts: List[str] = [sentences]
+        else:
+            texts = sentences
         if texts is None or len(texts) == 0:
             return np.array([], dtype=object)
@@ -135,33 +135,8 @@ class TextEmbeddingEncoder:
                 else:
                     raise ValueError(f"No embedding found for text index {original_idx}: {text[:50]}...")
-        # 返回 numpy 数组（dtype=object），元素为 np.ndarray 或 None
+        # 返回 numpy 数组（dtype=object），元素均为有效 np.ndarray 向量
         return np.array(embeddings, dtype=object)
-
-    def encode_batch(
-        self,
-        texts: List[str],
-        batch_size: int = 32,
-        device: str = 'cpu',
-        normalize_embeddings: bool = True,
-    ) -> np.ndarray:
-        """
-        Encode a batch of texts efficiently via network service.
-
-        Args:
-            texts: List of texts to encode
-            batch_size: Batch size for processing
-            device: Device parameter ignored for service compatibility
-
-        Returns:
-            numpy array of embeddings
-        """
-        return self.encode(
-            texts,
-            batch_size=batch_size,
-            device=device,
-            normalize_embeddings=normalize_embeddings,
-        )
     def _is_valid_embedding(self, embedding: np.ndarray) -> bool:
         """
@@ -641,7 +641,7 @@ class IncrementalIndexerService:
                             title_doc_indices.append(i)
                     if title_texts:
-                        embeddings = encoder.encode_batch(title_texts, batch_size=32)
+                        embeddings = encoder.encode(title_texts, batch_size=32)
                         if embeddings is None or len(embeddings) != len(title_texts):
                             raise RuntimeError(
                                 f"[IncrementalIndexing] Batch embedding length mismatch for tenant_id={tenant_id}: "
@@ -96,6 +96,9 @@ except Exception as e:
     logger.warning(f"Failed to initialize Redis for anchors cache: {e}")
     _anchor_redis = None
+# 中文版本提示词（请勿删除）：
+# "你是一名电商平台的商品标注员，你的工作是对输入的每个商品进行理解、分析和标注，"
+# "并按要求格式返回 Markdown 表格。所有输出内容必须为中文。"
 SYSTEM_MESSAGES = (
     "You are a product annotator for an e-commerce platform. "
@@ -163,6 +166,31 @@ def create_prompt(products: List[Dict[str, str]], target_lang: str = &quot;zh&quot;) -&gt; st
     """
     lang_name = SOURCE_LANG_CODE_MAP.get(target_lang, target_lang)
+# 中文版本提示词（请勿删除）
+# prompt = """请对输入的每条商品标题，分析并提取以下信息：
+
+# 1. 商品标题：将输入商品名称翻译为自然、完整的中文商品标题
+# 2. 品类路径：从大类到细分品类，用">"分隔（例如：服装>女装>裤子>工装裤）
+# 3. 细分标签：商品的风格、特点、功能等（例如：碎花，收腰，法式）
+# 4. 适用人群：性别/年龄段等（例如：年轻女性）
+# 5. 使用场景
+# 6. 适用季节
+# 7. 关键属性
+# 8. 材质说明 
+# 9. 功能特点
+# 10. 商品卖点：分析和提取一句话核心卖点，用于推荐理由
+# 11. 锚文本：生成一组能够代表该商品、并可能被用户用于搜索的词语或短语。这些词语应覆盖用户需求的各个维度，如品类、细分标签、功能特性、需求场景等等。
+
+# 输入商品列表：
+
+# """
+#         prompt_tail = """
+# 请严格按照以下markdown表格格式返回，每列内部的多值内容都用逗号分隔，不要添加任何其他说明：
+
+# | 序号 | 商品标题 | 品类路径 | 细分标签 | 适用人群 | 使用场景 | 适用季节 | 关键属性 | 材质说明 | 功能特点 | 商品卖点 | 锚文本 |
+# |----|----|----|----|----|----|----|----|----|----|----|----|
+# """
+
     prompt = """Please analyze each input product title and extract the following information:
 1. Product title: a natural English product name derived from the input title