embed norm

tangwang
1 parent 654f20d1
Showing 10 changed files with 128 additions and 44 deletions Show diff stats
embeddings/README.md
embeddings/clip_as_service_encoder.py
embeddings/clip_model.py
embeddings/config.py
embeddings/image_encoder.py
embeddings/protocols.py
embeddings/server.py
embeddings/text_encoder.py
restart.sh
scripts/service_ctl.sh
@@ -14,7 +14,7 @@
 - **clip-as-service 客户端**：`clip_as_service_encoder.py`（图片向量，推荐）
 - **向量化服务（FastAPI）**：`server.py`
 - **统一配置**：`config.py`
-- **接口契约**：`protocols.ImageEncoderProtocol`（图片编码统一为 `encode_image_urls(urls, batch_size)`，本地 CN-CLIP 与 clip-as-service 均实现该接口）
+- **接口契约**：`protocols.ImageEncoderProtocol`（图片编码统一为 `encode_image_urls(urls, batch_size, normalize_embeddings)`，本地 CN-CLIP 与 clip-as-service 均实现该接口）
 说明：历史上的云端 embedding 试验实现（DashScope）已从主仓库移除，当前仅维护 6005 这条统一向量服务链路。
@@ -29,10 +29,12 @@
 - `POST /embed/text`
   - 入参：`["文本1", "文本2", ...]`
+  - 可选 query 参数：`normalize=true|false`（不传则使用服务端默认）
   - 出参：`[[...], [...], ...]`（与输入按 index 对齐，失败直接报错）
 - `POST /embed/image`
   - 入参：`["url或本地路径1", ...]`
+  - 可选 query 参数：`normalize=true|false`（不传则使用服务端默认）
   - 出参：`[[...], [...], ...]`（与输入按 index 对齐，失败直接报错）
 ### 图片向量：clip-as-service（推荐）
@@ -77,6 +79,7 @@ TEI_USE_GPU=0 ./scripts/start_tei_service.sh
 - `PORT`: 服务端口（默认 6005）
 - `TEXT_MODEL_ID`, `TEXT_DEVICE`, `TEXT_BATCH_SIZE`, `TEXT_NORMALIZE_EMBEDDINGS`
+- `IMAGE_NORMALIZE_EMBEDDINGS`（默认 true）
 - `USE_CLIP_AS_SERVICE`, `CLIP_AS_SERVICE_SERVER`：图片向量（clip-as-service）
 - `IMAGE_MODEL_NAME`, `IMAGE_DEVICE`：本地 CN-CLIP（当 `USE_CLIP_AS_SERVICE=false` 时）
 - TEI 相关：`TEI_USE_GPU`、`TEI_VERSION`、`TEI_MAX_BATCH_TOKENS`、`TEI_MAX_CLIENT_BATCH_SIZE`、`TEI_HEALTH_TIMEOUT_SEC`
@@ -79,6 +79,7 @@ class ClipAsServiceImageEncoder:
         self,
         urls: List[str],
         batch_size: Optional[int] = None,
+        normalize_embeddings: bool = True,
     ) -> List[np.ndarray]:
         """
         Encode a list of image URLs to vectors.
@@ -117,10 +118,15 @@ class ClipAsServiceImageEncoder:
             vec = np.asarray(row, dtype=np.float32)
             if vec.ndim != 1 or vec.size == 0 or not np.isfinite(vec).all():
                 raise RuntimeError("clip-as-service returned invalid embedding vector")
+            if normalize_embeddings:
+                norm = float(np.linalg.norm(vec))
+                if not np.isfinite(norm) or norm <= 0.0:
+                    raise RuntimeError("clip-as-service returned zero/invalid norm vector")
+                vec = vec / norm
             out.append(vec)
         return out
-    def encode_image_from_url(self, url: str) -> Optional[np.ndarray]:
+    def encode_image_from_url(self, url: str, normalize_embeddings: bool = True) -> Optional[np.ndarray]:
         """Encode a single image URL. Returns 1024-dim vector or None."""
-        results = self.encode_image_urls([url], batch_size=1)
+        results = self.encode_image_urls([url], batch_size=1, normalize_embeddings=normalize_embeddings)
         return results[0] if results else None
@@ -76,25 +76,27 @@ class ClipImageModel(object):
             text_features /= text_features.norm(dim=-1, keepdim=True)
         return text_features
-    def encode_image(self, image: Image.Image) -> Optional[np.ndarray]:
+    def encode_image(self, image: Image.Image, normalize_embeddings: bool = True) -> Optional[np.ndarray]:
         if not isinstance(image, Image.Image):
             raise ValueError("ClipImageModel.encode_image input must be a PIL.Image")
         infer_data = self.preprocess(image).unsqueeze(0).to(self.device)
         with torch.no_grad():
             image_features = self.model.encode_image(infer_data)
-            image_features /= image_features.norm(dim=-1, keepdim=True)
+            if normalize_embeddings:
+                image_features /= image_features.norm(dim=-1, keepdim=True)
         return image_features.cpu().numpy().astype("float32")[0]
-    def encode_image_from_url(self, url: str) -> Optional[np.ndarray]:
+    def encode_image_from_url(self, url: str, normalize_embeddings: bool = True) -> Optional[np.ndarray]:
         image_data = self.download_image(url)
         image = self.validate_image(image_data)
         image = self.preprocess_image(image)
-        return self.encode_image(image)
+        return self.encode_image(image, normalize_embeddings=normalize_embeddings)
     def encode_image_urls(
         self,
         urls: List[str],
         batch_size: Optional[int] = None,
+        normalize_embeddings: bool = True,
     ) -> List[Optional[np.ndarray]]:
         """
         Encode a list of image URLs to vectors. Same interface as ClipAsServiceImageEncoder.
