| |
| |
| """ |
| 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 |
|
|
| |
| mx = None |
|
|
| 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: |
| 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: |
| |
| 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 |
| |
| 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 = input_ids.shape |
| scales_g = self.scales[ids] |
| 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) |
| n_tokens = core.sum(mask_f, axis=1) |
| n_tokens = core.maximum(n_tokens, core.array([1.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) |
| for i, v in zip(batch_idx, out): |
| outs[i] = v |
| |
| 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", |
| "København er hovedstaden i Danmark", |
| "Le machine learning change le monde", |
| "机器学习正在改变世界", |
| "機械学習は世界を変えつつある", |
| "¿Cómo llego al aeropuerto desde aquí?", |
| "How do I get to the airport from here?", |
| "Was kostet ein Ticket zum Hauptbahnhof?", |
| ] |
|
|
|
|
| def to_numpy(x): |
| """MLX array -> numpy。优先用 __array__ 互操作，老版本退回 tolist()。""" |
| try: |
| return np.asarray(x) |
| except Exception: |
| 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() |
|
|
