arlenecc/potion-multilingual-128M-onnx

potion-multilingual-128M — ONNX（FP16 / INT8 / INT4）

把 HF 上的多语言静态嵌入模型 minishlab/potion-multilingual-128M （mean-pool + L2 normalize，vocab 500353 × dim 256，fp32 约 512 MB）转成可直接用 onnxruntime 部署的版本，并提供 FP16 / INT8 / INT4 三个精简版本。

目录	大小	相比 fp32	说明
fp16/	256.18 MB	2.0x	半精度权重（查表后转 fp32 累加），几乎无损、速度最快
int8/	136.10 MB	3.77x	分块对称 INT8（block=32）
int4/	80.07 MB	6.40x	分块非对称 INT4（block=32，2 nibbles/byte 打包）

每个变体目录自包含：model.onnx + tokenizer.json + config.json + example.py。 （原仓库自带的 onnx/model.onnx 是作者官方导出的 fp32 版本，512 MB；本目录的三个版本体积分别为它的 1/2、1/3.8、1/6.4。）

potion-multilingual-128M-onnx/
├── fp16/{model.onnx, tokenizer.json, config.json, example.py}
├── int8/{...}
├── int4/{...}
├── convert.py                        # 转换脚本（含量化实现）
├── sweep.py                          # 量化参数搜索（本模型的默认值由它实测确定）
├── verify.py                         # 精度 / 可用性 / 跨语言检索评测
├── benchmark.py                      # 体积、加载、延迟、吞吐对比
├── eval_data_multilingual_bitext.json  # 2124 条英↔法/德/西/日/阿 平行句对（验证语料）
├── conversion_report.json
├── verification_report.json
├── quantization_sweep.json
├── benchmark_results.json
└── requirements.txt

下载源模型：

python - <<'EOF'
from huggingface_hub import snapshot_download
snapshot_download('minishlab/potion-multilingual-128M', repo_type='model',
                  local_dir='/Users/arlenecc/Downloads/minishlab/potion-multilingual-128M',
                  allow_patterns=['config.json','model.safetensors','tokenizer.json',
                                  'modules.json','tokenizer_config.json','special_tokens_map.json'])
EOF

---

1. 为什么不用官方 onnxruntime 动态量化

前向是 ids → Embedding 查表(Gather) → 序列维求均值 → L2 归一化，图里唯一的权重就是 embedding 矩阵本体， 没有任何 MatMul / Gemm / Conv。quantize_dynamic 只处理 MatMul 类算子，对 Gather 完全不动， 量化后文件大小与精度都不变（即"转换"实际没有生效）。

2. 采用方案：embedding 表分块权重量化 + Gather 后即时反量化

fp16: input_ids → Gather(fp16 表)      → Cast fp32 → mean-pool → normalize
int8: input_ids → Gather(int8 表)      → × scale(per block) → mean-pool → normalize
int4: input_ids → Gather(uint8 打包表) → 拆 nibble → × scale + bias → mean-pool → normalize

• 只反量化被命中的 token 行，计算量仍与 token 数成正比，不会退化成 vocab × dim 的大 MatMul；
• 按行切 block，每块单独算 scale（INT4 非对称另带 bias），scale/bias 用 fp16 存，元数据开销 <7%；
• INT4 相邻两个 4bit 打包进一个 byte，取出后用纯浮点算术拆包，不依赖 BitwiseAnd/BitShift；
• 图里没有任何动态 Reshape（per-block scale 用 Gather + 常量索引沿最后一维展开），整图保持静态 rank， onnx.checker(full_check=True) 可通过。opset = 18，输入输出：input_ids / attention_mask （int64，动态 batch × seq）→ embeddings（float32，[batch, 256]，已 L2 归一化）。

3. 量化参数（sweep.py 在 1800 句多语言真实语料上实测）

