| |
| """Exports Qwen2.5-{1.5B,3B}-Instruct to int4 ONNX models that ONNX Runtime runs on Android and iOS. |
| |
| The ONNX Runtime GenAI model builder writes the graph: GroupQueryAttention with a KV cache, 4-bit MatMulNBits |
| (blocks of 32, accuracy_level 4) and the tied embedding as a 4-bit GatherBlockQuantized that shares its weights |
| with the LM head. The LM head only computes the last position's logits. Next to the model it writes the tokenizer |
| as vocab.txt and merges.txt, the plain files the Kotlin and Swift tokenizers read. |
| |
| Usage: |
| .venv/bin/python -I export.py [--size 1.5B|3B ...] |
| """ |
|
|
| import argparse |
| import json |
| import shutil |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import onnx |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| from common import HERE, SIZES, out_dir, src_dir |
|
|
| EXTRA_OPTIONS = [ |
| "shared_embeddings=true", |
| "prune_lm_head=true", |
| "block_size=32", |
| "accuracy_level=4", |
| "hf_remote=false", |
| ] |
|
|
|
|
| def build(size: str) -> Path: |
| out = out_dir(size) |
| onnx_dir = out / "onnx" |
| if onnx_dir.exists(): |
| shutil.rmtree(onnx_dir) |
| cmd = [ |
| sys.executable, "-m", "onnxruntime_genai.models.builder", |
| "-i", str(src_dir(size)), "-o", str(onnx_dir), "-p", "int4", "-e", "cpu", |
| "-c", str(HERE / "data" / "cache"), "--extra_options", *EXTRA_OPTIONS, |
| ] |
| subprocess.run(cmd, check=True) |
| return onnx_dir |
|
|
|
|
| def write_tokenizer(size: str) -> None: |
| """vocab.txt: one byte-level token per line, by id (added tokens included). merges.txt: one 'a b' pair per line, by rank.""" |
| tok = json.loads((src_dir(size) / "tokenizer.json").read_text(encoding="utf-8")) |
| model = tok["model"] |
| assert model["type"] == "BPE" and not model.get("byte_fallback"), model["type"] |
| vocab = dict(model["vocab"]) |
| for added in tok["added_tokens"]: |
| vocab[added["content"]] = added["id"] |
| ids = sorted(vocab.values()) |
| assert ids == list(range(len(ids))), "token ids must be contiguous" |
| by_id = {i: t for t, i in vocab.items()} |
| for token in by_id.values(): |
| assert "\n" not in token and "\r" not in token, repr(token) |
| out = out_dir(size) |
| (out / "vocab.txt").write_text("".join(by_id[i] + "\n" for i in ids), encoding="utf-8") |
| merges = [m if isinstance(m, str) else " ".join(m) for m in model["merges"]] |
| for m in merges: |
| assert m.count(" ") == 1, repr(m) |
| (out / "merges.txt").write_text("".join(m + "\n" for m in merges), encoding="utf-8") |
| shutil.copy(src_dir(size) / "LICENSE", out / "LICENSE") |
|
|
|
|
| def describe(onnx_dir: Path) -> None: |
| model = onnx.load(onnx_dir / "model.onnx", load_external_data=False) |
|
|
| def shape(v): |
| return [d.dim_param or d.dim_value for d in v.type.tensor_type.shape.dim] |
|
|
| for v in list(model.graph.input)[:5] + list(model.graph.output)[:3]: |
| print(f" {v.name}: {onnx.TensorProto.DataType.Name(v.type.tensor_type.elem_type)} {shape(v)}") |
| ops = sorted({n.op_type for n in model.graph.node}) |
| print(" ops:", ", ".join(ops)) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--size", choices=SIZES, nargs="+", default=list(SIZES)) |
| args = parser.parse_args() |
| for size in args.size: |
| onnx_dir = build(size) |
| write_tokenizer(size) |
| print(size) |
| describe(onnx_dir) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
