chore: combine chunk indexes and embedding evaluation

backend/app/services/embedding.py

@@ -0,0 +1,165 @@
+"""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