@@ -106,19 +108,27 @@ class ClipImageModel(object):
         Returns:
             List of vectors (or None for failed items), same length as urls.
         """
-        return self.encode_batch(urls, batch_size=batch_size or 8)
+        return self.encode_batch(
+            urls,
+            batch_size=batch_size or 8,
+            normalize_embeddings=normalize_embeddings,
+        )
-    def encode_batch(self, images: List[Union[str, Image.Image]], batch_size: int = 8) -> List[Optional[np.ndarray]]:
+    def encode_batch(
+        self,
+        images: List[Union[str, Image.Image]],
+        batch_size: int = 8,
+        normalize_embeddings: bool = True,
+    ) -> List[Optional[np.ndarray]]:
         results: List[Optional[np.ndarray]] = []
         for i in range(0, len(images), batch_size):
             batch = images[i : i + batch_size]
             for img in batch:
                 if isinstance(img, str):
-                    results.append(self.encode_image_from_url(img))
+                    results.append(self.encode_image_from_url(img, normalize_embeddings=normalize_embeddings))
                 elif isinstance(img, Image.Image):
-                    results.append(self.encode_image(img))
+                    results.append(self.encode_image(img, normalize_embeddings=normalize_embeddings))
                 else:
                     results.append(None)
         return results
-
@@ -37,6 +37,7 @@ class EmbeddingConfig(object):
     # Service behavior
     IMAGE_BATCH_SIZE = 8
+    IMAGE_NORMALIZE_EMBEDDINGS = os.getenv("IMAGE_NORMALIZE_EMBEDDINGS", "true").lower() in ("1", "true", "yes")
 CONFIG = EmbeddingConfig()
@@ -26,7 +26,7 @@ class CLIPImageEncoder:
         self.endpoint = f"{self.service_url}/embed/image"
         logger.info("Creating CLIPImageEncoder instance with service URL: %s", self.service_url)
-    def _call_service(self, request_data: List[str]) -> List[Any]:
+    def _call_service(self, request_data: List[str], normalize_embeddings: bool = True) -> List[Any]:
         """
         Call the embedding service API.
@@ -39,6 +39,7 @@ class CLIPImageEncoder:
         try:
             response = requests.post(
                 self.endpoint,
+                params={"normalize": "true" if normalize_embeddings else "false"},
                 json=request_data,
                 timeout=60
             )
@@ -56,7 +57,7 @@ class CLIPImageEncoder:
         """
         raise NotImplementedError("encode_image with PIL Image is not supported by embedding service")
-    def encode_image_from_url(self, url: str) -> np.ndarray:
+    def encode_image_from_url(self, url: str, normalize_embeddings: bool = True) -> np.ndarray:
         """
         Generate image embedding via network service using URL.
@@ -66,7 +67,7 @@ class CLIPImageEncoder:
         Returns:
             Embedding vector
         """
-        response_data = self._call_service([url])
+        response_data = self._call_service([url], normalize_embeddings=normalize_embeddings)
         if not response_data or len(response_data) != 1 or response_data[0] is None:
             raise RuntimeError(f"No image embedding returned for URL: {url}")
         vec = np.array(response_data[0], dtype=np.float32)
@@ -77,7 +78,8 @@ class CLIPImageEncoder:
     def encode_batch(
         self,
         images: List[Union[str, Image.Image]],
-        batch_size: int = 8
+        batch_size: int = 8,
+        normalize_embeddings: bool = True,
     ) -> List[np.ndarray]:
         """
         Encode a batch of images efficiently via network service.
