code/export.py

"""
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())