Blame for _example_template.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
example.py — MLX 版 model2vec 静态嵌入模型的最小可用示例（任一变体目录下都有同一份）。

运行环境：Apple Silicon (M 系列) + macOS，安装 `pip install mlx`。
         MLX 不支持 Intel Mac / Linux / Windows。

运行:
    python example.py
    python example.py "Hello world" "Bonjour le monde" "你好世界"

API:
    from example import StaticEmbeddingMLX
    model = StaticEmbeddingMLX.from_dir(".")      # 自动加载同目录的 weights.npz / config.json
    vec = model.encode("Bonjour le monde")        # -> mx.array, shape (256,)
    mat = model.encode(["en", "fr", "zh"])        # -> (3, 256)
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np

# MLX 只支持 Apple Silicon；这里延迟到真正推理时才 import，便于在其它机器上做静态检查。
mx = None  # 由 _load_mlx() 注入

HERE = Path(__file__).resolve().parent


def _load_mlx():
    """延迟导入 mlx.core，并给出清晰的错误提示。"""
    global mx
    if mx is None:
        try:
            import mlx.core as _mx
        except ImportError as exc:  # pragma: no cover - 仅非 Apple Silicon 上触发
            raise SystemExit(
                "未能导入 mlx.core。MLX 需要 Apple Silicon (M 系列) + macOS，"
                "请先 `pip install mlx`。Intel Mac / Linux / Windows 无法运行。"
            ) from exc
        mx = _mx
    return mx


class Model2VecTokenizer:
    """复刻 model2vec.StaticModel.tokenize：add_special_tokens=False + 丢弃 unk + 截断。

    model2vec 在 tokenize 前会先按字符粗截一遍（max_length * median_token_length），
    再做 token 级截断；两者都在这里复刻，保证与官方 Hub 模型的输出对齐。
    """

    def __init__(self, tokenizer_path: Path, config: dict | None = None):
        cfg = config or {}
        from tokenizers import Tokenizer

        self.tokenizer = Tokenizer.from_file(str(tokenizer_path))
        if self.tokenizer.padding is not None:
            self.tokenizer.no_padding()
        self.max_length = cfg.get("max_length")
        self.median_token_length = int(cfg.get("median_token_length") or 0)
        self.unk_token_id = cfg.get("unk_token_id")
        self._char_limit = (
            None if (self.max_length is None or not self.median_token_length)
            else self.max_length * self.median_token_length
        )

    def __call__(self, sentences: list[str]) -> list[list[int]]:
        if self.max_length is not None:
            self.tokenizer.enable_truncation(self.max_length)
        else:
            self.tokenizer.no_truncation()
        texts = [s[: self._char_limit] for s in sentences] if self._char_limit else list(sentences)
        ids = [enc.ids for enc in self.tokenizer.encode_batch_fast(texts, add_special_tokens=False)]
        if self.unk_token_id is not None:
            ids = [[i for i in row if i != self.unk_token_id] for row in ids]
        return ids


