scripts/convert_minicpm5_tjs.py

#!/usr/bin/env python3
"""Convert openbmb/MiniCPM5-1B to a Transformers.js q4 ONNX repo."""

from __future__ import annotations

import argparse
import json
import logging
import os
import shutil
import subprocess
import sys
import tarfile
import tempfile
import urllib.request
from dataclasses import asdict
from pathlib import Path

from huggingface_hub import HfApi, create_repo
from optimum.exporters.onnx import main_export
from transformers import AutoConfig


SOURCE_MODEL = "openbmb/MiniCPM5-1B"
TRANSFORMERS_JS_TAG = "3.8.1"
TRANSFORMERS_JS_TARBALL = (
    f"https://github.com/huggingface/transformers.js/archive/refs/tags/{TRANSFORMERS_JS_TAG}.tar.gz"
)


def run(cmd: list[str], cwd: Path | None = None) -> None:
    print("+", " ".join(cmd), flush=True)
    subprocess.run(cmd, cwd=cwd, check=True)


def download_transformers_js(work_dir: Path) -> Path:
    archive_path = work_dir / "transformers.js.tar.gz"
    urllib.request.urlretrieve(TRANSFORMERS_JS_TARBALL, archive_path)
    with tarfile.open(archive_path) as archive:
        archive.extractall(work_dir)
    return work_dir / f"transformers.js-{TRANSFORMERS_JS_TAG}"


def patch_config(config_path: Path, dtype: str, q4_external_chunks: int) -> None:
    config = json.loads(config_path.read_text())
    config.setdefault("transformers.js_config", {})
    config["transformers.js_config"]["dtype"] = dtype
    if q4_external_chunks:
        config["transformers.js_config"]["use_external_data_format"] = {
            "model_q4.onnx": q4_external_chunks,
        }
    else:
        config["transformers.js_config"].pop("use_external_data_format", None)
    config_path.write_text(json.dumps(config, indent=2) + "\n")


def write_readme(output_dir: Path, source_model: str, target_repo: str) -> None:
    readme = f"""---
license: apache-2.0
library_name: transformers.js
pipeline_tag: text-generation
base_model: {source_model}
tags:
- transformers.js
- onnx
- onnxruntime-web
- llama
- minicpm5
- text-generation
- browser
- webgpu
---

# MiniCPM5-1B ONNX Web

Transformers.js q4 ONNX export of `{source_model}` for browser text generation.

## Files

- `onnx/model_q4.onnx`: ONNX Runtime 4-bit MatMul quantized decoder with KV cache.
- `config.json`: includes `transformers.js_config.dtype = "q4"` so Transformers.js loads the q4 artifact by default.
- tokenizer and generation config files copied from the source model export.

## Usage

```js
import {{ pipeline }} from "@huggingface/transformers";

const generator = await pipeline("text-generation", "{target_repo}", {{
  dtype: "q4",
  device: "webgpu",
}});

const output = await generator("Briefly introduce yourself.", {{
  max_new_tokens: 64,
  temperature: 0.2,
  do_sample: true,
}});
console.log(output[0].generated_text);
```

If WebGPU is unavailable, use `device: "wasm"` in the browser.
"""
    (output_dir / "README.md").write_text(readme)


def convert(args: argparse.Namespace) -> Path:
    work_dir = Path(args.work_dir or tempfile.mkdtemp(prefix="minicpm5-tjs-")).resolve()
    work_dir.mkdir(parents=True, exist_ok=True)
    print(f"Working directory: {work_dir}", flush=True)

    transformers_js_dir = download_transformers_js(work_dir)
    scripts_dir = transformers_js_dir / "scripts"
    sys.path.insert(0, str(transformers_js_dir))

    from scripts.quantize import QuantizationArguments, quantize
    logging.getLogger("onnxruntime.quantization.matmul_4bits_quantizer").setLevel(logging.WARNING)

    export_root = work_dir / "export"
    model_dir = export_root / args.source_model
    model_dir.mkdir(parents=True, exist_ok=True)

    device = args.device
    if device == "auto":
        try:
            import torch

            device = "cuda" if torch.cuda.is_available() else "cpu"
        except Exception:
            device = "cpu"
    print(f"Export device: {device}", flush=True)

    config = AutoConfig.from_pretrained(args.source_model)
    print(
        "Source config:",
        json.dumps(
            {
                "model_type": config.model_type,
                "architectures": getattr(config, "architectures", None),
                "hidden_size": getattr(config, "hidden_size", None),
                "num_hidden_layers": getattr(config, "num_hidden_layers", None),
                "torch_dtype": str(getattr(config, "torch_dtype", None)),
            },
            indent=2,
        ),
        flush=True,
    )

    main_export(
        model_name_or_path=args.source_model,
        output=model_dir,
        task="text-generation-with-past",
        opset=args.opset,
        device=device,
        dtype=args.export_dtype,
        do_validation=False,
        trust_remote_code=False,
        library_name="transformers",
        slim=False,
    )

    onnx_dir = model_dir / "onnx"
    onnx_dir.mkdir(exist_ok=True)

    quant_args = QuantizationArguments(
        modes=["q4"],
        per_channel=False,
        reduce_range=False,
        block_size=args.block_size,
        is_symmetric=True,
        accuracy_level=None,
        op_block_list=None,
    )
    quantize(str(model_dir), str(onnx_dir), quant_args)
    (model_dir / "quantize_config.json").write_text(json.dumps(asdict(quant_args), indent=2) + "\n")

    for path in model_dir.glob("*.onnx*"):
        path.unlink()

    q4_model = onnx_dir / "model_q4.onnx"
    if not q4_model.exists():
        raise FileNotFoundError(f"Missing expected quantized model: {q4_model}")

    q4_external_chunks = 1 if (onnx_dir / "model_q4.onnx_data").exists() else 0
    patch_config(model_dir / "config.json", "q4", q4_external_chunks)
    write_readme(model_dir, args.source_model, args.target_repo)

    if args.output_dir:
        final_dir = Path(args.output_dir).resolve()
        if final_dir.exists():
            shutil.rmtree(final_dir)
        shutil.copytree(model_dir, final_dir)
        model_dir = final_dir

    print("Final files:", flush=True)
    for file in sorted(p.relative_to(model_dir).as_posix() for p in model_dir.rglob("*") if p.is_file()):
        print(file, flush=True)

    if args.target_repo:
        token = os.environ.get("HF_TOKEN") or True
        create_repo(args.target_repo, repo_type="model", private=args.private, exist_ok=True, token=token)
        api = HfApi(token=token)
        api.upload_folder(
            repo_id=args.target_repo,
            repo_type="model",
            folder_path=str(model_dir),
            commit_message=f"Add q4 Transformers.js export of {args.source_model}",
        )
        print(f"Uploaded to https://huggingface.co/{args.target_repo}", flush=True)

    return model_dir


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-model", default=SOURCE_MODEL)
    parser.add_argument("--target-repo", default="")
    parser.add_argument("--output-dir", default="")
    parser.add_argument("--work-dir", default="")
    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"])
    parser.add_argument("--export-dtype", default="fp32", choices=["fp32", "fp16"])
    parser.add_argument("--opset", type=int, default=18)
    parser.add_argument("--block-size", type=int, default=32)
    parser.add_argument("--private", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    convert(parse_args())