[beautify] 优化远程embedding模式日志与依赖安装

Author: ziwenhahaha

.github/workflows/daily-paper-reader.yml

@@ -55,20 +55,23 @@ jobs:
         with:
           python-version: "3.11"
 
-      - name: Cache pip + torch
+      - name: Cache pip
         uses: actions/cache@v4
         with:
-          path: |
-            ~/.cache/pip
-            ~/.cache/torch
-          key: ${{ runner.os }}-dpr-hf-v2-${{ hashFiles('requirements.txt') }}
+          path: ~/.cache/pip
+          key: ${{ runner.os }}-dpr-remote-pip-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)]
+          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)
+          ]
           open("/tmp/req.txt", "w", encoding="utf-8").write("\n".join(lines))
           PY
           python -m pip install --upgrade pip

src/filter.py

@@ -1,18 +1,20 @@
 #!/usr/bin/env python
 # 通用向量检索工具：封装 sentence-transformers 的向量计算与粗筛逻辑
 
+from __future__ import annotations
+
 import os
 import numpy as np
-from typing import Any, Dict, List
+from typing import Any, Dict, List, TYPE_CHECKING
 import time
 from datetime import datetime, timezone
 
 os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS", "1")
 
-import torch
-from sentence_transformers import SentenceTransformer
+from model_loader import is_remote_embedding_enabled, load_sentence_transformer
 
-from model_loader import load_sentence_transformer
+if TYPE_CHECKING:
+  from sentence_transformers import SentenceTransformer
 
 # E5 系列推荐使用 query/passsage 前缀来区分检索侧与文档侧
 E5_QUERY_PREFIX = "query: "
@@ -30,6 +32,8 @@ def debug_hf_runtime(prefix: str) -> None:
   enable = (os.getenv("DPR_DEBUG_HF") == "1") or (os.getenv("GITHUB_ACTIONS") == "true")
   if not enable:
     return
+  if is_remote_embedding_enabled():
+    return
 
   log(f"[DEBUG][HF] {prefix}")
   keys = [
@@ -73,7 +77,7 @@ def ls_dir(path: str) -> None:
     ls_dir(hf_home)
 
 
-def _set_max_seq_length(model: SentenceTransformer, max_length: int | None) -> None:
+def _set_max_seq_length(model: Any, max_length: int | None) -> None:
   """尽量通过 SentenceTransformer 的 max_seq_length 控制截断长度。"""
   if max_length is None or max_length <= 0:
     return
@@ -93,7 +97,7 @@ def _set_max_seq_length(model: SentenceTransformer, max_length: int | None) -> N
 
 
 def encode_queries(
-  model: SentenceTransformer,
+  model: Any,
   texts: List[str],
   batch_size: int = 8,
   max_length: int | None = None,
@@ -128,7 +132,7 @@ def encode_queries(
 
 
 def compute_embeddings(
-  model: SentenceTransformer,
+  model: Any,
   items: List[Any],
   batch_size: int = 8,
   max_length: int | None = None,
@@ -206,15 +210,27 @@ def __init__(
     self.batch_size = batch_size
     self.max_length = max_length
 
+    remote_mode = is_remote_embedding_enabled()
     if device is None:
-      self.device = "cuda" if torch.cuda.is_available() else "cpu"
+      if remote_mode:
+        self.device = "remote"
+      else:
+        try:
+          import torch
+          self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        except Exception:
+          self.device = "cpu"
     else:
-      self.device = device
+      self.device = device if not remote_mode else "remote"
 
-    print(f"[INFO] 正在加载向量模型：{self.model_name}，device={self.device}")
-    debug_hf_runtime("before SentenceTransformer()")
+    if remote_mode:
+      print(f"[INFO] 正在初始化远程向量服务：{self.model_name}，device={self.device}")
+    else:
+      print(f"[INFO] 正在加载本地向量模型：{self.model_name}，device={self.device}")
+      debug_hf_runtime("before SentenceTransformer()")
     self.model = load_sentence_transformer(self.model_name, device=self.device)
-    debug_hf_runtime("after SentenceTransformer()")
+    if not remote_mode:
+      debug_hf_runtime("after SentenceTransformer()")
     _set_max_seq_length(self.model, self.max_length)
 
   def filter(self, items: List[Any], queries: List[Dict[str, Any]]) -> Dict[str, Any]:

src/model_loader.py

@@ -5,12 +5,13 @@
 from contextlib import contextmanager
 import os
 import time
-from typing import Callable, Optional
+from typing import Callable, Optional, TYPE_CHECKING
 
 import numpy as np
 import requests
 
-from sentence_transformers import SentenceTransformer
+if TYPE_CHECKING:
+  from sentence_transformers import SentenceTransformer
 
 
 HUGGINGFACE_ENDPOINT = "https://huggingface.co"
@@ -27,6 +28,10 @@ def _log_default(message: str) -> None:
   print(message, flush=True)
 
 
+def is_remote_embedding_enabled() -> bool:
+  return bool(str(_DEFAULT_REMOTE_EMBED_ENDPOINT or "").strip())
+
+
 class RemoteSentenceTransformer:
   """兼容 SentenceTransformer.encode 接口的远程 embedding 包装器。"""
 
@@ -299,6 +304,7 @@ def load_sentence_transformer(
           f"（provider={provider_name}，device={device}）"
         )
         with _hf_endpoint(endpoint), _hf_http_backoff(max_retries=hf_backoff_retries):
+          from sentence_transformers import SentenceTransformer
           return SentenceTransformer(model_name, device=device)
       except Exception as e:  # pragma: no cover - 仅异常路径
         last_err = e