kinit

/

canary-1b-v2-sk

Like 1

Follow

Kempelen Institute of Intelligent Technologies 48

output from training is missing ctc head, Ivbeen there too. here is script to add it from orignal, thank you for work.

#!/usr/bin/env python
r"""Graft canary-1b-v2's timestamps/aligner submodel into a finetuned .nemo.

The stock canary-1b-v2.nemo bundles a SEPARATE CTC forced-alignment model
(`timestamps_asr_model_*`) that the AED model needs for BOTH word timestamps AND long-audio
chunking (aed_multitask_models.py gates both on `self.timestamps_asr_model is not None`).
A finetune only re-exports the trained EncDecMultiTaskModel, so that aligner is dropped ->
the .nemo is ~2.5 GB smaller and `transcribe(..., timestamps=True)` on long audio OOMs
(full-file attention) and returns no word timestamps.

This tar-copies the aligner artifacts from the ORIGINAL .nemo into the finetuned one, leaving
the trained main model byte-identical. Run it ONCE per training round, on the best .nemo,
before deploying — needed whenever you intend to use timestamps / long-audio transcription.
(Pure short-clip WER benchmarking works without it, but the grafted file is the "proper" one.)

Usage:
    python graft_timestamps.py b:\ViBE_voice_ASR\_Training\Canary-1b\_experiments\Canary-1b-v2-SK-Finetune-v3\checkpoints\Canary-1b-v2-SK-Finetune-v3.nemo -o OUTPUT.nemo --base b:\ViBE_voice_ASR\_ASR_models_NV\canary-1b-v2\canary-1b-v2.nemo

Default output: alongside the input with a `-ts` suffix. Default base: the workspace's stock
canary-1b-v2.nemo in _ASR_models_NV/.
"""

from __future__ import annotations

import argparse
import os
import sys
import tarfile
import time
from pathlib import Path

WORKSPACE = Path(__file__).resolve().parent.parent.parent  # B:/ViBE_voice_ASR
DEFAULT_BASE = WORKSPACE / "_ASR_models_NV" / "canary-1b-v2" / "canary-1b-v2.nemo"


def graft(ft: Path, base: Path, out: Path) -> None:
    # Members present in `base` but missing from `ft` = the aligner artifacts
    # (timestamps_asr_model weights + config, and its own tokenizer/vocab files).
    with tarfile.open(base, "r:*") as b:
        base_names = set(b.getnames())
    with tarfile.open(ft, "r:*") as f:
        ft_names = set(f.getnames())

    missing = sorted(base_names - ft_names)
    if not missing:
        sys.exit("Nothing to graft: finetune already contains all base members.")
    if not any("timestamps_asr_model" in n for n in missing):
        sys.exit("Unexpected: no timestamps_asr_model member in diff:\n  " + "\n  ".join(missing))

    print("Grafting from base:")
    for n in missing:
        print("  +", n)

    t0 = time.time()
    with tarfile.open(out, "w") as o:                       # uncompressed, like .nemo
        with tarfile.open(ft, "r:*") as src:                # all finetune members first
            for m in src.getmembers():
                o.addfile(m, src.extractfile(m) if m.isfile() else None)
        with tarfile.open(base, "r:*") as src:              # then the grafted aligner set
            for n in missing:
                m = src.getmember(n)
                o.addfile(m, src.extractfile(m))
    print(f"\nWrote {out}  ({os.path.getsize(out)/1e9:.2f} GB)  in {time.time()-t0:.0f}s")
    print("Point the _Benchmark / transcribe model path at this file, or replace the original.")


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("finetuned", help="Finetuned .nemo to add the aligner to")
    p.add_argument("-o", "--output", help="Output .nemo (default: <input>-ts.nemo)")
    p.add_argument("--base", default=str(DEFAULT_BASE),
                   help=f"Stock canary .nemo to copy the aligner from (default: {DEFAULT_BASE})")
    a = p.parse_args()

    ft = Path(a.finetuned)
    base = Path(a.base)
    out = Path(a.output) if a.output else ft.with_name(ft.stem + "-ts.nemo")
    for label, pth in (("finetuned", ft), ("base", base)):
        if not pth.is_file():
            sys.exit(f"ERROR: {label} .nemo not found: {pth}")
    if out.resolve() == ft.resolve():
        sys.exit("ERROR: output would overwrite the input; choose a different -o.")

    graft(ft, base, out)


if __name__ == "__main__":
    main()