class StaticEmbeddingMLX:
    """量化 model2vec 的 MLX 实现：Gather → 反量化 → 均值池化 → L2 归一化。

    量化布局（与 convert_mlx.py 一一对应）:
      INT8: embeddings int8   [V, D]          , 对称:   x = q * scale
            scales     fp16   [V, D/group]
      INT4: embeddings uint8  [V, D/2]        , 非对称: x = q * scale + bias
            scales     fp16   [V, D/group]      低 4 位 = 偶数维, 高 4 位 = 奇数维
            biases     fp16   [V, D/group]

    这里刻意只用最基础的 MLX 算子（索引/reshape/逐元素运算），不依赖
    nn.QuantizedEmbedding 的内部打包格式，避免 mlx 版本差异导致加载失败。
    """

    def __init__(self, weights_npz: Path, config: dict, tokenizer_path: Path | None = None):
        self.config = config
        self.bits = int(config["quantization"]["bits"])
        self.group_size = int(config["quantization"]["group_size"])
        self.dim = int(config["dim"])
        self.vocab_size = int(config["vocab_size"])
        core = _load_mlx()
        arrays = core.load(str(weights_npz))
        self.embeddings = arrays["embeddings"]
        self.scales = arrays["scales"]
        self.biases = arrays.get("biases") if hasattr(arrays, "get") else None
        self.tokenizer = Model2VecTokenizer(tokenizer_path or Path(weights_npz).parent / "tokenizer.json", config)

    @classmethod
    def from_dir(cls, directory) -> "StaticEmbeddingMLX":
        d = Path(directory)
        cfg = json.loads((d / "config.json").read_text(encoding="utf-8"))
        return cls(d / "weights.npz", cfg, d / "tokenizer.json")

    # ------------------------------------------------------------------ #
    def _dequantize(self, rows, scales_g, biases_g, batch: int, seq: int):
        """把查表得到的量化行还原成 float32 的 token 向量 [batch, seq, dim]。"""
        core = mx
        x = rows.astype(core.float32)
        if self.bits == 4:
            # 低 4 位 / 高 4 位拆包（纯浮点算术，避免依赖整数取模语义）
            hi = core.floor(x / 16.0)
            lo = x - 16.0 * hi
            x = core.stack([lo, hi], axis=-1).reshape(batch, seq, self.dim)
        groups = self.dim // self.group_size
        # [batch, seq, dim] -> [batch, seq, groups, group_size] 后按组广播 scale/bias
        x = x.reshape(batch, seq, groups, self.group_size)
        x = x * scales_g.astype(core.float32)[..., None]
        if biases_g is not None:
            x = x + biases_g.astype(core.float32)[..., None]
        return x.reshape(batch, seq, self.dim)

    def embed_ids(self, input_ids: np.ndarray, attention_mask: np.ndarray):
        """input_ids / attention_mask: numpy int64 [batch, seq] -> [batch, dim]。"""
        core = mx
        ids = core.array(input_ids)
        rows = self.embeddings[ids]                      # [batch, seq, storage_dim]
        batch, seq = input_ids.shape
        scales_g = self.scales[ids]                      # [batch, seq, groups]
        biases_g = None if self.biases is None else self.biases[ids]
        x = self._dequantize(rows, scales_g, biases_g, batch, seq)

        mask_f = core.array(attention_mask.astype(np.float32))
        masked = x * mask_f.reshape(batch, seq, 1)
        token_sum = core.sum(masked, axis=1)                            # [batch, dim]
        n_tokens = core.sum(mask_f, axis=1)                             # [batch]
        n_tokens = core.maximum(n_tokens, core.array([1.0]))            # 防止 0 除
        mean = token_sum / n_tokens.reshape(batch, 1)

        n2 = core.sum(mean * mean, axis=1, keepdims=True)
        norm = core.sqrt(core.maximum(n2, core.array([[1e-12]])))
        return mean / norm

    def encode(self, sentences, batch_size: int = 256, sort_by_length: bool = True):
        """把句子编码成 L2 归一化的向量，返回 mx.array。"""
        mx_local = _load_mlx()
        single = isinstance(sentences, str)
        texts = [sentences] if single else list(sentences)
        ids = self.tokenizer(texts)
        order = sorted(range(len(texts)), key=lambda i: len(ids[i])) if sort_by_length else list(range(len(texts)))
        outs = {}
        for start in range(0, len(order), batch_size):
            batch_idx = order[start : start + batch_size]
            rows = [ids[i] for i in batch_idx]
            width = max([len(r) for r in rows] + [1])
            arr_ids = np.zeros((len(rows), width), dtype=np.int64)
            arr_mask = np.zeros((len(rows), width), dtype=np.int64)
            for i, r in enumerate(rows):
                if r:
                    arr_ids[i, : len(r)] = r
                    arr_mask[i, : len(r)] = 1
            out = self.embed_ids(arr_ids, arr_mask)          # mx.array [batch, dim]
            for i, v in zip(batch_idx, out):
                outs[i] = v
        # 全程保持在 MLX 张量域内拼结果，避免依赖不同版本的 MLX <-> numpy 互操作
        stacked = mx_local.stack([outs[i] for i in range(len(texts))], axis=0)
        return stacked[0] if single else stacked


DEMO_SENTENCES = [
    "The capital of Denmark is Copenhagen",  # en
    "København er hovedstaden i Danmark",  # da
    "Le machine learning change le monde",  # fr
    "机器学习正在改变世界",  # zh
    "機械学習は世界を変えつつある",  # ja
    "¿Cómo llego al aeropuerto desde aquí?",  # es
    "How do I get to the airport from here?",  # en
    "Was kostet ein Ticket zum Hauptbahnhof?",  # de
]


def to_numpy(x):
    """MLX array -> numpy。优先用 __array__ 互操作，老版本退回 tolist()。"""
    try:
        return np.asarray(x)
    except Exception:  # noqa: BLE001
        return np.array(x.tolist(), dtype=np.float32)


def main():
    model = StaticEmbeddingMLX.from_dir(HERE)
    cfg = model.config
    size_mb = (HERE / "weights.npz").stat().st_size / 1e6
    print(f"权重: {HERE / 'weights.npz'} ({size_mb:.2f} MB)")
    print(f"配置: variant={cfg.get('variant')} bits={model.bits} group_size={model.group_size} "
          f"dim={model.dim} vocab={model.vocab_size} scheme={cfg['quantization'].get('scheme')}")

    emb = to_numpy(model.encode(DEMO_SENTENCES))
    print(f"\n向量 shape={emb.shape}, ||v||={np.linalg.norm(emb[0]):.6f}")
    sim = emb @ emb.T
    print("\n余弦相似度矩阵:")
    print("            " + " ".join(f"{i:>6}" for i in range(len(DEMO_SENTENCES))))
    for i, row in enumerate(sim):
        print(f"[{i}] " + " ".join(f"{v:6.3f}" for v in row) + f"   {DEMO_SENTENCES[i][:40]}")
    off = sim[~np.eye(len(sim), dtype=bool)]
    print(f"\n非对角相似度范围: [{off.min():.3f}, {off.max():.3f}]")


if __name__ == "__main__":
    main()