@@ -98,7 +100,7 @@ class CLIPImageEncoder:
         results: List[np.ndarray] = []
         for i in range(0, len(images), batch_size):
             batch_urls = [str(u).strip() for u in images[i:i + batch_size]]
-            response_data = self._call_service(batch_urls)
+            response_data = self._call_service(batch_urls, normalize_embeddings=normalize_embeddings)
             if not response_data or len(response_data) != len(batch_urls):
                 raise RuntimeError(
                     f"Image embedding response length mismatch: expected {len(batch_urls)}, "
@@ -119,6 +121,7 @@ class CLIPImageEncoder:
         self,
         urls: List[str],
         batch_size: Optional[int] = None,
+        normalize_embeddings: bool = True,
     ) -> List[np.ndarray]:
         """
         与 ClipImageModel / ClipAsServiceImageEncoder 一致的接口，供索引器 document_transformer 调用。
@@ -130,4 +133,8 @@ class CLIPImageEncoder:
         Returns:
             与 urls 等长的向量列表
         """
-        return self.encode_batch(urls, batch_size=batch_size or 8)
+        return self.encode_batch(
+            urls,
+            batch_size=batch_size or 8,
+            normalize_embeddings=normalize_embeddings,
+        )
@@ -17,6 +17,7 @@ class ImageEncoderProtocol(Protocol):
         self,
         urls: List[str],
         batch_size: Optional[int] = None,
+        normalize_embeddings: bool = True,
     ) -> List[Optional[np.ndarray]]:
         """
         Encode a list of image URLs to vectors.
@@ -112,14 +112,24 @@ def load_models():
     logger.info("All embedding models loaded successfully, service ready")
-def _as_list(embedding: Optional[np.ndarray]) -> Optional[List[float]]:
+def _normalize_vector(vec: np.ndarray) -> np.ndarray:
+    norm = float(np.linalg.norm(vec))
+    if not np.isfinite(norm) or norm <= 0.0:
+        raise RuntimeError("Embedding vector has invalid norm (must be > 0)")
+    return vec / norm
+
+
+def _as_list(embedding: Optional[np.ndarray], normalize: bool = False) -> Optional[List[float]]:
     if embedding is None:
         return None
     if not isinstance(embedding, np.ndarray):
         embedding = np.array(embedding, dtype=np.float32)
     if embedding.ndim != 1:
         embedding = embedding.reshape(-1)
-    return embedding.astype(np.float32).tolist()
+    embedding = embedding.astype(np.float32, copy=False)
+    if normalize:
+        embedding = _normalize_vector(embedding).astype(np.float32, copy=False)
+    return embedding.tolist()
 @app.get("/health")
@@ -134,9 +144,10 @@ def health() -&gt; Dict[str, Any]:
 @app.post("/embed/text")
-def embed_text(texts: List[str]) -> List[Optional[List[float]]]:
+def embed_text(texts: List[str], normalize: Optional[bool] = None) -> List[Optional[List[float]]]:
     if _text_model is None:
         raise RuntimeError("Text model not loaded")
+    effective_normalize = bool(CONFIG.TEXT_NORMALIZE_EMBEDDINGS) if normalize is None else bool(normalize)
     normalized: List[str] = []
     for i, t in enumerate(texts):
         if not isinstance(t, str):
@@ -152,7 +163,7 @@ def embed_text(texts: List[str]) -&gt; List[Optional[List[float]]]:
                 normalized,
                 batch_size=int(CONFIG.TEXT_BATCH_SIZE),
                 device=CONFIG.TEXT_DEVICE,
-                normalize_embeddings=bool(CONFIG.TEXT_NORMALIZE_EMBEDDINGS),
+                normalize_embeddings=effective_normalize,
             )
     except Exception as e:
         logger.error("Text embedding backend failure: %s", e, exc_info=True)
@@ -167,7 +178,7 @@ def embed_text(texts: List[str]) -&gt; List[Optional[List[float]]]:
         )
     out: List[Optional[List[float]]] = []
     for i, emb in enumerate(embs):
-        vec = _as_list(emb)
+        vec = _as_list(emb, normalize=effective_normalize)
         if vec is None:
             raise RuntimeError(f"Text model returned empty embedding for index {i}")
         out.append(vec)
@@ -175,9 +186,10 @@ def embed_text(texts: List[str]) -&gt; List[Optional[List[float]]]:
 @app.post("/embed/image")
-def embed_image(images: List[str]) -> List[Optional[List[float]]]:
+def embed_image(images: List[str], normalize: Optional[bool] = None) -> List[Optional[List[float]]]:
     if _image_model is None:
         raise RuntimeError("Image model not loaded")
+    effective_normalize = bool(CONFIG.IMAGE_NORMALIZE_EMBEDDINGS) if normalize is None else bool(normalize)
     urls: List[str] = []
     for i, url_or_path in enumerate(images):
         if not isinstance(url_or_path, str):
@@ -188,7 +200,11 @@ def embed_image(images: List[str]) -&gt; List[Optional[List[float]]]:
         urls.append(s)
     with _image_encode_lock:
-        vectors = _image_model.encode_image_urls(urls, batch_size=CONFIG.IMAGE_BATCH_SIZE)
+        vectors = _image_model.encode_image_urls(
+            urls,
+            batch_size=CONFIG.IMAGE_BATCH_SIZE,
+            normalize_embeddings=effective_normalize,
+        )
     if vectors is None or len(vectors) != len(urls):
         raise RuntimeError(
             f"Image model response length mismatch: expected {len(urls)}, "
@@ -196,7 +212,7 @@ def embed_image(images: List[str]) -&gt; List[Optional[List[float]]]:
         )
     out: List[Optional[List[float]]] = []
     for i, vec in enumerate(vectors):
-        out_vec = _as_list(vec)
+        out_vec = _as_list(vec, normalize=effective_normalize)
         if out_vec is None:
             raise RuntimeError(f"Image model returned empty embedding for index {i}")
         out.append(out_vec)
@@ -50,7 +50,7 @@ class TextEmbeddingEncoder:
             logger.warning("Failed to initialize Redis cache for embeddings: %s, continuing without cache", e)
             self.redis_client = None
-    def _call_service(self, request_data: List[str]) -> List[Any]:
+    def _call_service(self, request_data: List[str], normalize_embeddings: bool = True) -> List[Any]:
         """
         Call the embedding service API.
