File size: 2,312 Bytes
9b32fb0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """Embedding 服务封装 — sentence-transformers 包装."""
import numpy as np
class Embedder:
"""文本向量化服务。
默认使用 paraphrase-multilingual-MiniLM-L12-v2，384 维，纯 CPU 推理，
支持 50+ 语言（含中文）。首次创建实例时下载模型（约 420MB），后续使用缓存。
如果模型不可用（网络不通等），操作会抛出 RuntimeError。
"""
_model = None
_model_name = None
def __init__(self, model_name: str = "paraphrase-multilingual-MiniLM-L12-v2"):
from sentence_transformers import SentenceTransformer
try:
self._model = SentenceTransformer(model_name)
self._model_name = model_name
except Exception as e:
raise RuntimeError(
f"Failed to load embedding model '{model_name}': {e}. "
f"Ensure network access to HuggingFace or pre-download the model."
) from e
@staticmethod
def is_available() -> bool:
"""检查 embedding 模型是否可用（不会触发下载，仅检查是否已缓存）。"""
try:
from sentence_transformers import SentenceTransformer
import os
# Check if model is cached locally
cache_dir = os.path.join(os.path.expanduser("~"), ".cache",
"torch", "sentence_transformers")
model_dir = os.path.join(cache_dir, "all-MiniLM-L6-v2")
if os.path.isdir(model_dir):
return True
# Try a quick offline check
return False
except ImportError:
return False
def encode(self, text: str) -> list[float]:
"""编码单条文本，返回 384 维浮点向量。"""
vec = self._model.encode(text, normalize_embeddings=True)
return vec.tolist()
def encode_batch(self, texts: list[str]) -> list[list[float]]:
"""批量编码，返回多个 384 维向量。"""
vecs = self._model.encode(texts, normalize_embeddings=True)
return vecs.tolist()
@staticmethod
def cosine_similarity(a: list[float], b: list[float]) -> float:
"""计算两个归一化向量的余弦相似度（已归一化时即为点积）。"""
return float(np.dot(a, b))
|