Init: 第一个版本

app/embeddings.py

@@ -0,0 +1,48 @@
+import logging
+import os
+import urllib.request
+
+from app.config import EMBEDDING_MODEL
+
+logger = logging.getLogger("learn_rag")
+QUERY_PREFIX = "为这个句子生成表示以用于检索相关文章："
+
+
+def configure_download() -> None:
+    """Hugging Face 不可达时改用国内镜像，避免首次下载卡住。"""
+    if os.getenv("HF_ENDPOINT"):
+        return
+    try:
+        with urllib.request.urlopen("https://huggingface.co", timeout=4) as response:
+            if response.status < 400:
+                return
+    except Exception:
+        logger.info("无法连接 huggingface.co，嵌入模型改从 hf-mirror.com 下载")
+        os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+
+class Embedder:
+    def __init__(self) -> None:
+        configure_download()
+        from fastembed import TextEmbedding
+
+        logger.info("正在加载嵌入模型 %s", EMBEDDING_MODEL)
+        self.model = TextEmbedding(model_name=EMBEDDING_MODEL)
+        probe = next(iter(self.model.embed(["书店"])))
+        self.dimension = len(probe)
+        logger.info("嵌入模型就绪，维度 %s", self.dimension)
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+        return [vector.tolist() for vector in self.model.embed(texts)]
+
+    def embed_query(self, text: str) -> list[float]:
+        query_embed = getattr(self.model, "query_embed", None)
+        if query_embed is not None:
+            try:
+                return next(iter(query_embed([text]))).tolist()
+            except Exception:
+                logger.info("query_embed 不可用，改用查询前缀")
+        prefixed = QUERY_PREFIX + text
+        return next(iter(self.model.embed([prefixed]))).tolist()