scripts/build_and_test_onnx.py

#!/usr/bin/env python3
"""
Build and locally TEST ONNX variants of SaT models (CPU, runs on Linux).

For each base model x vocab variant:
  - export to ONNX (fp32)
  - dynamic-quantize to int8
  - measure on-disk size
  - run REAL segmentation (max_length=80) via onnxruntime and compare the
    predicted chunk boundaries against the original full-vocab torch model.

This is the local stand-in for the Core ML models (which can't run on Linux).
It validates the pruning + quantization end-to-end on actual text.

Run with the conda libstdc++ preloaded so onnxruntime/numpy import:
  LD_PRELOAD=$CONDA_PREFIX/lib/libstdc++.so.6 python scripts/build_and_test_onnx.py
"""
import argparse
import sys
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent))

import wtpsplit.models  # noqa: F401
from transformers import AutoModelForTokenClassification, AutoTokenizer
from wtpsplit.utils import token_to_char_probs
from wtpsplit.utils.constraints import constrained_segmentation
from wtpsplit.utils.priors import create_prior_function
from build_ios_coreml import compute_keep_ids, prune_embedding, LogitsWrapper

NEWLINE_INDEX = 0
SPECIALS = {0, 1, 2}  # bos, pad, eos for xlm-r (after remap these shift; handled below)

TEST_TEXTS = [
    "Breaking News: Scientists at CERN have announced a groundbreaking discovery "
    "that could revolutionize our understanding of particle physics. The team observed "
    "unexpected behavior in proton collisions at energies never before achieved.",
    "This is English. 这是中文。Mixed text works too! 混合文本也可以正常处理。",
]


def export_onnx(wrapper, seq_len, path):
    ids = torch.ones((1, seq_len), dtype=torch.long)
    mask = torch.ones((1, seq_len), dtype=torch.long)
    torch.onnx.export(
        wrapper, (ids, mask), str(path),
        input_names=["input_ids", "attention_mask"], output_names=["logits"],
        dynamic_axes={"input_ids": {0: "b", 1: "s"},
                      "attention_mask": {0: "b", 1: "s"},
                      "logits": {0: "b", 1: "s"}},
        opset_version=17, do_constant_folding=True,
    )


def char_boundary_probs(session_or_model, tokenizer, text, remap=None, unk_new=None, is_onnx=True):
    enc = tokenizer([text], return_offsets_mapping=True, add_special_tokens=True)
    ids = np.array(enc["input_ids"], dtype=np.int64)
    mask = np.array(enc["attention_mask"], dtype=np.int64)
    feed_ids = ids.copy()
    if remap is not None:
        feed_ids = remap[ids]
        feed_ids[feed_ids == -1] = unk_new
    if is_onnx:
        import onnxruntime as ort  # noqa
        logits = session_or_model.run(
            ["logits"], {"input_ids": feed_ids.astype(np.int64),
                         "attention_mask": mask.astype(np.int64)})[0]
    else:
        with torch.no_grad():
            logits = session_or_model(input_ids=torch.tensor(feed_ids),
                                      attention_mask=torch.tensor(mask),
                                      return_dict=True).logits.numpy()
    # token_logits: (seq, 1); map to char probs using offsets
    token_logits = logits[0]
    offsets = enc["offset_mapping"][0]
    tokens = tokenizer.convert_ids_to_tokens(enc["input_ids"][0])
    char_logits = token_to_char_probs(text, tokens, token_logits,
                                      tokenizer.all_special_tokens, offsets)
    probs = 1.0 / (1.0 + np.exp(-char_logits[:, NEWLINE_INDEX]))
    return probs


def segment(probs, text, max_length=80, min_length=40):
    prior = create_prior_function("gaussian",
                                  {"target_length": 70, "spread": 12, "max_length": max_length})
    idx = constrained_segmentation(probs, prior, min_length=min_length,
                                   max_length=max_length, algorithm="viterbi")
    idx = [0] + list(idx) + [len(text)]
    return [text[idx[i]:idx[i+1]] for i in range(len(idx)-1)]


