[fix] 恢复远程embedding失败时的本地兜底

ziwenhahaha · ziwenhahaha · commit 202829da9e9a · 2026-03-13T01:34:44.000+08:00
diff --git a/.github/workflows/daily-paper-reader.yml b/.github/workflows/daily-paper-reader.yml
@@ -55,23 +55,21 @@ jobs:
         with:
           python-version: "3.11"
 
-      - name: Cache pip
+      - name: Cache Python deps
         uses: actions/cache@v4
         with:
-          path: ~/.cache/pip
-          key: ${{ runner.os }}-dpr-remote-pip-v1-${{ hashFiles('requirements.txt') }}
+          path: |
+            ~/.cache/pip
+            ~/.cache/uv
+            ~/.cache/torch
+          key: ${{ runner.os }}-dpr-embed-deps-v1-${{ hashFiles('requirements.txt') }}
 
       - name: Install deps (skip sqlite3)
         run: |
           python - <<'PY'
           import re
           lines = open("requirements.txt", "r", encoding="utf-8").read().splitlines()
-          lines = [
-            l for l in lines
-            if l.strip()
-            and not re.match(r"^sqlite3\\b", l)
-            and not re.match(r"^sentence-transformers\\b", l)
-          ]
+          lines = [l for l in lines if l.strip() and not re.match(r"^sqlite3\\b", l)]
           open("/tmp/req.txt", "w", encoding="utf-8").write("\n".join(lines))
           PY
           python -m pip install --upgrade pip
diff --git a/src/model_loader.py b/src/model_loader.py
@@ -5,7 +5,7 @@
 from contextlib import contextmanager
 import os
 import time
-from typing import Callable, Optional, TYPE_CHECKING
+from typing import Any, Callable, Optional, TYPE_CHECKING
 
 import numpy as np
 import requests
@@ -44,6 +44,12 @@ def __init__(
     api_key: str = "",
     timeout: int = _DEFAULT_REMOTE_TIMEOUT_SECONDS,
     default_batch_size: int = 8,
+    local_device: str = "cpu",
+    local_retries: int | None = None,
+    local_providers: tuple[tuple[str, str], ...] = (
+      ("huggingface", HUGGINGFACE_ENDPOINT),
+      ("modelscope", MODELSCOPE_ENDPOINT),
+    ),
     log: Callable[[str], None] = _log_default,
   ):
     self.model_name = model_name