@@ -63,6 +63,7 @@ class TextEmbeddingEncoder:
         try:
             response = requests.post(
                 self.endpoint,
+                params={"normalize": "true" if normalize_embeddings else "false"},
                 json=request_data,
                 timeout=60
             )
@@ -84,7 +85,7 @@ class TextEmbeddingEncoder:
         Args:
             sentences: Single string or list of strings to encode
-            normalize_embeddings: Whether to normalize embeddings (ignored for service)
+            normalize_embeddings: Whether to request normalized embeddings from service
             device: Device parameter ignored for service compatibility
             batch_size: Batch size for processing (used for service requests)
@@ -103,7 +104,7 @@ class TextEmbeddingEncoder:
         embeddings: List[Optional[np.ndarray]] = [None] * len(sentences)
         for i, text in enumerate(sentences):
-            cached = self._get_cached_embedding(text, "generic")
+            cached = self._get_cached_embedding(text, "generic", normalize_embeddings)
             if cached is not None:
                 embeddings[i] = cached
             else:
@@ -115,7 +116,7 @@ class TextEmbeddingEncoder:
         # If there are uncached texts, call service
         if uncached_texts:
-            response_data = self._call_service(request_data)
+            response_data = self._call_service(request_data, normalize_embeddings=normalize_embeddings)
             # Process response
             for i, text in enumerate(uncached_texts):
@@ -129,7 +130,7 @@ class TextEmbeddingEncoder:
                     embedding_array = np.array(embedding, dtype=np.float32)
                     if self._is_valid_embedding(embedding_array):
                         embeddings[original_idx] = embedding_array
-                        self._set_cached_embedding(text, "generic", embedding_array)
+                        self._set_cached_embedding(text, "generic", embedding_array, normalize_embeddings)
                     else:
                         raise ValueError(
                             f"Invalid embedding returned from service for text index {original_idx}"
@@ -144,7 +145,8 @@ class TextEmbeddingEncoder:
         self,
         texts: List[str],
         batch_size: int = 32,
-        device: str = 'cpu'
+        device: str = 'cpu',
+        normalize_embeddings: bool = True,
     ) -> np.ndarray:
         """
         Encode a batch of texts efficiently via network service.
@@ -157,11 +159,17 @@ class TextEmbeddingEncoder:
         Returns:
             numpy array of embeddings
         """
