dsh-data-domain-knowledge 0.1.0-alpha.1

package/src/python/embedding_worker.py

@@ -0,0 +1,471 @@
+#!/usr/bin/env python3
+"""Small JSONL bridge for the local WeMM SentenceTransformers backend.
+
+The Host deliberately talks to this process over stdin/stdout instead of
+importing Python from Node.  That keeps the DSH plugin installable without a
+Python dependency while still making the real multimodal model a first-class
+local backend on macOS (MPS) and Linux (CPU/CUDA as selected by PyTorch).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import sys
+import tempfile
+from contextlib import redirect_stdout
+from typing import Any
+
+
+MODEL_CACHE: dict[tuple[str, str], tuple[Any, str]] = {}
+# (model, device, health) — the health dict is filled once, at load, by
+# `reranker_health` and reused for every request in this process.
+RERANK_CACHE: dict[tuple[str, str], tuple[Any, str, dict[str, Any]]] = {}
+
+# A real XLM-R vocabulary is ~250k entries. A *degenerate* one still loads and
+# still returns numbers: with `sentencepiece.bpe.model` / `tokenizer.json` missing
+# from the local snapshot, `XLMRobertaTokenizer` came up with `vocab_size = 5`
+# (the five special tokens only) and tokenized every input to `<s> <unk>… </s>`.
+# Measured: three texts (relevant / unrelated-topic / pure noise) scored
+# 0.999, 0.998, 0.9994 — i.e. the "reranker" was ranking `<unk>` soup, and stage 4
+# reported `status: ready` with `dropped: 0` the whole time. Silent, so it has to
+# be a hard error at load rather than something a caller might notice.
+RERANK_MIN_VOCAB = 1000
+
+
+def emit(value: dict[str, Any]) -> None:
+    sys.stdout.write(json.dumps(value, ensure_ascii=False, separators=(",", ":")) + "\n")
+    sys.stdout.flush()
+
+
+def choose_device(requested: str) -> str:
+    if requested in {"mps", "cpu"}:
+        return requested
+    try:
+        import torch
+
+        if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
+            return "mps"
+    except Exception:
+        pass
+    return "cpu"
+
+
+def load_local_first(build):
+    """Build a HuggingFace-backed model **from the local cache first**.
+
+    A machine that already has the weights must not need huggingface.co to answer
+    a query. `SentenceTransformer(...)` / `CrossEncoder(...)` resolve a bare repo id
+    by contacting the Hub, and with no egress that call does not fail fast — it
+    *hangs* until the plugin's timeout. Measured on this machine: the reranker
+    weights were fully present in `~/.cache/huggingface/hub` (1.0G), yet the first
+    rerank after startup spent the whole 60s timeout trying to reach the Hub and
+    came back as `EMBEDDING_TIMEOUT` — the wrong subsystem *and* the wrong cause.
+    `local_files_only=True` turns a cached model into an instant offline load; the
+    networked attempt stays for the genuinely-uncached first run.
+
+    Note what this flag does **not** prove: the first attempt below succeeds on a
+    *partially* downloaded snapshot too. That is exactly how the missing reranker
+    tokenizer went unnoticed — `local_files_only=True` loaded in 0.1s and returned
+    a tokenizer with an empty vocabulary. "Loads offline" is not "is usable";
+    `reranker_health` is what checks the second half.
+
+    `build(local_only)` receives the flag because the argument has to reach the
+    constructor, and only the caller knows its exact signature.
+    """
+    try:
+        return build(True)
+    except Exception:
+        # Uncached cache, or a cache that cannot satisfy the load: let the normal
+        # path run so its own error — which names the real problem — surfaces.
+        return build(False)
+
+def load_model(model_name: str, requested_device: str):
+    import torch
+    from sentence_transformers import SentenceTransformer
+
+    device = choose_device(requested_device)
+    # Do not force torch_dtype here.  The quantized checkpoint owns its weight
+    # loading policy; forcing a dtype breaks some custom WeMM loaders.
+    if device == "mps":
+        torch.set_float32_matmul_precision("high")
+    def load(path: str, target: str, local_only: bool = False):
+        kwargs = {"local_files_only": True} if local_only else {}
+        return SentenceTransformer(path, trust_remote_code=True, device=target, **kwargs)
+
+    def move_float8_storage_to_mps(model: Any) -> None:
+        # MPS does not expose Float8_e4m3fn tensors. The quantized checkpoint
+        # uses FP8 buffers for a subset of weights, while its forward methods
+        # already dequantize them to float32/bfloat16. Converting those buffers
+        # before the device move preserves the computation and makes the
+        # checkpoint usable on Apple Silicon; the W4A8 uint8 buffers stay packed.
+        float8 = getattr(torch, "float8_e4m3fn", None)
+        if float8 is None:
+            raise RuntimeError("MPS_FLOAT8_UNAVAILABLE")
+        for module in model.modules():
+            for name, buffer in list(module._buffers.items()):
+                if buffer is not None and buffer.dtype == float8:
+                    module._buffers[name] = buffer.to(dtype=torch.bfloat16)
+
+    # The released W4A8/FP8 checkpoint intentionally omits lm_head.weight, but
+    # its config still says tie_word_embeddings=true. Transformers then tries
+    # to assign the FP8 input buffer to an nn.Parameter and aborts. Keep the
+    # downloaded model immutable and load through a tiny symlinked directory
+    # with only this inference-irrelevant config flag corrected.
+    patched_dir = None
+    load_path = model_name
+    if os.path.isdir(model_name):
+        config_path = os.path.join(model_name, "config.json")
+        try:
+            with open(config_path, "r", encoding="utf-8") as handle:
+                config = json.load(handle)
+            if config.get("tie_word_embeddings") is True:
+                patched_dir = tempfile.mkdtemp(prefix="tita-wemm-")
+                for name in os.listdir(model_name):
+                    source = os.path.join(model_name, name)
+                    target = os.path.join(patched_dir, name)
+                    if name == "config.json":
+                        config["tie_word_embeddings"] = False
+                        with open(target, "w", encoding="utf-8") as handle:
+                            json.dump(config, handle)
+                    else:
+                        os.symlink(source, target)
+                load_path = patched_dir
+        except Exception:
+            if patched_dir is not None:
+                shutil.rmtree(patched_dir, ignore_errors=True)
+                patched_dir = None
+
+    def attempt(local_only: bool):
+        try:
+            if device == "mps":
+                model = load(load_path, "cpu", local_only)
+                move_float8_storage_to_mps(model)
+                model.to("mps")
+                return model, device
+            return load(load_path, device, local_only), device
+        except Exception:
+            # MPS support for custom/quantized checkpoints can vary by torch
+            # build. `auto` remains useful where CPU is the safe fallback.
+            if requested_device == "auto" and device == "mps":
+                return load(load_path, "cpu", local_only), "cpu"
+            raise
+
+    try:
+        return load_local_first(attempt)
+    finally:
+        if patched_dir is not None:
+            shutil.rmtree(patched_dir, ignore_errors=True)
+
+
+def make_input(item: Any) -> Any:
+    if isinstance(item, str):
+        return item
+    if not isinstance(item, dict):
+        raise ValueError("embedding input must be a string or object")
+    text = str(item.get("text") or "")
+    image = item.get("image")
+    if image:
+        # WeMM uses a shared text/image space.  Keeping the markdown image and
+        # its surrounding text in one input is what makes image-in-markdown
+        # searchable by a natural-language query.
+        value = {"image": os.fspath(image), "text": text or "Represent this image for retrieval."}
+        return value
+    return text
+
+
+def cached_model(model_name: str, requested_device: str):
+    key = (model_name, requested_device)
+    cached = MODEL_CACHE.get(key)
+    if cached is not None:
+        return cached
+    # Some model-loading libraries log informational messages to stdout. The
+    # stdout channel is the JSONL protocol, so keep those logs on stderr.
+    with redirect_stdout(sys.stderr):
+        loaded = load_model(model_name, requested_device)
+    MODEL_CACHE[key] = loaded
+    return loaded
+
+
+def has_official_path(model_wrapper: Any) -> bool:
+    """Whether this checkpoint exposes WeMM's own processor + embedding head.
+
+    WeMM's card is explicit that the official path is the processor + `model.embedding`
+    pair, and here it is not a preference. `SentenceTransformer.encode` on this
+    checkpoint returns **the same vector for every input** (measured: four unrelated
+    texts — Chinese prose, English SQL — pairwise cosine 1.000000, at both 2048 and
+    1024 dimensions, on CPU as well as MPS). The official path on the same texts gives
+    0.208 / 0.361 / 0.461. So the wrapper path is not a fallback for this checkpoint,
+    it is a trap; the Host's encoder health gate is the second line of defence.
+    """
+    # `SentenceTransformer` subclasses `torch.nn.Sequential`, so `isinstance(..., (list, tuple))`
+    # is False even though `model_wrapper[0]` is exactly the transformer module. Indexing is
+    # the capability test that actually holds.
+    try:
+        transformer = model_wrapper[0]
+    except Exception:
+        transformer = None
+    processor = getattr(transformer, "processor", None)
+    inner_model = getattr(transformer, "model", None)
+    return processor is not None and inner_model is not None and hasattr(inner_model, "embedding")
+
+
+def _render_text(processor: Any, text: str) -> str:
+    return processor.apply_chat_template(
+        [{"role": "user", "content": [{"type": "text", "text": text}]}],
+        tokenize=False,
+        add_generation_prompt=False,
+    )
+
+
+def encode_official_text(model_wrapper: Any, texts: list[str], device: str):
+    """Encode text-only inputs through the checkpoint's own processor path.
+
+    Batched on purpose: the processor accepts a list of rendered conversations and
+    pads them, and the result matches the per-item path to cos 0.9999 (measured) at a
+    fifth of the time. Per-item would be correct but needlessly slow — a 44-note store
+    re-encodes on every `rebuild all`.
+    """
+    import torch
+
+    transformer = model_wrapper[0]
+    processor = transformer.processor
+    inner_model = transformer.model
+    rendered = [_render_text(processor, text) for text in texts]
+    inputs = processor(text=rendered, return_tensors="pt", padding=True)
+    inputs = {key: (value.to(device) if isinstance(value, torch.Tensor) else value) for key, value in inputs.items()}
+    with torch.inference_mode():
+        embedding = inner_model.embedding(**inputs)
+    return embedding.detach().to("cpu").float().numpy()
+
+
+def encode_multimodal(model_wrapper: Any, items: list[dict[str, Any]], device: str):
+    """Encode image-bearing inputs through the official WeMM processor path.
+
+    SentenceTransformers 6.0's generic `apply_chat_template(tokenize=True)` path can
+    produce image features without the matching `<|image_pad|>` token for this
+    checkpoint, so the explicit processor path is used for Markdown images.
+    """
+    import torch
+    from qwen_vl_utils import process_vision_info
+
+    transformer = model_wrapper[0]
+    processor = transformer.processor
+    inner_model = transformer.model
+    outputs = []
+    for item in items:
+        messages = [{"role": "user", "content": [
+            {"type": "image", "image": os.fspath(item["image"])},
+            {"type": "text", "text": str(item.get("text") or "Represent this image for retrieval.")},
+        ]}]
+        rendered = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
+        images, videos, video_kwargs = process_vision_info(
+            messages,
+            image_patch_size=getattr(transformer, "image_patch_size", 16),
+            return_video_kwargs=True,
+            return_video_metadata=True,
+        )
+        kwargs = {"text": [rendered], "images": images, "videos": videos, "return_tensors": "pt", "padding": True, **video_kwargs}
+        inputs = processor(**kwargs)
+        inputs = {key: value.to(device) if isinstance(value, torch.Tensor) else value for key, value in inputs.items()}
+        with torch.inference_mode():
+            embedding = inner_model.embedding(**inputs)
+        outputs.append(embedding.detach().to("cpu").float().numpy()[0])
+    return outputs
+
+
+def encode(request: dict[str, Any]) -> dict[str, Any]:
+    model_name = str(request.get("model") or "ewin-reg/WeMM-Embedding-2B-Quantized")
+    requested_device = str(request.get("device") or "auto")
+    dimension = int(request.get("dimension") or 1024)
+    batch_size = max(1, int(request.get("batchSize") or 8))
+    raw_inputs = request.get("inputs")
+    if not isinstance(raw_inputs, list) or not raw_inputs:
+        raise ValueError("inputs must be a non-empty array")
+    if dimension < 1:
+        raise ValueError("dimension must be positive")
+
+    model, device = cached_model(model_name, requested_device)
+    import numpy as np
+
+    vectors: list[list[float]] = []
+    # A single model process handles all batches, which is important for large
+    # knowledge stores: loading a 1.8GB checkpoint once is acceptable; loading it once
+    # per document is not.
+    official = has_official_path(model)
+    for start in range(0, len(raw_inputs), batch_size):
+        raw_batch = raw_inputs[start : start + batch_size]
+        batch = [make_input(item) for item in raw_batch]
+        has_images = any(isinstance(item, dict) and item.get("image") for item in raw_batch)
+        if official and not has_images:
+            # The common case: a whole batch of text through the official path.
+            encoded = encode_official_text(model, [str(item) for item in batch], device)
+        elif official:
+            # A batch may mix ordinary text with an image-bearing unit. The low-level
+            # multimodal processor requires an image for every item, so keep the model
+            # loaded once but route each item through the path it needs.
+            rows = []
+            for raw_item, text_item in zip(raw_batch, batch):
+                if isinstance(raw_item, dict) and raw_item.get("image"):
+                    item_encoded = np.asarray(encode_multimodal(model, [raw_item], device), dtype=np.float32)
+                else:
+                    item_encoded = np.asarray(encode_official_text(model, [str(text_item)], device), dtype=np.float32)
+                if item_encoded.ndim == 1:
+                    item_encoded = item_encoded.reshape(1, -1)
+                rows.extend(item_encoded.tolist())
+            encoded = rows
+        else:
+            # A generic SentenceTransformers checkpoint: there is no WeMM head to use,
+            # so the wrapper's own encode is the only option. `has_official_path`
+            # documents why that path cannot be trusted for WeMM itself.
+            encoded = model.encode(
+                batch,
+                batch_size=batch_size,
+                normalize_embeddings=True,
+                truncate_dim=dimension,
+                convert_to_numpy=True,
+                show_progress_bar=False,
+            )
+        array = np.asarray(encoded, dtype=np.float32)
+        if array.ndim == 1:
+            array = array.reshape(1, -1)
+        # The explicit multimodal WeMM path returns the native 2048-d vector; the
+        # generic SentenceTransformers path can apply MRL truncation for us. Keep both
+        # paths identical by applying the requested dimension here too.
+        if array.ndim == 2 and array.shape[1] > dimension:
+            array = array[:, :dimension]
+        if array.ndim != 2 or array.shape[1] != dimension:
+            raise RuntimeError(f"embedding shape mismatch: expected (*, {dimension}), got {tuple(array.shape)}")
+        # The model promises normalized vectors, but re-normalize after MRL truncation
+        # so cosine / int8 scoring stays mathematically consistent.
+        norms = np.linalg.norm(array, axis=1, keepdims=True)
+        array = array / np.maximum(norms, 1e-12)
+        vectors.extend(array.tolist())
+    return {"ok": True, "device": device, "dimension": dimension, "vectors": vectors}
+
+
+def reranker_health(model) -> dict[str, Any]:
+    """Refuse a cross-encoder whose tokenizer has no vocabulary.
+
+    Runs once per loaded model. Both checks are cheap and catch the same failure
+    mode from two directions: the vocabulary size (a *present* tokenizer) and the
+    probe sentence (a tokenizer that is present but resolves nothing).
+    """
+    tokenizer = getattr(model, "tokenizer", None)
+    vocab_size = int(getattr(tokenizer, "vocab_size", 0) or 0)
+    if vocab_size < RERANK_MIN_VOCAB:
+        raise RuntimeError(
+            f"RERANK_TOKENIZER_DEGENERATE: vocab_size={vocab_size} (< {RERANK_MIN_VOCAB}); "
+            "本地模型快照缺少分词器词表（bge-reranker-base 需要 tokenizer.json 或 "
+            "sentencepiece.bpe.model，见 docs/0024-本地模型运行时.md）"
+        )
+    unknown = getattr(tokenizer, "unk_token_id", None)
+    special = set(getattr(tokenizer, "all_special_ids", None) or [])
+    probe = tokenizer("订单 指标 排查 口径 DAU", add_special_tokens=False)["input_ids"]
+    # 特判 `add_special_tokens` 的两种行为：只排除 <unk> 是不够的，`<s>` / `</s>`
+    # 也不是「真实词表条目」，否则空词表能靠这两个 id 蒙过检查。
+    if not [token for token in probe if token != unknown and token not in special]:
+        raise RuntimeError("RERANK_TOKENIZER_DEGENERATE: 探针句全部落到 <unk>，词表里没有真实条目")
+    return {"vocabSize": vocab_size}
+
+
+def predict_raw(model, pairs):
+    """Cross-encoder 分数，**在原始 logit 刻度上**。
+
+    `CrossEncoder.predict` 对 `num_labels == 1` 的模型默认套 sigmoid，返回值落在
+    (0, 1)；而 `rerankThreshold` 的默认值 0.2 是 WeKnora 在 **logit** 刻度上定的。
+    两个刻度混用会让阈值形同虚设：实测 sigmoid 输出恒 > 0.2，于是 `dropped` 永远是 0，
+    「重排阈值」这个绝对闸门从来没生效过；Host 再 sigmoid 一次做展示，又把无关文档
+    （logit −10.2 → 0.0000454）显示成 0.5。
+
+    展示用的那一次 sigmoid 由 Host 做（`recall-rank.ts`），这里只给 logit ——
+    阈值与排序都留在一个可以跨查询比较的刻度上。
+    """
+    try:
+        import torch
+    except Exception as exc:  # pragma: no cover - torch is required to have loaded the model
+        raise RuntimeError(f"RERANK_TORCH_MISSING: {exc}") from exc
+    try:
+        return model.predict(pairs, show_progress_bar=False, activation_fn=torch.nn.Identity())
+    except TypeError as exc:
+        raise RuntimeError(
+            "RERANK_ACTIVATION_UNSUPPORTED: 需要 sentence-transformers 支持 predict(activation_fn=...) "
+            "才能拿到原始 logit（阈值刻度）"
+        ) from exc
+
+
+def load_reranker(model_name: str, requested_device: str):
+    """Load a BERT cross-encoder for stage 4 rerank.
+
+    Separate from the embedding model: the reranker scores (query, note) pairs
+    instead of producing vectors, and lives in its own Host-side worker process
+    so a missing reranker model cannot disable embeddings.
+    """
+    from sentence_transformers import CrossEncoder
+
+    device = choose_device(requested_device)
+
+    def attempt(local_only: bool):
+        kwargs = {"local_files_only": True} if local_only else {}
+        try:
+            return CrossEncoder(model_name, device=device, trust_remote_code=True, **kwargs), device
+        except TypeError:
+            # Older sentence-transformers releases do not expose trust_remote_code.
+            return CrossEncoder(model_name, device=device, **kwargs), device
+
+    return load_local_first(attempt)
+
+
+def cached_reranker(model_name: str, requested_device: str):
+    key = (model_name, requested_device)
+    cached = RERANK_CACHE.get(key)
+    if cached is not None:
+        return cached
+    with redirect_stdout(sys.stderr):
+        model, device = load_reranker(model_name, requested_device)
+        health = reranker_health(model)
+    RERANK_CACHE[key] = (model, device, health)
+    return RERANK_CACHE[key]
+
+
+def rerank(request: dict[str, Any]) -> dict[str, Any]:
+    model_name = str(request.get("model") or "BAAI/bge-reranker-base")
+    requested_device = str(request.get("device") or "auto")
+    query = str(request.get("query") or "")
+    documents = request.get("documents")
+    if not query:
+        raise ValueError("rerank requires a non-empty query")
+    if not isinstance(documents, list) or not documents:
+        raise ValueError("rerank requires a non-empty documents array")
+    model, device, health = cached_reranker(model_name, requested_device)
+    pairs = [(query, str(document)) for document in documents]
+    with redirect_stdout(sys.stderr):
+        scores = predict_raw(model, pairs)
+    return {
+        "ok": True,
+        "device": device,
+        "vocabSize": health.get("vocabSize"),
+        "scores": [float(score) for score in scores],
+    }
+
+
+def main() -> None:
+    for line in sys.stdin:
+        if not line.strip():
+            continue
+        try:
+            request = json.loads(line)
+            # Stage 4 rerank and stage 2/3 embedding share one worker script so
+            # the plugin keeps a single local Python runtime; the `task` field
+            # selects the path and both models are cached per process.
+            if str(request.get("task") or "") == "rerank":
+                emit(rerank(request))
+            else:
+                emit(encode(request))
+        except Exception as exc:  # return machine-readable errors to the Host
+            emit({"ok": False, "error": f"{type(exc).__name__}: {exc}"})
+
+
+if __name__ == "__main__":
+    main()