bits	block	scheme	size(MB)	压缩	cosine mean	cosine min
int8	row	sym	129.09	3.97x	0.999978	0.999939
int8	32	sym	136.10	3.77x	0.999991	0.999981
int8	64	sym	132.09	3.88x	0.999989	0.999974
int8	128	sym	130.09	3.94x	0.999985	0.999960
int4	row	sym	65.05	7.88x	0.992917	0.979802
int4	row	asym	66.05	7.76x	0.995074	0.987909
int4	32	sym	72.05	7.11x	0.997057	0.993861
int4	32	asym	80.06	6.40x	0.998110	0.996149
int4	64	asym	72.05	7.11x	0.997362	0.994258
int4	128	asym	68.05	7.53x	0.996119	0.990986

→ 默认值 INT8 block=32 对称 / INT4 block=32 非对称，与上一轮 128M 同规模模型结论一致。 想要更高保真度可 --int4-block 16。

---

4. 验证结果

verify.py 在 4248 句真实多语言语料上做端到端校验：2124 条平行句对 （英↔法/德/西/日/阿，取自 sentence-transformers/parallel-sentences-news-commentary）， 语料已随目录保存，结果可完全复现。

4.1 一致性 / 可用性

检查项	结果
分词与 model2vec 官方一致	✅ 4248/4248 句 token ids 完全一致
onnx.checker(full_check) + ORT 加载	✅ 三个模型全部通过
边界输入（空串 / 单 token / CJK / emoji 特殊符号 / 超长 / 混合 padding batch）	✅ 无 NaN / 无 Inf，输出均为单位向量（‖v‖偏差 ≤ 1.8e-07）
跨语言语义自检（example.py demo）	✅ 英/西同义疑问句 0.789，中/日同义句 0.676，无关句 ≈ 0

4.2 精度损失（对照官方 model2vec fp32 实现）

变体	cosine mean	cosine min	cosine p1	max|Δ|	mean|Δ|
fp16	1.000000	1.000000	1.000000	1.17e-04	~1e-06
int8	0.999991	0.999980	0.999984	1.97e-03	9e-05
int4	0.998050	0.995969	0.996610	2.59e-02	1.3e-03

4.3 跨语言检索（query = 英文句子，候选 = 目标语言的 500 条译文，正确项为其对应译文）

模型	R@1 均值	R@5	R@10	MRR@10	en-fr	en-de	en-es	en-ja	en-ar
model2vec 官方（参考）	0.7582	0.8853	0.9163	0.8124	0.836	0.822	0.858	0.594	0.681
fp16 ONNX	0.7582	0.8853	0.9163	0.8124	0.836	0.822	0.858	0.594	0.681
int8 ONNX	0.7589	0.8853	0.9163	0.8126	0.836	0.822	0.858	0.597	0.681
int4 ONNX	0.7587	0.8835	0.9162	0.8121	0.832	0.826	0.860	0.601	0.675

排序偏移（5 个语言对汇总，对比官方参考实现）：

变体	名次完全一致率	Top-1 一致率	R@1 掉点	MRR@10 掉点
fp16	0.9971	1.0000	0.0000	0.0000
int8	0.9671	0.9993	-0.0007	-0.0001
int4	0.8780	0.9921	-0.0005	+0.0003

→ 三个变体在真实跨语言检索任务上的指标全部在噪声范围内（|Δ| < 0.001）， fp16 与官方结果逐位一致，int8/int4 也只掉千分之一以内。

5. 性能（MacBook x86_64 CPU，16GB，ORT 1.23 默认线程；3 次运行中位）

变体	大小	加载	单句延迟	QPS	短句吞吐*
fp16	256.18 MB	2.4 s	0.144 ms	~7000	~21k–40k 句/秒
int8	136.10 MB	2.7 s	0.162 ms	~6160	~32k–36k 句/秒
int4	80.07 MB	2.4 s	0.201 ms	~5000	~27k 句/秒
model2vec 官方 numpy	512.36 MB	6.6 s	0.212 ms	~4800	~39k 句/秒

* 吞吐类指标受 CPU 调度影响波动较大（单次 ±30%），单句延迟更稳定：延迟排序 fp16 ≲ int8 < int4 < 官方实现。 这个模型词表达 50 万、表比上一个模型大得多，访存压力上升，所以 int8 凭借 136MB vs 256MB 的内存优势 追平了 fp16 的算术优势，两者延迟差距已进入抖动范围。