@@ -52,6 +58,10 @@ def __init__(
     self.timeout = max(int(timeout or _DEFAULT_REMOTE_TIMEOUT_SECONDS), 1)
     self.default_batch_size = max(int(default_batch_size or 1), 1)
     self.max_seq_length = None
+    self.local_device = str(local_device or "cpu")
+    self.local_retries = local_retries
+    self.local_providers = local_providers
+    self._local_model = None
     self._log = log
 
   @staticmethod
@@ -71,6 +81,26 @@ def _headers(self) -> dict[str, str]:
       headers["Authorization"] = f"Bearer {self.api_key}"
     return headers
 
+  def _get_local_model(self):
+    if self._local_model is None:
+      self._log(
+        f"[WARN] 远程 embedding 不可用，回退本地模型：{self.model_name} "
+        f"(device={self.local_device})"
+      )
+      self._local_model = _load_local_sentence_transformer(
+        self.model_name,
+        device=self.local_device,
+        retries=self.local_retries,
+        log=self._log,
+        providers=self.local_providers,
+      )
+      if self.max_seq_length is not None and hasattr(self._local_model, "max_seq_length"):
+        try:
+          self._local_model.max_seq_length = self.max_seq_length
+        except Exception:
+          pass
+    return self._local_model
+
   def encode(
     self,
     texts,
@@ -80,7 +110,6 @@ def encode(
     show_progress_bar: bool = False,
     **kwargs,
   ):
-    del show_progress_bar, kwargs
     if isinstance(texts, str):
       texts = [texts]
     if not isinstance(texts, list):
@@ -90,60 +119,78 @@ def encode(
       return empty if convert_to_numpy else empty.tolist()
 
     safe_batch_size = max(int(batch_size or self.default_batch_size), 1)
-    chunks = [texts[i : i + safe_batch_size] for i in range(0, len(texts), safe_batch_size)]
-    outputs: list[np.ndarray] = []
-
-    self._log(
-      f"[INFO] 远程 embedding：model={self.model_name} "
-      f"endpoint={self.endpoint} total={len(texts)} batch={safe_batch_size}"
-    )
+    try:
+      chunks = [texts[i : i + safe_batch_size] for i in range(0, len(texts), safe_batch_size)]
+      outputs: list[np.ndarray] = []
 
-    for chunk_index, chunk in enumerate(chunks, start=1):
-      headers = self._headers()
-      response = requests.post(
-        self.endpoint,
-        headers=headers,
-        json={"texts": chunk},
-        timeout=self.timeout,
+      self._log(
+        f"[INFO] 远程 embedding：model={self.model_name} "
+        f"endpoint={self.endpoint} total={len(texts)} batch={safe_batch_size}"
       )
-      if response.status_code == 401 and headers.get("Authorization"):
-        self._log("[WARN] 远程 embedding 鉴权失败，自动回退为无鉴权请求重试一次。")
-        headers = {
-          "Content-Type": "application/json",
-        }
+
+      for chunk_index, chunk in enumerate(chunks, start=1):
+        headers = self._headers()
         response = requests.post(
           self.endpoint,
           headers=headers,
           json={"texts": chunk},
           timeout=self.timeout,
         )
-      response.raise_for_status()
-      data = response.json()
-      embeddings = data.get("embeddings")
-      if not isinstance(embeddings, list):
-        raise RuntimeError("远程 embedding 服务返回缺少 embeddings 字段")
-      try:
-        arr = np.asarray(embeddings, dtype=np.float32)
-      except Exception as exc:
-        raise RuntimeError(f"远程 embedding 返回无法转换为 float32：{exc}") from exc
-
-      if arr.ndim != 2:
-        raise RuntimeError(f"远程 embedding 返回维度异常：shape={getattr(arr, 'shape', None)}")
-      if arr.shape[0] != len(chunk):
-        raise RuntimeError(
-          f"远程 embedding 返回条数异常：expected={len(chunk)} actual={arr.shape[0]}"
+        if response.status_code == 401 and headers.get("Authorization"):
+          self._log("[WARN] 远程 embedding 鉴权失败，自动回退为无鉴权请求重试一次。")
+          headers = {
+            "Content-Type": "application/json",
+          }
+          response = requests.post(
+            self.endpoint,
+            headers=headers,
+            json={"texts": chunk},
+            timeout=self.timeout,
+          )
+        response.raise_for_status()
+        data = response.json()
+        embeddings = data.get("embeddings")
+        if not isinstance(embeddings, list):
+          raise RuntimeError("远程 embedding 服务返回缺少 embeddings 字段")
+        try:
+          arr = np.asarray(embeddings, dtype=np.float32)
+        except Exception as exc:
+          raise RuntimeError(f"远程 embedding 返回无法转换为 float32：{exc}") from exc
+
+        if arr.ndim != 2:
+          raise RuntimeError(f"远程 embedding 返回维度异常：shape={getattr(arr, 'shape', None)}")
+        if arr.shape[0] != len(chunk):
+          raise RuntimeError(
+            f"远程 embedding 返回条数异常：expected={len(chunk)} actual={arr.shape[0]}"
+          )
+        if normalize_embeddings:
+          norms = np.linalg.norm(arr, axis=1, keepdims=True)
+          arr = arr / np.clip(norms, 1e-12, None)
+        outputs.append(arr)
+        self._log(
+          f"[INFO] 远程 embedding 批次完成：{chunk_index}/{len(chunks)} "
+          f"count={len(chunk)} dim={arr.shape[1]}"
         )
-      if normalize_embeddings:
-        norms = np.linalg.norm(arr, axis=1, keepdims=True)
-        arr = arr / np.clip(norms, 1e-12, None)
-      outputs.append(arr)
-      self._log(
-        f"[INFO] 远程 embedding 批次完成：{chunk_index}/{len(chunks)} "
-        f"count={len(chunk)} dim={arr.shape[1]}"
-      )
 
-    merged = np.vstack(outputs) if outputs else np.zeros((0, 0), dtype=np.float32)
-    return merged if convert_to_numpy else merged.tolist()
+      merged = np.vstack(outputs) if outputs else np.zeros((0, 0), dtype=np.float32)
+      return merged if convert_to_numpy else merged.tolist()
+    except Exception as exc:
+      self._log(f"[WARN] 远程 embedding 请求失败，将自动回退本地模型：{exc}")
+      local_model = self._get_local_model()
+      result = local_model.encode(
+        texts,
+        convert_to_numpy=convert_to_numpy,
+        normalize_embeddings=normalize_embeddings,
+        batch_size=safe_batch_size,
+        show_progress_bar=show_progress_bar,
+        **kwargs,
+      )
+      if convert_to_numpy and not isinstance(result, np.ndarray):
+        try:
+          result = np.asarray(result, dtype=np.float32)
+        except Exception:
+          pass
+      return result
 
   def start_multi_process_pool(self, target_devices=None):
     del target_devices
@@ -266,9 +313,32 @@ def load_sentence_transformer(
       endpoint=str(remote_endpoint).strip(),
       api_key=remote_api_key,
       timeout=remote_timeout,
+      local_device=device,
+      local_retries=retries,
+      local_providers=providers,
       log=log,
     )
 
+  return _load_local_sentence_transformer(
+    model_name,
+    device=device,
+    retries=retries,
+    log=log,
+    providers=providers,
+  )
+
+
+def _load_local_sentence_transformer(
+  model_name: str,
+  *,
+  device: str,
+  retries: int | None = None,
+  log: Callable[[str], None] = _log_default,
+  providers: tuple[tuple[str, str], ...] = (
+    ("huggingface", HUGGINGFACE_ENDPOINT),
+    ("modelscope", MODELSCOPE_ENDPOINT),
+  ),
+):
   if retries is None:
     env_retries = os.getenv("LLM_EMBED_MODEL_RETRIES")
     if env_retries is None:
diff --git a/tests/test_model_loader_remote.py b/tests/test_model_loader_remote.py
@@ -3,6 +3,7 @@
 from unittest.mock import MagicMock, patch
 
 import numpy as np
+import requests
 
 from src.model_loader import RemoteSentenceTransformer, load_sentence_transformer
 
@@ -51,6 +52,33 @@ def test_remote_encode_batches_and_normalizes(self, mock_post):
         self.assertEqual(first_call.kwargs["headers"]["Authorization"], "Bearer test-key")
         self.assertEqual(first_call.kwargs["timeout"], 30)
 
+    @patch("src.model_loader._load_local_sentence_transformer")
+    @patch("src.model_loader.requests.post")
+    def test_remote_encode_falls_back_to_local_model_when_remote_fails(self, mock_post, mock_load_local):
+        mock_post.side_effect = requests.exceptions.Timeout("remote timeout")
+        local_model = MagicMock()
+        local_model.encode.return_value = np.asarray([[0.1, 0.2]], dtype=np.float32)
+        mock_load_local.return_value = local_model
+
+        model = RemoteSentenceTransformer(
+            model_name="BAAI/bge-small-en-v1.5",
+            endpoint="https://embed.zwwen.online",
+            api_key="test-key",
+            timeout=30,
+            default_batch_size=2,
+        )
+        arr = model.encode(
+            ["a"],
+            convert_to_numpy=True,
+            normalize_embeddings=True,
+            batch_size=2,
+        )
+
+        self.assertEqual(mock_post.call_count, 1)
+        mock_load_local.assert_called_once()
+        local_model.encode.assert_called_once()
+        self.assertEqual(arr.shape, (1, 2))
+
     @patch.dict(
         os.environ,
         {