-        return self.encode(texts, batch_size=batch_size, device=device)
+        return self.encode(
+            texts,
+            batch_size=batch_size,
+            device=device,
+            normalize_embeddings=normalize_embeddings,
+        )
-    def _get_cache_key(self, query: str, language: str) -> str:
+    def _get_cache_key(self, query: str, language: str, normalize_embeddings: bool = True) -> str:
         """Generate a cache key for the query"""
-        return f"embedding:{language}:{query}"
+        norm_flag = "norm1" if normalize_embeddings else "norm0"
+        return f"embedding:{language}:{norm_flag}:{query}"
     def _is_valid_embedding(self, embedding: np.ndarray) -> bool:
         """
@@ -184,13 +192,18 @@ class TextEmbeddingEncoder:
             return False
         return True
-    def _get_cached_embedding(self, query: str, language: str) -> Optional[np.ndarray]:
+    def _get_cached_embedding(
+        self,
+        query: str,
+        language: str,
+        normalize_embeddings: bool = True,
+    ) -> 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)
+            cache_key = self._get_cache_key(query, language, normalize_embeddings)
             cached_data = self.redis_client.get(cache_key)
             if cached_data:
                 embedding = pickle.loads(cached_data)
@@ -216,13 +229,19 @@ class TextEmbeddingEncoder:
             logger.error(f"Error retrieving embedding from cache: {e}")
             return None
-    def _set_cached_embedding(self, query: str, language: str, embedding: np.ndarray) -> bool:
+    def _set_cached_embedding(
+        self,
+        query: str,
+        language: str,
+        embedding: np.ndarray,
+        normalize_embeddings: bool = True,
+    ) -> bool:
         """Store embedding in cache"""
         if not self.redis_client:
             return False
         try:
-            cache_key = self._get_cache_key(query, language)
+            cache_key = self._get_cache_key(query, language, normalize_embeddings)
             serialized_data = pickle.dumps(embedding)
             self.redis_client.setex(
                 cache_key,
@@ -4,4 +4,7 @@
 cd "$(dirname "$0")"
-./scripts/service_ctl.sh restart
+START_EMBEDDING=1 START_TRANSLATOR=1 START_RERANKER=1 START_TEI=1 CNCLIP_DEVICE=cuda TEI_USE_GPU=1 ./scripts/service_ctl.sh restart
+
+# ./scripts/service_ctl.sh restart
+
@@ -190,7 +190,14 @@ start_one() {
       nohup bash -lc "${cmd}" > "${lf}" 2>&1 &
       local pid=$!
       echo "${pid}" > "${pf}"
-      if wait_for_health "${service}"; then
+      # Some services (notably reranker with vLLM backend) can take longer
+      # to load models / compile graphs on first start. Give them a longer
+      # health check window to avoid false negatives.
+      local retries=30
+      if [ "${service}" = "reranker" ]; then
+        retries=90
+      fi
+      if wait_for_health "${service}" "${retries}"; then
         echo "[ok] ${service} healthy (pid=${pid}, log=${lf})"
       else
         echo "[error] ${service} health check timeout, inspect ${lf}" >&2
@@ -320,9 +327,14 @@ resolve_targets() {
       if [ "${START_RERANKER:-0}" = "1" ]; then targets+=("reranker"); fi
       echo "${targets[@]}"
       ;;
-    stop|restart|status)
+    stop|status)
       echo "$(all_services)"
       ;;
+    restart)
+      # Restart with no explicit services should preserve start-order dependency
+      # behavior (e.g. tei/cnclip before embedding).
+      echo "$(resolve_targets start)"
+      ;;
     *)
       echo ""
       ;;
@@ -340,7 +352,7 @@ Usage:
 Default target set (when no service provided):
   start   -> backend indexer frontend (+ optional by env flags)
   stop    -> all known services
-  restart -> all known services
+  restart -> stop all known services, then start with start targets
   status  -> all known services
 Optional startup flags:
@@ -361,7 +373,13 @@ main() {
   shift || true
   load_env_file "${PROJECT_ROOT}/.env"
+  local stop_targets=""
   local targets
+  # For restart without explicit services, stop everything first, then start
+  # with dependency-aware start targets.
+  if [ "${action}" = "restart" ] && [ "$#" -eq 0 ]; then
+    stop_targets="$(resolve_targets stop)"
+  fi
   targets="$(resolve_targets "${action}" "$@")"
   if [ -z "${targets}" ]; then
     usage
@@ -380,7 +398,7 @@ main() {
       done
       ;;
     restart)
-      for svc in ${targets}; do
+      for svc in ${stop_targets:-${targets}}; do
         stop_one "${svc}"
       done
       for svc in ${targets}; do