怎么选：**精度/速度优先 → fp16/（几乎无损）；综合性价比 → int8/**（体积小一半、延迟相当、指标不掉）； **极致体积（端侧 / 容器镜像 / 冷启动）→ int4/**。

---

6. 使用

6.1 Python

from example import StaticEmbeddingONNX   # 位于任意变体目录内

model = StaticEmbeddingONNX.from_dir("int8")        # 或 "fp16" / "int4"
vec = model.encode("Bonjour le monde")              # -> (256,) 单位向量
mat = model.encode(["Hello world", "Bonjour le monde", "你好世界"])
sim = float(mat[0] @ mat[1])                        # 余弦相似度（跨语言亦可比较）

命令行 demo：

python fp16/example.py
python int8/example.py "What time does the train leave?" "¿A qué hora sale el tren?"

6.2 裸 onnxruntime（不依赖本目录代码）

import numpy as np, onnxruntime as ort
from tokenizers import Tokenizer

tok = Tokenizer.from_file("int8/tokenizer.json"); tok.enable_truncation(512)
ids = [i for i in tok.encode("Bonjour le monde", add_special_tokens=False).ids if i != 1]  # 1 = <unk>

x = np.array([ids], dtype=np.int64); mask = np.ones_like(x)
sess = ort.InferenceSession("int8/model.onnx", providers=["CPUExecutionProvider"])
emb = sess.run(None, {"input_ids": x, "attention_mask": mask})[0]   # (1, 256)，已 L2 归一化

⚠️ 必须对齐的 3 个细节（否则结果与原模型不一致）：

1. add_special_tokens=False；
2. 丢弃 unk token（id = 1）或用 attention_mask 屏蔽；
3. 超长文本需截断：max_length=512 token，且 model2vec 还会先按字符截到 512 × 7 = 3584 个字符 （median_token_length 由 config.json 给出）。 以上逻辑都封装在 example.py 的 Model2VecTokenizer 里，直接调用即可。

---

7. 复现

pip install -r requirements.txt

# 1) 转换（默认导出 fp16 / int8 / int4）
python convert.py

# 2) 量化参数搜索（可选）
python sweep.py --limit 900

# 3) 验证：分词对齐 + 数值保真 + 边界用例 + 跨语言检索评测
python verify.py

# 4) 性能对比
python benchmark.py

参数	默认	说明
--variants	fp16 int8 int4	导出哪些变体（fp32 版本同样支持，512 MB）
--int8-block	32	INT8 分块大小，0 = 整行一块（体积最小、精度略差）
--int4-block	32	INT4 分块大小（需整除 dim/2 = 128）
--int4-sym	关	改用对称 INT4（80 MB → 72 MB，精度下降）

8. 依赖

numpy>=1.24
onnx>=1.15
onnxruntime>=1.17
tokenizers>=0.15
safetensors>=0.4        # 仅 convert.py 读取源权重需要
model2vec>=0.4          # 仅 verify.py / benchmark.py 的对拍基线需要（缺失会自动跳过）

转换过程不需要 PyTorch / sentence-transformers。源模型的 tokenizer.json 已是 HF tokenizers 格式，无需再做 TF → HF 词表转换。

9. 已知限制

• normalize=true 已烘焙进图里；输出维度固定 256；超过 512 token 的文本会被截断。
• config.json 里 INT4 的 row_cosine_min = 0.0 是统计假象：源词表中 [PAD]（id 0）这一行的 embedding 本来就是全零，0/0 记为 0；该行量化后仍为全零，不影响任何结果。
• INT4 的 embedding 表在导出时做过列重排（偶数维/奇数维分离以便打包），通过图推理不受影响； 若直接用外部工具读取 int4 的 Wq 初值，需按 config.json 记录的量化元信息还原。
• fp16/int8/int4 的体积收益来自 embedding 表；推理时的激活部分始终是 fp32， 因此延迟收益主要来自访存减半而非计算加速。