def dir_size_mb(p: Path):
    return p.stat().st_size / 1e6


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="onnx_models")
    ap.add_argument("--models", nargs="+", default=["sat-1l-sm", "sat-3l-sm"])
    ap.add_argument("--vocabs", nargs="+", default=["full", "en_zh"])
    ap.add_argument("--seq-len", type=int, default=256)
    args = ap.parse_args()

    import onnxruntime as ort
    from onnxruntime.quantization import quantize_dynamic, QuantType

    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
    tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")

    # reference: original full-vocab torch model per architecture
    ref_models = {}
    sizes, accuracy = [], []

    for short in args.models:
        repo = f"segment-any-text/{short}"
        ref_models[short] = AutoModelForTokenClassification.from_pretrained(repo).eval()

        for vocab in args.vocabs:
            print(f"\n=== {short} / {vocab} ===", flush=True)
            model = AutoModelForTokenClassification.from_pretrained(repo).eval()
            remap = unk_new = None
            if vocab == "en_zh":
                keep = compute_keep_ids(tokenizer)
                remap = prune_embedding(model, keep)
                unk_new = int(remap[tokenizer.unk_token_id])
                print(f"  pruned vocab {len(keep)}/{len(remap)}", flush=True)

            vdir = out / f"{short}-{vocab}"; vdir.mkdir(exist_ok=True)
            fp32 = vdir / "model.onnx"
            int8 = vdir / "model.int8.onnx"
            export_onnx(LogitsWrapper(model).eval(), args.seq_len, fp32)
            quantize_dynamic(str(fp32), str(int8), weight_type=QuantType.QInt8,
                             op_types_to_quantize=["MatMul", "Gather"])
            for tag, p in [("fp32", fp32), ("int8", int8)]:
                sz = dir_size_mb(p); sizes.append((f"{short}-{vocab}-{tag}", round(sz, 1)))
                print(f"  [{tag}] {sz:7.1f} MB", flush=True)

            # accuracy: compare onnx variants vs original full torch on the test texts
            ref = ref_models[short]
            for tag, p in [("fp32", fp32), ("int8", int8)]:
                sess = ort.InferenceSession(str(p), providers=["CPUExecutionProvider"])
                all_same, max_d = True, 0.0
                for text in TEST_TEXTS:
                    pr = char_boundary_probs(ref, tokenizer, text, is_onnx=False)
                    po = char_boundary_probs(sess, tokenizer, text, remap, unk_new, is_onnx=True)
                    max_d = max(max_d, float(np.abs(pr - po).max()))
                    sr = segment(pr, text); so = segment(po, text)
                    all_same &= ([len(x) for x in sr] == [len(x) for x in so])
                accuracy.append((f"{short}-{vocab}-{tag}", all_same, round(max_d, 4)))
                print(f"  [{tag}] boundaries==orig: {all_same}  maxΔp={max_d:.4f}", flush=True)

    print("\n================ SIZES ================")
    for n, s in sizes:
        print(f"  {n:28s} {s:8.1f} MB")
    print("\n============== ACCURACY vs original torch ==============")
    for n, ok, d in accuracy:
        print(f"  {n:28s} same_chunks={ok!s:5s}  maxΔp={d}")

    # show a real segmentation example from the smallest pruned int8 model
    print("\n========= SAMPLE OUTPUT: sat-1l-sm-en_zh int8, max_length=80 =========")
    small = AutoModelForTokenClassification.from_pretrained(
        f"segment-any-text/{args.models[0]}").eval()
    keep = compute_keep_ids(tokenizer); rm = prune_embedding(small, keep)
    un = int(rm[tokenizer.unk_token_id])
    sess = ort.InferenceSession(str(out / f"{args.models[0]}-en_zh/model.int8.onnx"),
                                providers=["CPUExecutionProvider"])
    for text in TEST_TEXTS:
        probs = char_boundary_probs(sess, tokenizer, text, rm, un, is_onnx=True)
        chunks = segment(probs, text)
        print(f"\n  text ({len(text)} chars):")
        for c in chunks:
            print(f"    [{len(c):3d}] {c.strip()[:75]}")


if __name__ == "__main__":
    main()