#!/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()