#!/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()