Files
doctate/canary/sweep_chunklen.py
T
Brummel 6b7a1cea49 feat: Add Canary ASR model service and sweep experiments
This commit introduces the `doctate-canary` service and associated sweep
experiments.

The service provides access to the `nvidia/canary-1b-v2` ASR model via a
native FastAPI API. It includes deployment scripts, Docker
configurations, and detailed documentation on installation, API usage,
and environment variables.

The sweep experiments aim to thoroughly evaluate the Canary model's
performance under various configurations. This includes testing
different inference pipelines (single-shot vs. buffered), decoding
strategies (greedy vs. beam search), and parameter tuning (chunk length,
overlap, batch size, precision). The goal is to reproduce previous
findings and identify optimal settings.

The commit also includes:
- Utility scripts for audio transcoding and manipulation.
- Comprehensive logging and result collection mechanisms for the sweep
  runs.
- Detailed analysis of `dur=0` occurrences and word confidence,
  concluding they are not reliable indicators of hallucination.
- Documentation on the interaction between decoding parameters,
  especially `return_hypotheses`, and the availability of confidence
  scores.
2026-04-30 14:30:12 +02:00

126 lines
4.3 KiB
Python

"""Mini-sweep: does chunk_len == audio_len fix Path 2 for short audios?
If yes → Path 2 is OK as long as audio fits in exactly one chunk.
The bug is in the multi-chunk merging or pad/split logic.
If no → Path 2 has a structural problem with single-chunk inputs
independent of length matching. Confirms "different inference
pipeline" theory.
"""
import copy
import json
import os
import subprocess
import tempfile
import time
import wave
import torch
from omegaconf import OmegaConf
from nemo.collections.asr.models import ASRModel
from nemo.collections.asr.parts.utils.streaming_utils import FrameBatchMultiTaskAED
from nemo.collections.asr.parts.utils.transcribe_utils import get_buffered_pred_feat_multitaskAED
AUDIO_DIR = "/sweep/threshold_audios"
RESULTS_PATH = "/sweep/chunklen_results.jsonl"
def transcode(src):
dst = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
dst.close()
subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", src,
"-ac", "1", "-ar", "16000", "-f", "wav", dst.name], check=True)
return dst.name
def wav_dur(p):
with wave.open(p, "rb") as w:
return w.getnframes() / float(w.getframerate())
def write_manifest(wav, dur):
f = tempfile.NamedTemporaryFile(suffix=".json", mode="w", delete=False)
json.dump({"audio_filepath": wav, "duration": dur, "taskname": "asr",
"source_lang": "de", "target_lang": "de", "pnc": "yes"}, f)
f.write("\n")
f.close()
return f.name
def run_path2(model, cfg, wav, dur, chunk_len):
manifest = write_manifest(wav, dur)
feature_stride = cfg.preprocessor["window_stride"]
model_stride = feature_stride * 8
fa = FrameBatchMultiTaskAED(asr_model=model, frame_len=chunk_len,
total_buffer=chunk_len, batch_size=8)
t0 = time.monotonic()
with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
with torch.no_grad():
hyps = get_buffered_pred_feat_multitaskAED(
fa, cfg.preprocessor, model_stride, model.device,
manifest=manifest, filepaths=None, timestamps=False,
)
elapsed = time.monotonic() - t0
os.unlink(manifest)
return (getattr(hyps[0], "text", None) or str(hyps[0])).strip(), elapsed
def run_path1(model, wav):
t0 = time.monotonic()
with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
with torch.no_grad():
hyps = model.transcribe([wav], source_lang="de", target_lang="de",
pnc="yes", batch_size=4)
elapsed = time.monotonic() - t0
return (getattr(hyps[0], "text", None) or str(hyps[0])).strip(), elapsed
def main():
open(RESULTS_PATH, "w").close()
print("loading model bf16", flush=True)
model = ASRModel.from_pretrained("nvidia/canary-1b-v2").to("cuda").to(torch.bfloat16)
model.eval()
cfg = copy.deepcopy(model._cfg)
OmegaConf.set_struct(cfg.preprocessor, False)
cfg.preprocessor.dither = 0.0
cfg.preprocessor.pad_to = 0
OmegaConf.set_struct(cfg.preprocessor, True)
print("model loaded", flush=True)
audios = ["trim_20s.m4a", "trim_30s.m4a"]
chunk_lens = {
"trim_20s.m4a": [21.0, 25.0, 40.0], # exact, slightly larger, default
"trim_30s.m4a": [31.0, 35.0, 40.0],
}
def record(d):
with open(RESULTS_PATH, "a") as f:
f.write(json.dumps(d, ensure_ascii=False) + "\n")
for name in audios:
src = os.path.join(AUDIO_DIR, name)
wav = transcode(src)
dur = wav_dur(wav)
print(f"\n=== {name} ({dur:.1f}s) ===", flush=True)
# baseline path1
text, t = run_path1(model, wav)
print(f"path1: {len(text)} chars in {t:.1f}s", flush=True)
record({"audio": name, "duration": dur, "path": "path1",
"chunk_len": None, "chars": len(text),
"duration_infer": round(t, 2), "text": text})
for cl in chunk_lens[name]:
text, t = run_path2(model, cfg, wav, dur, cl)
print(f"path2 chunk_len={cl}: {len(text)} chars in {t:.1f}s", flush=True)
record({"audio": name, "duration": dur, "path": "path2",
"chunk_len": cl, "chars": len(text),
"duration_infer": round(t, 2), "text": text})
print("done", flush=True)
if __name__ == "__main__":
main()