File size: 4,972 Bytes
e4805d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """Export Zap: bf16 safetensors (the published weights) + ONNX (fp32, optional int8) with parity checks.
Usage: python export.py MODEL_DIR DATA_DIR OUT_DIR
"""
import json
import os
import shutil
import sys
import numpy as np
import torch
sys.path.insert(0, "/job/work")
from train.zapdata import N_FORM, windows # noqa: E402
from train.zaptrain import pad, read_jsonl # noqa: E402
def main(model_dir, data_dir, out):
from transformers import AutoTokenizer, BertForTokenClassification
os.makedirs(os.path.join(out, "onnx"), exist_ok=True)
tok = AutoTokenizer.from_pretrained(model_dir)
m32 = BertForTokenClassification.from_pretrained(model_dir, dtype=torch.float32).eval()
m_exp = BertForTokenClassification.from_pretrained(model_dir, dtype=torch.float32).eval() # tracing mutates it
# 1) bf16 weights
m16 = BertForTokenClassification.from_pretrained(model_dir, dtype=torch.bfloat16).eval()
m16.save_pretrained(out)
tok.save_pretrained(out)
# portable tokenizer class: loads tokenizer.json as-is in transformers 4.x/5.x and transformers.js
tc_path = os.path.join(out, "tokenizer_config.json")
tc = json.load(open(tc_path))
tc["tokenizer_class"] = "PreTrainedTokenizerFast"
json.dump(tc, open(tc_path, "w"), indent=1)
tok2 = AutoTokenizer.from_pretrained(out)
probe = ["[TITLE] Anmeldung – Café Müller [FLD] [TYPE] input password [NAME] passwort [LBL] Kennwort",
"[TITLE] 登录 - 账户 [FLD] [TYPE] input text [PH] 手机号/邮箱 [FLD] [TYPE] input password",
"[TITLE] Вход [FLD] [TYPE] input email [LBL] Электронная почта [REQ]"]
for t in probe:
assert tok(t)["input_ids"] == tok2(t)["input_ids"], "tokenizer round-trip mismatch"
print("tokenizer round-trip ok", tok2.__class__.__name__)
# 2) ONNX fp32
ids = torch.tensor([[tok.cls_token_id] + [5] * 30 + [tok.sep_token_id]] * 2)
att = torch.ones_like(ids)
class Wrap(torch.nn.Module):
def __init__(self, m):
super().__init__()
self.m = m
def forward(self, input_ids, attention_mask):
return self.m(input_ids=input_ids, attention_mask=attention_mask).logits
onnx_path = os.path.join(out, "onnx", "model.onnx")
torch.onnx.export(Wrap(m_exp), (ids, att), onnx_path, input_names=["input_ids", "attention_mask"],
output_names=["logits"], opset_version=17, dynamo=False,
dynamic_axes={"input_ids": {0: "batch", 1: "seq"}, "attention_mask": {0: "batch", 1: "seq"},
"logits": {0: "batch", 1: "seq"}})
import onnxruntime as ort
from onnxruntime.quantization import QuantType, quantize_dynamic
q_path = os.path.join(out, "onnx", "model_quantized.onnx")
quantize_dynamic(onnx_path, q_path, weight_type=QuantType.QInt8)
sess = {"fp32": ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]),
"int8": ort.InferenceSession(q_path, providers=["CPUExecutionProvider"])}
# 3) parity on test windows: argmax agreement at [CLS] / [FLD] positions vs torch fp32
rows = read_jsonl(os.path.join(data_dir, "test.jsonl"))[:3000]
items = [w for r in rows for w in windows(tok, r["ctx"], r["fields"])]
agree = {k: [0, 0] for k in ("bf16_torch", "fp32_onnx", "int8_onnx")}
maxdiff = 0.0
for s in range(0, len(items), 64):
chunk = items[s:s + 64]
x = pad([c[0] for c in chunk], tok.pad_token_id)
a = (x != tok.pad_token_id).long()
with torch.no_grad():
ref = m32(input_ids=x, attention_mask=a).logits.float()
b16 = m16(input_ids=x, attention_mask=a).logits.float()
outs = {"bf16_torch": b16}
for k, ss in sess.items():
outs[f"{k}_onnx"] = torch.from_numpy(ss.run(None, {"input_ids": x.numpy().astype(np.int64),
"attention_mask": a.numpy().astype(np.int64)})[0])
maxdiff = max(maxdiff, float((outs["fp32_onnx"] - ref).abs().max()))
for b, (_, pos, _) in enumerate(chunk):
for p, sl in [(0, slice(0, N_FORM))] + [(p, slice(N_FORM, None)) for p in pos]:
r = int(ref[b, p, sl].argmax())
for k, o in outs.items():
agree[k][0] += int(o[b, p, sl].argmax()) == r
agree[k][1] += 1
report = {"positions": agree["fp32_onnx"][1], "fp32_onnx_max_abs_logit_diff": maxdiff,
**{f"{k}_argmax_agreement": v[0] / max(1, v[1]) for k, v in agree.items()}}
print(json.dumps(report, indent=1))
json.dump(report, open(os.path.join(out, "export_parity.json"), "w"), indent=1)
if report["int8_onnx_argmax_agreement"] < 0.995:
os.remove(q_path)
print("int8 ONNX dropped: agreement below 99.5%")
shutil.copy(os.path.join(model_dir, "val_metrics.json"), out)
if __name__ == "__main__":
main(*sys.argv[1:4])
|