hermes_core/embedder.py

"""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))