49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
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()
|