""" Step 5: export the trained model for the app. python export.py --config config/pilot.yaml python export.py --config config/full.yaml --best # mid-training: the best checkpoint so far Averages the last checkpoints (or takes best.pt with --best), exports ONNX in the interface sherpa-onnx loads as a "zipformer CTC" model, quantizes it to int8, then checks on test clips that sherpa-onnx (the engine on the phone) produces the same text as PyTorch. Finally copies model.int8.onnx and tokens.txt into the app (paths.app_assets). """ import copy import os import shutil import sys import onnx import torch from quran_asr import config as config_mod from quran_asr import data, features, model as model_mod, text def load_averaged(run_dir, average_last): steps = sorted(f for f in os.listdir(run_dir) if f.startswith("step_") and f.endswith(".pt")) chosen = steps[-average_last:] if steps and average_last > 1 else [] if not chosen: chosen = ["best.pt"] if os.path.exists(os.path.join(run_dir, "best.pt")) else steps[-1:] print("averaging:", ", ".join(chosen)) states = [torch.load(os.path.join(run_dir, f), map_location="cpu", weights_only=False) for f in chosen] avg = {} for key, value in states[0]["model"].items(): if value.is_floating_point(): avg[key] = sum(s["model"][key] for s in states) / len(states) else: avg[key] = value return avg, states[0]["model_config"] class Exportable(torch.nn.Module): def __init__(self, model): super().__init__() self.model = model def forward(self, x, x_lens): log_probs, lens = self.model(x, x_lens) return log_probs, lens.to(torch.int64) def export_onnx(model, path): model.eval() x = torch.randn(1, 600, features.NUM_BINS) x_lens = torch.tensor([600], dtype=torch.int64) torch.onnx.export( Exportable(model), (x, x_lens), path, opset_version=17, dynamo=False, input_names=["x", "x_lens"], output_names=["log_probs", "log_probs_len"], dynamic_axes={"x": {0: "N", 1: "T"}, "x_lens": {0: "N"}, "log_probs": {0: "N", 1: "T_out"}, "log_probs_len": {0: "N"}}, ) m = onnx.load(path) meta = {"model_type": "zipformer2_ctc", "vocab_size": str(len(text.VOCAB)), "subsampling_factor": "4", "comment": "Quran recitation, Arabic letters without harakat (Warsh + Hafs)"} for k, v in meta.items(): entry = m.metadata_props.add() entry.key, entry.value = k, v onnx.save(m, path) def quantize(src, dst): from onnxruntime.quantization import QuantType, quantize_dynamic quantize_dynamic(src, dst, op_types_to_quantize=["MatMul"], weight_type=QuantType.QInt8) def sherpa_recognizer(model_path, tokens_path): import sherpa_onnx return sherpa_onnx.OfflineRecognizer.from_zipformer_ctc(model=model_path, tokens=tokens_path, num_threads=4) def sherpa_text(recognizer, wave): stream = recognizer.create_stream() stream.accept_waveform(features.SAMPLE_RATE, wave) recognizer.decode_stream(stream) return stream.result.text.replace(text.WORD, " ").strip() def main(): parser = config_mod.from_args(__doc__) parser.add_argument("--best", action="store_true", help="export best.pt (lowest validation error so far)") args = parser.parse_args() cfg = config_mod.load(args.config) torch.set_num_threads(4) # leave the rest of the CPU to a running training run = cfg.paths.run_dir if args.best: ckpt = torch.load(os.path.join(run, "best.pt"), map_location="cpu", weights_only=False) print(f"best.pt: step {ckpt.get('step')}, validation letter error {ckpt.get('valid_cer', 0) * 100:.2f}%") state, model_config = ckpt["model"], ckpt["model_config"] else: state, model_config = load_averaged(run, cfg.export.average_last) model = model_mod.build(config_mod._wrap(model_config), len(text.VOCAB)) model.load_state_dict(state) model_mod.use_plain_attention(model).eval() out = os.path.join(run, "export") os.makedirs(out, exist_ok=True) fp32 = os.path.join(out, "model.onnx") export_onnx(copy.deepcopy(model), fp32) # tracing changes the model it exports: keep ours clean for the parity check final = fp32 if cfg.export.quantize == "int8": final = os.path.join(out, "model.int8.onnx") quantize(fp32, final) tokens = os.path.join(out, "tokens.txt") text.write_tokens(tokens) print(f"exported {final} ({os.path.getsize(final) / 1e6:.1f} MB)") # Parity: the phone's engine must hear exactly what PyTorch hears. from quran_asr import audio test = data.read_manifest(os.path.join(config_mod.manifest_dir(cfg), "test.jsonl")) clips = test[:: max(1, len(test) // cfg.export.parity_clips)][: cfg.export.parity_clips] rec_final = sherpa_recognizer(final, tokens) rec_fp32 = sherpa_recognizer(fp32, tokens) same = 0 errs = {"pytorch": [0, 0], "phone": [0, 0]} for e in clips: wave = audio.load(e["path"]) feats = features.fbank(torch.from_numpy(wave)) with torch.no_grad(): log_probs, lens = model(feats[None], torch.tensor([feats.shape[0]])) ref_text = text.decode(model_mod.greedy_decode(log_probs, lens)[0]) phone_text = sherpa_text(rec_final, wave) same += ref_text == sherpa_text(rec_fp32, wave) # the export itself must be exact for name, hyp in (("pytorch", ref_text), ("phone", phone_text)): ce, ct, *_ = text.error_counts(hyp, e["text"]) errs[name][0] += ce errs[name][1] += ct cer = {k: v[0] / max(1, v[1]) for k, v in errs.items()} print(f"parity on {len(clips)} test clips: sherpa-onnx fp32 = PyTorch on {same}/{len(clips)}; letter error " f"PyTorch fp32 {cer['pytorch'] * 100:.2f}% vs sherpa-onnx {cfg.export.quantize} {cer['phone'] * 100:.2f}%") if same < 0.9 * len(clips): print("WARNING: sherpa-onnx and PyTorch disagree: the export is broken; not copying it into the app") return 1 if cer["phone"] > cer["pytorch"] + 0.02: print("WARNING: the exported model is noticeably worse than PyTorch; not copying it into the app") return 1 os.makedirs(cfg.paths.app_assets, exist_ok=True) shutil.copy(final, os.path.join(cfg.paths.app_assets, "model.int8.onnx")) shutil.copy(tokens, os.path.join(cfg.paths.app_assets, "tokens.txt")) print("copied into", cfg.paths.app_assets, "- rebuild the app to ship it") return 0 if __name__ == "__main__": sys.exit(main())