border-red-post-nalati/backend/app/services/embedding.py

"""Optional local embeddings; importing this module never loads model weights or PyTorch."""

import hashlib
import json
import math
import re
from dataclasses import asdict, dataclass
from typing import Protocol

QUERY_INSTRUCTION = "为这个句子生成表示以用于检索相关文章："


class EmbeddingError(ValueError):
    pass


@dataclass(frozen=True)
class EmbeddingSpec:
    model: str
    revision: str
    dimension: int
    token_limit: int = 512
    query_instruction: str = QUERY_INSTRUCTION
    implementation: str = "hf-cls-l2-cpu-fp32-v1"
    checkpoint_format: str = "safetensors"

    def __post_init__(self):
        if not isinstance(self.revision, str) or not re.fullmatch(r"[0-9a-f]{40}", self.revision):
            raise EmbeddingError("模型必须固定到完整 revision SHA")
        if (
            type(self.dimension) is not int
            or type(self.token_limit) is not int
            or self.dimension < 1
            or self.token_limit < 2
        ):
            raise EmbeddingError("向量维度和 token 上限无效")
        if (
            not isinstance(self.model, str)
            or not self.model
            or not isinstance(self.query_instruction, str)
        ):
            raise EmbeddingError("模型名称或查询指令无效")
        if self.checkpoint_format not in {"safetensors", "pytorch-weights-only"}:
            raise EmbeddingError("不支持的权重格式")

    @property
    def space_id(self) -> str:
        return hashlib.sha256(
            json.dumps(asdict(self), sort_keys=True, ensure_ascii=False).encode()
        ).hexdigest()


SMALL_ZH = EmbeddingSpec("BAAI/bge-small-zh-v1.5", "7999e1d3359715c523056ef9478215996d62a620", 512)
BASE_ZH = EmbeddingSpec(
    "BAAI/bge-base-zh-v1.5",
    "f03589ceff5aac7111bd60cfc7d497ca17ecac65",
    768,
    checkpoint_format="pytorch-weights-only",
)
LOCAL_MODELS = {"small": SMALL_ZH, "base": BASE_ZH}


class EmbeddingProvider(Protocol):
    spec: EmbeddingSpec

    def encode(self, texts: list[str], *, query: bool = False) -> list[list[float]]: ...


def validate_vectors(vectors: list[list[float]], count: int, spec: EmbeddingSpec):
    if len(vectors) != count:
        raise EmbeddingError("模型返回数量与输入不一致")
    for vector in vectors:
        try:
            finite = all(type(x) in {float, int} and math.isfinite(x) for x in vector)
        except (TypeError, OverflowError):
            finite = False
        if len(vector) != spec.dimension or not finite:
            raise EmbeddingError("模型返回维度错误或非有限数值")
        if not math.isclose(sum(x * x for x in vector), 1, abs_tol=1e-4):
            raise EmbeddingError("模型输出没有正确 L2 归一化")


class LocalEmbeddingProvider:
    def __init__(
        self,
        spec: EmbeddingSpec,
        *,
        batch_size: int = 8,
        threads: int = 2,
        cache_dir=None,
        offline: bool = True,
    ):
        if (
            type(batch_size) is not int
            or type(threads) is not int
            or not 1 <= batch_size <= 32
            or not 1 <= threads <= 4
        ):
            raise EmbeddingError("CPU batch_size 必须为 1–32，threads 必须为 1–4")
        import torch
        from huggingface_hub import snapshot_download
        from transformers import AutoModel, AutoTokenizer

        torch.set_num_threads(threads)
        self.spec, self.batch_size, self.torch = spec, batch_size, torch
        path = snapshot_download(
            repo_id=spec.model,
            revision=spec.revision,
            cache_dir=cache_dir,
            local_files_only=offline,
            token=False,
            allow_patterns=[
                "config.json",
                "model.safetensors"
                if spec.checkpoint_format == "safetensors"
                else "pytorch_model.bin",
                "tokenizer.json",
                "tokenizer_config.json",
                "special_tokens_map.json",
                "vocab.txt",
                "added_tokens.json",
            ],
        )
        self.tokenizer = AutoTokenizer.from_pretrained(
            path, local_files_only=True, trust_remote_code=False
        )
        self.model = (
            AutoModel.from_pretrained(
                path,
                local_files_only=True,
                trust_remote_code=False,
                use_safetensors=spec.checkpoint_format == "safetensors",
                weights_only=True,
                dtype=torch.float32,
            )
            .to("cpu")
            .eval()
        )
        if self.model.config.hidden_size != spec.dimension:
            raise EmbeddingError("模型实际维度与配置不一致")
        if self.tokenizer.model_max_length != spec.token_limit:
            raise EmbeddingError("模型 tokenizer 上限与配置不一致")

    def encode(self, texts: list[str], *, query: bool = False) -> list[list[float]]:
        if not isinstance(texts, list) or any(
            not isinstance(text, str) or not text.strip() or len(text) > 10_000 for text in texts
        ):
            raise EmbeddingError("每条输入须为非空文本且最多 10000 字符")
        torch = self.torch
        vectors = []
        for offset in range(0, len(texts), self.batch_size):
            batch = texts[offset : offset + self.batch_size]
            if query:
                batch = [self.spec.query_instruction + text for text in batch]
            encoded = self.tokenizer(
                batch, padding=True, truncation=False, return_tensors="pt", verbose=False
            )
            if encoded["input_ids"].shape[1] > self.spec.token_limit:
                raise EmbeddingError("输入超出模型 token 上限；请重新切分，禁止静默截断")
            with torch.inference_mode():
                features = self.model(**encoded).last_hidden_state[:, 0]
                normalized = torch.nn.functional.normalize(features, p=2, dim=1)
            vectors.extend(normalized.tolist())
        validate_vectors(vectors, len(texts), self.spec)
        return vectors