dish-image-toolkit/retrieval/embed_local.py

# -*- coding: utf-8 -*-
"""本地 BGE-M3 向量(BAAI/bge-m3,1024 维 dense,多语言/中文强,离线无限额)。

与 cloud 版 embed.py 同接口(embed_corpus / matrix / embed_query / last_failed / last_error),
server.py 直接换用。首次会下载约 2.3G 权重(默认走 hf-mirror.com 镜像,国内快);
检测到 CUDA 自动用 GPU。无任何调用配额/限速问题。
"""
from __future__ import annotations

import json
import os
import sys
import threading

import numpy as np

_HERE = os.path.dirname(os.path.abspath(__file__))
# 默认缓存在 C 盘 .cache,但本机 C: 空间不足 → 放到项目所在盘(bge-m3 约 2.3G)。
# 想放别处:启动前设环境变量 HF_HOME 覆盖即可。
os.environ.setdefault("HF_HOME", os.path.join(_HERE, "hf_home"))
# 国内默认走 HF 镜像(已设则不覆盖),下载更稳;并阻止 transformers 误加载 TensorFlow。
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
# 关掉 hf-xet 传输:它绕过 HF_ENDPOINT 直连 huggingface.co 的 xet 服务器会超时。
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
# 系统代理(Clash 7897)会解压 zstd 并改写 hf-mirror 的 etag,破坏 hf_hub 元数据校验
# (报"couldn't connect to hf-mirror.com")。hf-mirror 是国内站,直连即可 → 清掉本进程代理。
for _k in ("HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy", "ALL_PROXY", "all_proxy"):
    os.environ.pop(_k, None)
os.environ.setdefault("NO_PROXY", "*")

MODEL_ID = os.environ.get("BGE_MODEL", "BAAI/bge-m3")


class LocalEmbedder:
    def __init__(self, model: str = MODEL_ID, cache_dir: str = "cache",
                 device: str | None = None):
        self.model_id = model
        tag = model.replace("/", "_")
        os.makedirs(cache_dir, exist_ok=True)
        self.cache_path = os.path.join(cache_dir, f"emb_{tag}.jsonl")
        self.cache: dict[str, list[float]] = {}
        self.last_failed = 0
        self.last_error = ""
        self.device = device
        self._model = None
        self._lock = threading.Lock()
        self._load_cache()

    # ---------- 缓存(与 cloud 版一致) ----------
    def _load_cache(self) -> None:
        if not os.path.exists(self.cache_path):
            return
        with open(self.cache_path, encoding="utf-8") as f:
            for line in f:
                try:
                    o = json.loads(line)
                    self.cache[o["t"]] = o["v"]
                except Exception:
                    pass

    def _append_cache(self, items) -> None:
        with self._lock:
            with open(self.cache_path, "a", encoding="utf-8") as f:
                for t, v in items:
                    f.write(json.dumps({"t": t, "v": v}, ensure_ascii=False) + "\n")
                    self.cache[t] = v

    # ---------- 模型 ----------
    def load_model(self) -> str:
        """加载模型(首次会下载权重),返回 device 字符串。"""
        if self._model is None:
            import torch
            # bge-m3 仅发布 pytorch_model.bin(无 safetensors);transformers 5.x 禁止
            # torch<2.6 加载 .bin(CVE-2025-32434,防不可信 pickle 执行代码)。本机 torch 2.5.1,
            # 模型来自 BAAI 官方(可信)、本地离线加载 → 关掉该检查,免升级 torch / 免下 safetensors。
            try:
                import transformers.modeling_utils as _mu
                _mu.check_torch_load_is_safe = lambda *a, **k: None
            except Exception:
                pass
            from sentence_transformers import SentenceTransformer
            dev = self.device or ("cuda" if torch.cuda.is_available() else "cpu")
            self._model = SentenceTransformer(self.model_id, device=dev)
            self.device = dev
        return self.device

    def _encode(self, texts: list[str]) -> np.ndarray:
        return np.asarray(
            self._model.encode(texts, normalize_embeddings=True, batch_size=64,
                               convert_to_numpy=True, show_progress_bar=False),
            dtype=np.float32)

    # ---------- 编码 ----------
    def embed_corpus(self, names, progress=None) -> None:
        """把 names 全部编码进缓存。progress(done, total) 回调可选。"""
        todo = [n for n in dict.fromkeys(names) if n not in self.cache]
        total = len(todo)
        self.last_failed = 0
        self.last_error = ""
        if progress:
            progress(0, total)
        if not total:
            return
        try:
            self.load_model()
        except Exception as e:  # noqa: BLE001
            self.last_failed = total
            self.last_error = f"加载模型失败: {type(e).__name__}: {e}"
            print(f"[embed_local] {self.last_error}", file=sys.stderr, flush=True)
            return
        step = 256
        done = 0
        for i in range(0, total, step):
            chunk = todo[i:i + step]
            try:
                vecs = self._encode(chunk)
                self._append_cache([(t, v.tolist()) for t, v in zip(chunk, vecs)])
            except Exception as e:  # noqa: BLE001
                self.last_failed += len(chunk)
                self.last_error = str(e)[:200]
                print(f"[embed_local] 编码失败: {e}", file=sys.stderr, flush=True)
            done += len(chunk)
            if progress:
                progress(done, total)

    # ---------- 取用(与 cloud 版一致) ----------
    def matrix(self, names) -> tuple[list[str], np.ndarray]:
        keep = [n for n in names if n in self.cache]
        if not keep:
            return [], np.zeros((0, 0), dtype=np.float32)
        m = np.array([self.cache[n] for n in keep], dtype=np.float32)
        norm = np.linalg.norm(m, axis=1, keepdims=True)
        norm[norm == 0] = 1.0
        return keep, m / norm

    def embed_query(self, q: str) -> np.ndarray:
        self.load_model()
        v = self._encode([q])[0]
        n = np.linalg.norm(v)
        return v / (n if n else 1.0)