#!/usr/bin/env python3 """Export a trained two-head model to int8 ONNX for CPU scoring. Reads a train.py output directory (model.pt, tokenizer, meta.json), exports ONNX (inputs input_ids / attention_mask, outputs relevance_logit / direction_logit), quantises the embedding table to int8 (weight-only; compute stays fp32), and writes the runtime bundle news_transformer.py loads: models/xau_transformer/model.onnx models/xau_transformer/tokenizer.json models/xau_transformer/meta.json Why only the embeddings: on DeBERTa-v3 any dynamic int8 quantisation of the transformer MatMuls (all, per-channel, reduce-range, FFN only) moved probabilities by up to 0.55-0.72 and flipped 13-16% of decisions (activation outliers). The 128k x 768 embedding table is 70% of the weights; quantising it alone halves the file (566 -> 271 MB) with mean |diff| 0.004 and ~1% flips, all near 0.5. Latency is set by the fp32 encoder: ~18 ms per headline on 1 thread, ~9 ms on 4. Then checks, on up to 500 headlines, that the runtime tokenizer produces the same ids as the training tokenizer, that probabilities stay within --tolerance of the PyTorch fp32 ones and flip at most --max-flip-rate of decisions, and times single-headline CPU latency. Usage: .venv-train/bin/python transformer/export_onnx.py # torch_final if present .venv-train/bin/python transformer/export_onnx.py --src models/xau_transformer/torch """ from __future__ import annotations import argparse import json import os import shutil import sys import time from pathlib import Path import numpy as np import torch ML_NEWS_DIR = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ML_NEWS_DIR)) sys.path.insert(0, str(Path(__file__).resolve().parent)) import news_transformer # noqa: E402 from train import TwoHeadModel, load_rows, split # noqa: E402 BUNDLE_DIR = ML_NEWS_DIR / "models" / "xau_transformer" class Exportable(torch.nn.Module): def __init__(self, model: TwoHeadModel): super().__init__() self.model = model def forward(self, input_ids, attention_mask): return self.model(input_ids, attention_mask) def export(model: TwoHeadModel, tokenizer, path: Path) -> None: sample = tokenizer(["Gold rises as the dollar weakens", "Fed hikes"], padding=True, return_tensors="pt") args = (sample["input_ids"], sample["attention_mask"]) wrapper = Exportable(model).eval() # TorchScript exporter: the dynamo exporter's value_info shapes fail # onnxruntime's shape inference during quantisation (768 vs 1). torch.onnx.export( wrapper, args, str(path), input_names=["input_ids", "attention_mask"], output_names=["relevance_logit", "direction_logit"], dynamic_axes={"input_ids": {0: "batch", 1: "sequence"}, "attention_mask": {0: "batch", 1: "sequence"}, "relevance_logit": {0: "batch"}, "direction_logit": {0: "batch"}}, opset_version=17, dynamo=False) def main() -> int: ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) ap.add_argument("--src", type=Path, default=None, help="train.py output dir (default: torch_final, else torch)") ap.add_argument("--out", type=Path, default=BUNDLE_DIR) ap.add_argument("--tolerance", type=float, default=0.05, help="max |probability difference| vs PyTorch fp32") ap.add_argument("--max-flip-rate", type=float, default=0.02, help="max share of relevance/direction decisions that may flip") ap.add_argument("--check-rows", type=int, default=500) args = ap.parse_args() src = args.src or next(p for p in (BUNDLE_DIR / "torch_final", BUNDLE_DIR / "torch") if (p / "model.pt").is_file()) meta = json.loads((src / "meta.json").read_text()) from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(src) model = TwoHeadModel(meta["base_model"]) model.load_state_dict(torch.load(src / "model.pt", map_location="cpu")) model.eval() args.out.mkdir(parents=True, exist_ok=True) fp32 = args.out / "model.fp32.onnx" int8 = args.out / "model.onnx" print(f"exporting {src} -> {fp32}") export(model, tokenizer, fp32) from onnxruntime.quantization import QuantType, quantize_dynamic quantize_dynamic(str(fp32), str(int8), weight_type=QuantType.QInt8, op_types_to_quantize=["Gather"]) print(f"embeddings int8 -> {int8} ({int8.stat().st_size / 1e6:.0f} MB, " f"fp32 {fp32.stat().st_size / 1e6:.0f} MB)") shutil.copy(src / "tokenizer.json", args.out / "tokenizer.json") runtime_meta = { **meta, "onnx_file": int8.name, "quantization": "int8 weight-only embedding table; fp32 compute", "pad_token": tokenizer.pad_token, "pad_token_id": tokenizer.pad_token_id, "exported_from": os.path.relpath(src.resolve(), ML_NEWS_DIR), } (args.out / "meta.json").write_text(json.dumps(runtime_meta, indent=2, default=str) + "\n") # ── parity: tokenizer ids and probabilities on held-out headlines ── _, _, val = split(load_rows(), meta["hyperparameters"]["val_from"], meta["hyperparameters"]["seed"]) texts = [r["text"] for r in val[: args.check_rows]] bundle = news_transformer.load_model(args.out) tok_mismatch = sum( bundle["tokenizer"].encode(t).ids != tokenizer(t, truncation=True, max_length=news_transformer.MAX_LENGTH)["input_ids"] for t in texts) with torch.no_grad(): ref_rel, ref_dir = [], [] for start in range(0, len(texts), 64): enc = tokenizer(texts[start:start + 64], padding=True, truncation=True, max_length=news_transformer.MAX_LENGTH, return_tensors="pt") r, d = model(enc["input_ids"], enc["attention_mask"]) ref_rel.append(torch.sigmoid(r)) ref_dir.append(torch.sigmoid(d)) ref_rel = torch.cat(ref_rel).numpy() ref_dir = torch.cat(ref_dir).numpy() got = [] for start in range(0, len(texts), 64): got += news_transformer.score_batch(bundle, texts[start:start + 64]) diff_rel = np.abs(ref_rel - np.array([g["p_relevant"] for g in got])) diff_dir = np.abs(ref_dir - np.array([g["p_bullish"] for g in got])) thr = runtime_meta["relevance_threshold"] flips = int(sum((a >= thr) != g["relevant"] for a, g in zip(ref_rel, got)) + sum((b >= 0.5) != (g["direction"] == "bullish") for b, g in zip(ref_dir, got))) print(f"tokenizer mismatches: {tok_mismatch}/{len(texts)}") print(f"int8 vs fp32 over {len(texts)}: p_relevant |diff| mean {diff_rel.mean():.4f} max " f"{diff_rel.max():.4f}; p_bullish mean {diff_dir.mean():.4f} max {diff_dir.max():.4f}; " f"decision flips {flips}") # ── latency: one headline at a time, as the live feed would call it ── latency = {} for threads in (1, os.cpu_count()): b = news_transformer.load_model(args.out, threads=threads) for t in texts[:20]: news_transformer.score_news(b, t) t0 = time.perf_counter() for t in texts[:200]: news_transformer.score_news(b, t) latency[threads] = (time.perf_counter() - t0) * 1000 / min(200, len(texts)) print(f"CPU latency, {threads} thread(s): {latency[threads]:.1f} ms per headline") runtime_meta["export_checks"] = { "rows": len(texts), "tokenizer_mismatches": tok_mismatch, "p_relevant_max_abs_diff": float(diff_rel.max()), "p_bullish_max_abs_diff": float(diff_dir.max()), "decision_flips": flips, "latency_ms_per_headline": {str(k): round(v, 2) for k, v in latency.items()}, } (args.out / "meta.json").write_text(json.dumps(runtime_meta, indent=2, default=str) + "\n") fp32.unlink() flip_rate = flips / (2 * len(texts)) ok = (tok_mismatch == 0 and max(diff_rel.max(), diff_dir.max()) <= args.tolerance and flip_rate <= args.max_flip_rate) print("PARITY OK" if ok else f"PARITY FAILED (tolerance {args.tolerance}, flip rate {flip_rate:.1%} " f"> {args.max_flip_rate:.0%}?)") return 0 if ok else 1 if __name__ == "__main__": sys.exit(main())