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.
This commit is contained in:
@@ -0,0 +1,178 @@
|
||||
"""Pull cross-attention scores and word confidence on the dosing audio.
|
||||
|
||||
Goal: see whether the AED's cross-attention distribution and per-word
|
||||
confidence flag the hallucinated "5 milligramm 101" cluster that we
|
||||
already identified via zero-duration word offsets.
|
||||
|
||||
Two new signals beyond what probe_canary_ts.py captured:
|
||||
1. word_confidence — per-word probability proxy. Hallucinations *should*
|
||||
have lower confidence than real words, though not guaranteed because
|
||||
AEDs hallucinate with high confidence too.
|
||||
2. xatt_scores — cross-attention from decoder tokens to encoder frames.
|
||||
For each output token we summarize where it "looked" in the audio:
|
||||
- peak_frame (argmax) and peak_time_s
|
||||
- entropy of the distribution (flat = no clear acoustic anchor)
|
||||
Hallucinated tokens should show high entropy (no clear acoustic peak).
|
||||
|
||||
Output:
|
||||
/sweep/probe_attn_conf.json — text + timestamp + word_confidence + per-token
|
||||
xatt summary (argmax frame, entropy)
|
||||
"""
|
||||
import json
|
||||
import math
|
||||
import subprocess
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from nemo.collections.asr.models import EncDecMultiTaskModel
|
||||
from nemo.collections.asr.models.aed_multitask_models import (
|
||||
MultiTaskTranscriptionConfig,
|
||||
parse_multitask_prompt,
|
||||
)
|
||||
|
||||
|
||||
WAV_SRC = "/sweep/case_c414cf52_15-01-39.m4a"
|
||||
OUT_PATH = "/sweep/probe_attn_conf_kwargs_returnhyp.json"
|
||||
|
||||
|
||||
def transcode(src: str) -> str:
|
||||
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 to_jsonable(o):
|
||||
if isinstance(o, torch.Tensor):
|
||||
return o.detach().cpu().tolist()
|
||||
if hasattr(o, "tolist"):
|
||||
return o.tolist()
|
||||
if hasattr(o, "item"):
|
||||
return o.item()
|
||||
return str(o)
|
||||
|
||||
|
||||
def entropy_of(probs: torch.Tensor) -> float:
|
||||
p = probs.clamp_min(1e-12)
|
||||
return float(-(p * p.log()).sum().item())
|
||||
|
||||
|
||||
def summarize_xatt(xatt_scores, frame_stride_s: float):
|
||||
"""xatt_scores: List[Tensor] over decoder layers.
|
||||
|
||||
Multi-head attention output shape is typically [heads, tokens, frames]
|
||||
or [tokens, frames] depending on the decoder. We mean-pool away any
|
||||
leading dimensions until we get a 2D [tokens, frames] matrix.
|
||||
"""
|
||||
if not xatt_scores:
|
||||
return None
|
||||
layers = [x.float() for x in xatt_scores]
|
||||
print(f" xatt layer shapes: {[tuple(l.shape) for l in layers]}", flush=True)
|
||||
avg = torch.stack(layers, dim=0).mean(dim=0)
|
||||
# Collapse all leading dims (batch, heads, ...) until 2D
|
||||
while avg.dim() > 2:
|
||||
avg = avg.mean(dim=0)
|
||||
n_tokens, n_frames = avg.shape
|
||||
summaries = []
|
||||
for t in range(n_tokens):
|
||||
row = avg[t]
|
||||
norm = row / row.sum().clamp_min(1e-12)
|
||||
peak = int(norm.argmax().item())
|
||||
summaries.append({
|
||||
"token_idx": t,
|
||||
"peak_frame": peak,
|
||||
"peak_time_s": round(peak * frame_stride_s, 3),
|
||||
"peak_prob": round(float(norm[peak].item()), 4),
|
||||
"entropy": round(entropy_of(norm), 4),
|
||||
"max_log_entropy": round(math.log(n_frames), 4),
|
||||
})
|
||||
return {"n_tokens": n_tokens, "n_frames": n_frames,
|
||||
"frame_stride_s": frame_stride_s,
|
||||
"n_layers": len(layers), "per_token": summaries}
|
||||
|
||||
|
||||
def main():
|
||||
print("loading Canary 1B v2", flush=True)
|
||||
model = EncDecMultiTaskModel.from_pretrained("nvidia/canary-1b-v2")
|
||||
model = model.to("cuda").to(torch.bfloat16)
|
||||
model.eval()
|
||||
|
||||
# Patch the live decoding config:
|
||||
# - greedy strategy (confidence is only filled in greedy path, not beam)
|
||||
# - return_xattn_scores=True
|
||||
# - confidence_cfg with all preserve flags
|
||||
# - enable_chunking=False is set via override_config below (NeMo gates xatt
|
||||
# on chunking; chunking also blocks word-level confidence aggregation).
|
||||
cfg = OmegaConf.to_container(model.cfg.decoding, resolve=True)
|
||||
cfg["strategy"] = "greedy"
|
||||
cfg["return_xattn_scores"] = True
|
||||
cfg.setdefault("confidence_cfg", {})
|
||||
cfg["confidence_cfg"]["preserve_word_confidence"] = True
|
||||
cfg["confidence_cfg"]["preserve_token_confidence"] = True
|
||||
cfg["confidence_cfg"]["preserve_frame_confidence"] = True
|
||||
cfg["confidence_cfg"]["aggregation"] = "min" # most signal-preserving
|
||||
cfg.setdefault("greedy", {})
|
||||
cfg["greedy"]["preserve_token_confidence"] = True
|
||||
cfg["greedy"]["preserve_alignments"] = True
|
||||
print("applying decoding override:", json.dumps(cfg, indent=2), flush=True)
|
||||
model.change_decoding_strategy(OmegaConf.create(cfg))
|
||||
|
||||
# Encoder frame stride: feature_stride * subsampling
|
||||
feature_stride = float(model.cfg.preprocessor["window_stride"])
|
||||
subsampling = 8 # FastConformer default; verified empirically against
|
||||
# 80ms-per-frame in probe_tdt.md
|
||||
frame_stride_s = feature_stride * subsampling
|
||||
print(f"frame stride: {frame_stride_s} s", flush=True)
|
||||
|
||||
wav = transcode(WAV_SRC)
|
||||
# Clean single-knob ablation: same override_config for both runs,
|
||||
# only `enable_chunking` differs. All other fields (batch_size, prompt,
|
||||
# timestamps, return_hypotheses) held constant.
|
||||
print(f"running transcribe(KWARGS path, return_hypotheses=True, timestamps=True) on {WAV_SRC}",
|
||||
flush=True)
|
||||
# Final clean test: kwargs path (like the original Mode A run that
|
||||
# showed None confidence), but with return_hypotheses=True forced.
|
||||
# If confidence appears, then `return_hypotheses` was the real
|
||||
# confounder all along — not timestamps, not the call form.
|
||||
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", timestamps=True, batch_size=1,
|
||||
return_hypotheses=True)
|
||||
h = hyps[0]
|
||||
|
||||
out = {
|
||||
"text": getattr(h, "text", None),
|
||||
"score": getattr(h, "score", None),
|
||||
"timestamp": getattr(h, "timestamp", None),
|
||||
"tokens": getattr(h, "tokens", None),
|
||||
"word_confidence": getattr(h, "word_confidence", None),
|
||||
"token_confidence": getattr(h, "token_confidence", None),
|
||||
"xatt_summary": summarize_xatt(getattr(h, "xatt_scores", None),
|
||||
frame_stride_s),
|
||||
}
|
||||
|
||||
with open(OUT_PATH, "w") as f:
|
||||
json.dump(out, f, default=to_jsonable, ensure_ascii=False, indent=2)
|
||||
|
||||
# Quick sanity print
|
||||
print(f"\ntext: {h.text}", flush=True)
|
||||
print(f"score: {h.score}", flush=True)
|
||||
wc = getattr(h, "word_confidence", None)
|
||||
print(f"word_confidence: {len(wc) if wc is not None else 'None'} entries", flush=True)
|
||||
tc = getattr(h, "token_confidence", None)
|
||||
print(f"token_confidence: {len(tc) if tc is not None else 'None'} entries", flush=True)
|
||||
xs = getattr(h, "xatt_scores", None)
|
||||
if xs is not None:
|
||||
print(f"xatt_scores: {len(xs)} layer(s); shapes={[tuple(x.shape) for x in xs]}", flush=True)
|
||||
else:
|
||||
print("xatt_scores: None", flush=True)
|
||||
print(f"\nfull dump -> {OUT_PATH}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,69 @@
|
||||
"""Show forced-alignment timestamps from Canary's main transcribe(timestamps=True).
|
||||
|
||||
Demonstrates the full Hypothesis structure produced by:
|
||||
AED text + CTC submodel acoustic alignment via Viterbi.
|
||||
|
||||
Output:
|
||||
/sweep/canary_timestamps_demo.json — full pretty-printed dict
|
||||
stdout — short summary
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
|
||||
from nemo.collections.asr.models import EncDecMultiTaskModel
|
||||
|
||||
|
||||
WAV_SRC = "/sweep/case_c414cf52_15-01-39.m4a"
|
||||
OUT_PATH = "/sweep/canary_timestamps_dosing.json"
|
||||
|
||||
|
||||
def transcode(src: str) -> str:
|
||||
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 to_jsonable(o):
|
||||
if hasattr(o, "tolist"):
|
||||
return o.tolist()
|
||||
if hasattr(o, "item"):
|
||||
return o.item()
|
||||
return str(o)
|
||||
|
||||
|
||||
def main():
|
||||
print("loading Canary 1B v2", flush=True)
|
||||
model = EncDecMultiTaskModel.from_pretrained("nvidia/canary-1b-v2")
|
||||
model = model.to("cuda").to(torch.bfloat16)
|
||||
model.eval()
|
||||
|
||||
wav = transcode(WAV_SRC)
|
||||
print(f"running model.transcribe(timestamps=True) on {WAV_SRC}", flush=True)
|
||||
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", timestamps=True, batch_size=4)
|
||||
h = hyps[0]
|
||||
|
||||
out = {"text": getattr(h, "text", None),
|
||||
"timestamp": getattr(h, "timestamp", None)}
|
||||
with open(OUT_PATH, "w") as f:
|
||||
json.dump(out, f, default=to_jsonable, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"\ntext: {h.text}", flush=True)
|
||||
ts = h.timestamp if isinstance(h.timestamp, dict) else {}
|
||||
for k, v in ts.items():
|
||||
n = len(v) if hasattr(v, "__len__") else "?"
|
||||
print(f"timestamp[{k}]: {n} entries", flush=True)
|
||||
print(f"\nfull dump -> {OUT_PATH}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Standalone CTC-submodel transcription on the dosing audio.
|
||||
|
||||
Loads Canary, accesses model.timestamps_asr_model directly, calls its own
|
||||
transcribe() — bypassing the AED decoder entirely. Goal: see whether the CTC
|
||||
submodel renders the spoken digit sequence "eins null eins" as separate
|
||||
tokens or also collapses it.
|
||||
"""
|
||||
import json
|
||||
import subprocess
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
|
||||
from nemo.collections.asr.models import EncDecMultiTaskModel
|
||||
|
||||
|
||||
WAV_SRC = "/sweep/case_c414cf52_15-01-39.m4a"
|
||||
OUT_PATH = "/sweep/ctc_standalone_dosing.json"
|
||||
|
||||
|
||||
def transcode(src: str) -> str:
|
||||
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 to_jsonable(o):
|
||||
if hasattr(o, "tolist"):
|
||||
return o.tolist()
|
||||
if hasattr(o, "item"):
|
||||
return o.item()
|
||||
return str(o)
|
||||
|
||||
|
||||
def main():
|
||||
print("loading Canary 1B v2 (just for the CTC submodel)", flush=True)
|
||||
model = EncDecMultiTaskModel.from_pretrained("nvidia/canary-1b-v2")
|
||||
model = model.to("cuda").to(torch.bfloat16)
|
||||
model.eval()
|
||||
|
||||
tdt = model.timestamps_asr_model
|
||||
print(f"submodel: {type(tdt).__name__}", flush=True)
|
||||
|
||||
wav = transcode(WAV_SRC)
|
||||
print(f"running tdt.transcribe(timestamps=True) on {WAV_SRC}", flush=True)
|
||||
with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
|
||||
with torch.no_grad():
|
||||
hyps = tdt.transcribe([wav], timestamps=True)
|
||||
h = hyps[0]
|
||||
|
||||
out = {"text": getattr(h, "text", None),
|
||||
"timestamp": getattr(h, "timestamp", None)}
|
||||
with open(OUT_PATH, "w") as f:
|
||||
json.dump(out, f, default=to_jsonable, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"\nstandalone CTC text:\n{h.text}\n", flush=True)
|
||||
ts = h.timestamp if isinstance(h.timestamp, dict) else {}
|
||||
for k, v in ts.items():
|
||||
n = len(v) if hasattr(v, "__len__") else "?"
|
||||
print(f"timestamp[{k}]: {n} entries", flush=True)
|
||||
print(f"\nfull dump -> {OUT_PATH}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,215 @@
|
||||
"""TDT probe — can timestamps_asr_model handle audio standalone, including >40 s?
|
||||
|
||||
Background:
|
||||
EncDecMultiTaskModel.timestamps_asr_model is a Token-Duration-Transducer
|
||||
(RNN-T variant) with hybrid CTC head, FastConformer encoder enc_hidden=1024.
|
||||
Bundled inside Canary 1B v2 specifically to provide forced-alignment
|
||||
timestamps for the AED head, but it is itself a complete ASR model.
|
||||
|
||||
Open question: does the relative positional encoding lift the 40 s wall that
|
||||
AED hits, or does TDT inherit the same training-time max_duration limit?
|
||||
|
||||
Test set:
|
||||
- 20 s : sanity check, well within both models
|
||||
- 80 s : full source — Canary path1 truncates here, baseline expectation
|
||||
- 160 s : concat(80, 80), no re-encode — pushes 4x past TDT training horizon
|
||||
|
||||
Per case:
|
||||
- TDT.transcribe(timestamps=False)
|
||||
- TDT.transcribe(timestamps=True) — also verify timestamp shape
|
||||
- Canary path1 — reference, will truncate >40 s
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
import wave
|
||||
|
||||
import torch
|
||||
|
||||
from nemo.collections.asr.models import EncDecMultiTaskModel
|
||||
|
||||
|
||||
AUDIO_DIR = "/sweep/threshold_audios"
|
||||
RESULTS_PATH = "/sweep/probe_tdt.jsonl"
|
||||
LOG_PATH = "/sweep/probe_tdt.log"
|
||||
|
||||
|
||||
def log(msg: str) -> None:
|
||||
line = f"{time.strftime('%H:%M:%S')} {msg}"
|
||||
print(line, flush=True)
|
||||
with open(LOG_PATH, "a") as f:
|
||||
f.write(line + "\n")
|
||||
|
||||
|
||||
def transcode(src: str, duration_secs: float | None = None) -> str:
|
||||
dst = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
|
||||
dst.close()
|
||||
cmd = ["ffmpeg", "-y", "-loglevel", "error", "-i", src,
|
||||
"-ac", "1", "-ar", "16000"]
|
||||
if duration_secs is not None:
|
||||
cmd += ["-t", str(duration_secs)]
|
||||
cmd += ["-f", "wav", dst.name]
|
||||
subprocess.run(cmd, check=True)
|
||||
return dst.name
|
||||
|
||||
|
||||
def concat_wavs(wavs: list[str], repeats: int = 1) -> str:
|
||||
list_file = tempfile.NamedTemporaryFile(suffix=".txt", mode="w", delete=False)
|
||||
for w in wavs * repeats:
|
||||
list_file.write(f"file '{w}'\n")
|
||||
list_file.close()
|
||||
dst = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
|
||||
dst.close()
|
||||
subprocess.run(["ffmpeg", "-y", "-loglevel", "error",
|
||||
"-f", "concat", "-safe", "0", "-i", list_file.name,
|
||||
"-c", "copy", dst.name], check=True)
|
||||
os.unlink(list_file.name)
|
||||
return dst.name
|
||||
|
||||
|
||||
def wav_duration(p: str) -> float:
|
||||
with wave.open(p, "rb") as w:
|
||||
return w.getnframes() / float(w.getframerate())
|
||||
|
||||
|
||||
def vram_mb() -> float:
|
||||
return torch.cuda.max_memory_allocated() / (1024 * 1024) if torch.cuda.is_available() else 0.0
|
||||
|
||||
|
||||
def run_tdt(tdt, wav: str, timestamps: bool) -> dict:
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
|
||||
with torch.no_grad():
|
||||
hyps = tdt.transcribe([wav], timestamps=timestamps)
|
||||
elapsed = time.monotonic() - t0
|
||||
h = hyps[0]
|
||||
text = (getattr(h, "text", None) or "").strip()
|
||||
ts = getattr(h, "timestamp", None)
|
||||
ts_summary = None
|
||||
if isinstance(ts, dict):
|
||||
ts_summary = {k: (len(v) if hasattr(v, "__len__") else str(type(v).__name__))
|
||||
for k, v in ts.items()}
|
||||
return {"ok": True, "text": text, "chars": len(text),
|
||||
"duration_infer": round(elapsed, 2),
|
||||
"vram_peak_mb": round(vram_mb(), 0),
|
||||
"timestamp_summary": ts_summary}
|
||||
except Exception as e:
|
||||
return {"ok": False, "error": f"{type(e).__name__}: {e}",
|
||||
"duration_infer": round(time.monotonic() - t0, 2)}
|
||||
|
||||
|
||||
def run_canary(model, wav: str) -> dict:
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
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
|
||||
text = (getattr(hyps[0], "text", None) or str(hyps[0])).strip()
|
||||
return {"ok": True, "text": text, "chars": len(text),
|
||||
"duration_infer": round(elapsed, 2),
|
||||
"vram_peak_mb": round(vram_mb(), 0)}
|
||||
except Exception as e:
|
||||
return {"ok": False, "error": f"{type(e).__name__}: {e}",
|
||||
"duration_infer": round(time.monotonic() - t0, 2)}
|
||||
|
||||
|
||||
def find_source() -> str | None:
|
||||
for c in ("trim_full.m4a", "full.m4a", "source.m4a"):
|
||||
p = os.path.join(AUDIO_DIR, c)
|
||||
if os.path.exists(p):
|
||||
return p
|
||||
if not os.path.isdir(AUDIO_DIR):
|
||||
return None
|
||||
m4as = [f for f in os.listdir(AUDIO_DIR) if f.endswith(".m4a")]
|
||||
if not m4as:
|
||||
return None
|
||||
m4as.sort(key=lambda f: os.path.getsize(os.path.join(AUDIO_DIR, f)), reverse=True)
|
||||
return os.path.join(AUDIO_DIR, m4as[0])
|
||||
|
||||
|
||||
def main():
|
||||
open(RESULTS_PATH, "w").close()
|
||||
open(LOG_PATH, "w").close()
|
||||
|
||||
log("loading EncDecMultiTaskModel (Canary 1B v2)")
|
||||
model = EncDecMultiTaskModel.from_pretrained("nvidia/canary-1b-v2")
|
||||
model = model.to("cuda").to(torch.bfloat16)
|
||||
model.eval()
|
||||
log(f" vram after load: {vram_mb():.0f} MB")
|
||||
|
||||
tdt = getattr(model, "timestamps_asr_model", None)
|
||||
if tdt is None:
|
||||
log("FATAL: model has no timestamps_asr_model attribute")
|
||||
return
|
||||
log(f" tdt class: {type(tdt).__name__}")
|
||||
log(f" tdt device: {next(tdt.parameters()).device}, dtype: {next(tdt.parameters()).dtype}")
|
||||
|
||||
src = find_source()
|
||||
if src is None:
|
||||
log(f"FATAL: no source audio found in {AUDIO_DIR}")
|
||||
return
|
||||
log(f"source audio: {src}")
|
||||
|
||||
wav_full = transcode(src)
|
||||
dur_full = wav_duration(wav_full)
|
||||
log(f" full wav: {dur_full:.1f} s")
|
||||
|
||||
wav_20 = transcode(src, duration_secs=20.0)
|
||||
log(f" 20 s wav: {wav_duration(wav_20):.1f} s")
|
||||
|
||||
wav_80 = transcode(src, duration_secs=80.0) if dur_full >= 80 else wav_full
|
||||
log(f" 80 s wav: {wav_duration(wav_80):.1f} s")
|
||||
|
||||
wav_160 = concat_wavs([wav_80], repeats=2)
|
||||
log(f" 160 s concat wav: {wav_duration(wav_160):.1f} s")
|
||||
|
||||
cases = [("20s", wav_20), ("80s", wav_80), ("160s_concat", wav_160)]
|
||||
|
||||
def record(d: dict) -> None:
|
||||
with open(RESULTS_PATH, "a") as f:
|
||||
f.write(json.dumps(d, ensure_ascii=False) + "\n")
|
||||
|
||||
for label, wav in cases:
|
||||
dur = wav_duration(wav)
|
||||
log(f"\n=== case {label} ({dur:.1f} s) ===")
|
||||
|
||||
log(" TDT timestamps=False")
|
||||
r = run_tdt(tdt, wav, timestamps=False)
|
||||
record({"case": label, "duration_audio": dur, "model": "tdt",
|
||||
"timestamps": False, **r})
|
||||
if r["ok"]:
|
||||
log(f" -> {r['chars']} chars in {r['duration_infer']} s, "
|
||||
f"{r['vram_peak_mb']} MB")
|
||||
else:
|
||||
log(f" FAIL: {r['error']}")
|
||||
|
||||
log(" TDT timestamps=True")
|
||||
r = run_tdt(tdt, wav, timestamps=True)
|
||||
record({"case": label, "duration_audio": dur, "model": "tdt",
|
||||
"timestamps": True, **r})
|
||||
if r["ok"]:
|
||||
log(f" -> {r['chars']} chars, ts_summary={r['timestamp_summary']}")
|
||||
else:
|
||||
log(f" FAIL: {r['error']}")
|
||||
|
||||
log(" Canary path1 (reference)")
|
||||
r = run_canary(model, wav)
|
||||
record({"case": label, "duration_audio": dur, "model": "canary_path1",
|
||||
**r})
|
||||
if r["ok"]:
|
||||
log(f" -> {r['chars']} chars in {r['duration_infer']} s")
|
||||
else:
|
||||
log(f" FAIL: {r['error']}")
|
||||
|
||||
log("done")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+425
@@ -0,0 +1,425 @@
|
||||
"""Canary parameter sweep — direct NeMo calls, no HTTP layer.
|
||||
|
||||
Runs many parameter combinations against a fixed set of audios and writes
|
||||
JSONL + a human-readable summary. Designed to run inside the doctate-canary
|
||||
Docker image (or its base) with the model volume mounted at /models.
|
||||
|
||||
Outputs:
|
||||
/sweep/results.jsonl — one JSON object per run
|
||||
/sweep/log.txt — chronological log
|
||||
"""
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import wave
|
||||
from contextlib import contextmanager
|
||||
|
||||
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/audios"
|
||||
OUT_DIR = "/sweep"
|
||||
RESULTS_PATH = f"{OUT_DIR}/results.jsonl"
|
||||
LOG_PATH = f"{OUT_DIR}/log.txt"
|
||||
|
||||
# Subset for most axes — one short, one medium, one long (>40 s each).
|
||||
SUBSET = [
|
||||
"2026-04-27T15-00-48Z.m4a", # 48.4 s
|
||||
"2026-04-27T14-55-41Z.m4a", # 65.9 s
|
||||
"2026-04-27T14-58-17Z.m4a", # 80.9 s
|
||||
]
|
||||
|
||||
# Single-audio subset for speed/VRAM-only axes.
|
||||
SINGLE = ["2026-04-27T14-58-17Z.m4a"] # 80.9 s
|
||||
|
||||
DTYPE_MAP = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
|
||||
|
||||
|
||||
def log(msg: str) -> None:
|
||||
line = f"{time.strftime('%H:%M:%S')} {msg}"
|
||||
print(line, flush=True)
|
||||
with open(LOG_PATH, "a") as f:
|
||||
f.write(line + "\n")
|
||||
|
||||
|
||||
def transcode_to_wav16k(src: str) -> str:
|
||||
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_duration(path: str) -> float:
|
||||
with wave.open(path, "rb") as w:
|
||||
return w.getnframes() / float(w.getframerate())
|
||||
|
||||
|
||||
def vram_mb() -> float:
|
||||
if torch.cuda.is_available():
|
||||
return torch.cuda.max_memory_allocated() / (1024 * 1024)
|
||||
return 0.0
|
||||
|
||||
|
||||
@contextmanager
|
||||
def vram_track():
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
pass
|
||||
|
||||
|
||||
def load_model(precision: str):
|
||||
log(f"loading model precision={precision}")
|
||||
t0 = time.monotonic()
|
||||
model = ASRModel.from_pretrained("nvidia/canary-1b-v2")
|
||||
model = model.to("cuda")
|
||||
if precision != "fp32":
|
||||
model = model.to(DTYPE_MAP[precision])
|
||||
model.eval()
|
||||
log(f"model loaded in {time.monotonic() - t0:.1f}s, vram={vram_mb():.0f}MB")
|
||||
return model
|
||||
|
||||
|
||||
def prep_cfg(model):
|
||||
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)
|
||||
return cfg
|
||||
|
||||
|
||||
def write_manifest(wav_path: str, duration: float, pnc: str = "yes",
|
||||
src_lang: str = "de", tgt_lang: str = "de",
|
||||
task: str = "asr") -> str:
|
||||
f = tempfile.NamedTemporaryFile(suffix=".json", mode="w", delete=False)
|
||||
json.dump({
|
||||
"audio_filepath": wav_path,
|
||||
"duration": duration,
|
||||
"taskname": task,
|
||||
"source_lang": src_lang,
|
||||
"target_lang": tgt_lang,
|
||||
"pnc": pnc,
|
||||
}, f)
|
||||
f.write("\n")
|
||||
f.close()
|
||||
return f.name
|
||||
|
||||
|
||||
def run_buffered(model, model_cfg, wav_path: str, duration: float,
|
||||
chunk_len: float, total_buffer: float, batch_size: int,
|
||||
precision: str, pnc: str = "yes", timestamps: bool = False):
|
||||
manifest = write_manifest(wav_path, duration, pnc=pnc)
|
||||
feature_stride = model_cfg.preprocessor["window_stride"]
|
||||
model_stride_secs = feature_stride * 8
|
||||
|
||||
frame_asr = FrameBatchMultiTaskAED(
|
||||
asr_model=model,
|
||||
frame_len=chunk_len,
|
||||
total_buffer=total_buffer,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
amp_dtype = DTYPE_MAP[precision]
|
||||
use_amp = precision != "fp32"
|
||||
t0 = time.monotonic()
|
||||
with vram_track():
|
||||
with torch.amp.autocast("cuda", enabled=use_amp, dtype=amp_dtype):
|
||||
with torch.no_grad():
|
||||
hyps = get_buffered_pred_feat_multitaskAED(
|
||||
frame_asr, model_cfg.preprocessor, model_stride_secs,
|
||||
model.device, manifest=manifest, filepaths=None,
|
||||
timestamps=timestamps,
|
||||
)
|
||||
elapsed = time.monotonic() - t0
|
||||
peak_vram = vram_mb()
|
||||
os.unlink(manifest)
|
||||
text = (getattr(hyps[0], "text", None) or str(hyps[0])).strip()
|
||||
return text, elapsed, peak_vram
|
||||
|
||||
|
||||
def run_transcribe_path1(model, wav_path: str, batch_size: int, pnc: str,
|
||||
timestamps: bool, precision: str):
|
||||
"""Path 1 — direct .transcribe() with batch_size=1 enables auto-chunking
|
||||
in NeMo >= 2.5 (per docstring of speech_to_text_aed_chunked_infer.py)."""
|
||||
amp_dtype = DTYPE_MAP[precision]
|
||||
use_amp = precision != "fp32"
|
||||
t0 = time.monotonic()
|
||||
with vram_track():
|
||||
with torch.amp.autocast("cuda", enabled=use_amp, dtype=amp_dtype):
|
||||
with torch.no_grad():
|
||||
hyps = model.transcribe(
|
||||
[wav_path],
|
||||
source_lang="de", target_lang="de",
|
||||
pnc=pnc, timestamps=timestamps,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
elapsed = time.monotonic() - t0
|
||||
peak_vram = vram_mb()
|
||||
text = (getattr(hyps[0], "text", None) or str(hyps[0])).strip()
|
||||
return text, elapsed, peak_vram
|
||||
|
||||
|
||||
def write_result(record: dict) -> None:
|
||||
with open(RESULTS_PATH, "a") as f:
|
||||
f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
|
||||
|
||||
def sweep():
|
||||
open(LOG_PATH, "w").close()
|
||||
open(RESULTS_PATH, "w").close()
|
||||
|
||||
log("=== preparing audios ===")
|
||||
wavs = {}
|
||||
for name in set(SUBSET + SINGLE):
|
||||
src = os.path.join(AUDIO_DIR, name)
|
||||
wav = transcode_to_wav16k(src)
|
||||
dur = wav_duration(wav)
|
||||
wavs[name] = (wav, dur)
|
||||
log(f" {name} → {wav} ({dur:.1f}s)")
|
||||
|
||||
# ===== precision axis (B.2) — separate model loads =====
|
||||
for precision in ["bf16", "fp16", "fp32"]:
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
model = load_model(precision)
|
||||
cfg = prep_cfg(model)
|
||||
for name in SINGLE:
|
||||
wav, dur = wavs[name]
|
||||
axis = f"B2_precision={precision}"
|
||||
log(f"running {axis} on {name}")
|
||||
try:
|
||||
text, t, vram = run_buffered(
|
||||
model, cfg, wav, dur,
|
||||
chunk_len=40.0, total_buffer=40.0, batch_size=8,
|
||||
precision=precision,
|
||||
)
|
||||
write_result({
|
||||
"axis": "precision", "precision": precision,
|
||||
"chunk_len": 40.0, "overlap": 0, "batch_size": 8,
|
||||
"decoding": "greedy", "pnc": "yes", "timestamps": False,
|
||||
"audio": name, "duration_audio": dur,
|
||||
"duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
|
||||
"chars": len(text), "text": text,
|
||||
})
|
||||
except Exception as e:
|
||||
log(f" FAIL {axis} on {name}: {e}")
|
||||
write_result({"axis": "precision", "precision": precision,
|
||||
"audio": name, "error": str(e)})
|
||||
del model
|
||||
torch.cuda.empty_cache()
|
||||
except Exception as e:
|
||||
log(f"FAIL precision={precision}: {e}")
|
||||
|
||||
# All remaining axes use bf16 default model.
|
||||
torch.cuda.empty_cache()
|
||||
model = load_model("bf16")
|
||||
cfg = prep_cfg(model)
|
||||
|
||||
# ===== chunk_len axis (A.1) =====
|
||||
for chunk_len in [10.0, 20.0, 40.0, 60.0]:
|
||||
for name in SUBSET:
|
||||
wav, dur = wavs[name]
|
||||
axis = f"A1_chunk_len={chunk_len}"
|
||||
log(f"running {axis} on {name}")
|
||||
try:
|
||||
text, t, vram = run_buffered(
|
||||
model, cfg, wav, dur,
|
||||
chunk_len=chunk_len, total_buffer=chunk_len, batch_size=8,
|
||||
precision="bf16",
|
||||
)
|
||||
write_result({
|
||||
"axis": "chunk_len", "chunk_len": chunk_len,
|
||||
"overlap": 0, "batch_size": 8, "precision": "bf16",
|
||||
"decoding": "greedy", "pnc": "yes", "timestamps": False,
|
||||
"audio": name, "duration_audio": dur,
|
||||
"duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
|
||||
"chars": len(text), "text": text,
|
||||
})
|
||||
except Exception as e:
|
||||
log(f" FAIL: {e}")
|
||||
write_result({"axis": "chunk_len", "chunk_len": chunk_len,
|
||||
"audio": name, "error": str(e)})
|
||||
|
||||
# ===== overlap axis (A.2) — chunk_len=40, total_buffer=chunk_len+overlap =====
|
||||
for overlap in [1.0, 2.0, 5.0, 10.0]:
|
||||
chunk_len = 40.0
|
||||
total_buffer = chunk_len + overlap
|
||||
for name in SUBSET:
|
||||
wav, dur = wavs[name]
|
||||
axis = f"A2_overlap={overlap}"
|
||||
log(f"running {axis} on {name}")
|
||||
try:
|
||||
text, t, vram = run_buffered(
|
||||
model, cfg, wav, dur,
|
||||
chunk_len=chunk_len, total_buffer=total_buffer,
|
||||
batch_size=8, precision="bf16",
|
||||
)
|
||||
write_result({
|
||||
"axis": "overlap", "chunk_len": chunk_len,
|
||||
"overlap": overlap, "batch_size": 8, "precision": "bf16",
|
||||
"decoding": "greedy", "pnc": "yes", "timestamps": False,
|
||||
"audio": name, "duration_audio": dur,
|
||||
"duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
|
||||
"chars": len(text), "text": text,
|
||||
})
|
||||
except Exception as e:
|
||||
log(f" FAIL: {e}")
|
||||
write_result({"axis": "overlap", "overlap": overlap,
|
||||
"audio": name, "error": str(e)})
|
||||
|
||||
# ===== batch_size axis (B.1) =====
|
||||
for batch_size in [1, 2, 4, 8, 16]:
|
||||
for name in SINGLE:
|
||||
wav, dur = wavs[name]
|
||||
axis = f"B1_batch={batch_size}"
|
||||
log(f"running {axis} on {name}")
|
||||
try:
|
||||
text, t, vram = run_buffered(
|
||||
model, cfg, wav, dur,
|
||||
chunk_len=40.0, total_buffer=40.0, batch_size=batch_size,
|
||||
precision="bf16",
|
||||
)
|
||||
write_result({
|
||||
"axis": "batch_size", "batch_size": batch_size,
|
||||
"chunk_len": 40.0, "overlap": 0, "precision": "bf16",
|
||||
"decoding": "greedy", "pnc": "yes", "timestamps": False,
|
||||
"audio": name, "duration_audio": dur,
|
||||
"duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
|
||||
"chars": len(text), "text": text,
|
||||
})
|
||||
except Exception as e:
|
||||
log(f" FAIL: {e}")
|
||||
write_result({"axis": "batch_size", "batch_size": batch_size,
|
||||
"audio": name, "error": str(e)})
|
||||
|
||||
# ===== pnc axis (C.2) =====
|
||||
for pnc in ["no", "yes"]:
|
||||
for name in SUBSET:
|
||||
wav, dur = wavs[name]
|
||||
axis = f"C2_pnc={pnc}"
|
||||
log(f"running {axis} on {name}")
|
||||
try:
|
||||
text, t, vram = run_buffered(
|
||||
model, cfg, wav, dur,
|
||||
chunk_len=40.0, total_buffer=40.0, batch_size=8,
|
||||
precision="bf16", pnc=pnc,
|
||||
)
|
||||
write_result({
|
||||
"axis": "pnc", "pnc": pnc,
|
||||
"chunk_len": 40.0, "overlap": 0, "batch_size": 8,
|
||||
"precision": "bf16", "decoding": "greedy",
|
||||
"timestamps": False,
|
||||
"audio": name, "duration_audio": dur,
|
||||
"duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
|
||||
"chars": len(text), "text": text,
|
||||
})
|
||||
except Exception as e:
|
||||
log(f" FAIL: {e}")
|
||||
write_result({"axis": "pnc", "pnc": pnc,
|
||||
"audio": name, "error": str(e)})
|
||||
|
||||
# ===== timestamps axis (C.3) — when on, NeMo recommends chunk_len=10 =====
|
||||
for ts_flag, chunk_len in [(False, 40.0), (True, 10.0)]:
|
||||
for name in SUBSET:
|
||||
wav, dur = wavs[name]
|
||||
axis = f"C3_timestamps={ts_flag}"
|
||||
log(f"running {axis} (chunk_len={chunk_len}) on {name}")
|
||||
try:
|
||||
text, t, vram = run_buffered(
|
||||
model, cfg, wav, dur,
|
||||
chunk_len=chunk_len, total_buffer=chunk_len, batch_size=8,
|
||||
precision="bf16", timestamps=ts_flag,
|
||||
)
|
||||
write_result({
|
||||
"axis": "timestamps", "timestamps": ts_flag,
|
||||
"chunk_len": chunk_len, "overlap": 0, "batch_size": 8,
|
||||
"precision": "bf16", "decoding": "greedy", "pnc": "yes",
|
||||
"audio": name, "duration_audio": dur,
|
||||
"duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
|
||||
"chars": len(text), "text": text,
|
||||
})
|
||||
except Exception as e:
|
||||
log(f" FAIL: {e}")
|
||||
write_result({"axis": "timestamps", "timestamps": ts_flag,
|
||||
"audio": name, "error": str(e)})
|
||||
|
||||
# ===== beam search (C.1) =====
|
||||
try:
|
||||
from nemo.collections.asr.parts.submodules.multitask_decoding import MultiTaskDecodingConfig
|
||||
for strategy, beam_size in [("beam", 4), ("greedy", 1)]:
|
||||
decoding_cfg = MultiTaskDecodingConfig()
|
||||
decoding_cfg.strategy = strategy
|
||||
if hasattr(decoding_cfg, "beam"):
|
||||
decoding_cfg.beam.beam_size = beam_size
|
||||
try:
|
||||
model.change_decoding_strategy(decoding_cfg)
|
||||
except Exception as e:
|
||||
log(f"change_decoding_strategy failed for {strategy}: {e}")
|
||||
continue
|
||||
for name in SUBSET:
|
||||
wav, dur = wavs[name]
|
||||
axis = f"C1_decoding={strategy}_beam={beam_size}"
|
||||
log(f"running {axis} on {name}")
|
||||
try:
|
||||
text, t, vram = run_buffered(
|
||||
model, cfg, wav, dur,
|
||||
chunk_len=40.0, total_buffer=40.0, batch_size=8,
|
||||
precision="bf16",
|
||||
)
|
||||
write_result({
|
||||
"axis": "decoding", "decoding": strategy,
|
||||
"beam_size": beam_size,
|
||||
"chunk_len": 40.0, "overlap": 0, "batch_size": 8,
|
||||
"precision": "bf16", "pnc": "yes", "timestamps": False,
|
||||
"audio": name, "duration_audio": dur,
|
||||
"duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
|
||||
"chars": len(text), "text": text,
|
||||
})
|
||||
except Exception as e:
|
||||
log(f" FAIL: {e}")
|
||||
write_result({"axis": "decoding", "decoding": strategy,
|
||||
"audio": name, "error": str(e)})
|
||||
except Exception as e:
|
||||
log(f"beam search axis skipped: {e}")
|
||||
|
||||
# ===== Path-1 (.transcribe() auto-chunking) — D.1 =====
|
||||
log("=== Path 1: model.transcribe() with auto-chunking ===")
|
||||
for name in SUBSET:
|
||||
wav, dur = wavs[name]
|
||||
axis = "D1_path1_transcribe"
|
||||
log(f"running {axis} on {name}")
|
||||
try:
|
||||
text, t, vram = run_transcribe_path1(
|
||||
model, wav, batch_size=1, pnc="yes",
|
||||
timestamps=False, precision="bf16",
|
||||
)
|
||||
write_result({
|
||||
"axis": "path1", "path": "transcribe_auto_chunking",
|
||||
"batch_size": 1, "precision": "bf16",
|
||||
"audio": name, "duration_audio": dur,
|
||||
"duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
|
||||
"chars": len(text), "text": text,
|
||||
})
|
||||
except Exception as e:
|
||||
log(f" FAIL: {e}")
|
||||
write_result({"axis": "path1", "audio": name, "error": str(e)})
|
||||
|
||||
log("=== sweep complete ===")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sweep()
|
||||
@@ -0,0 +1,125 @@
|
||||
"""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()
|
||||
@@ -0,0 +1,178 @@
|
||||
"""Threshold sweep — varying audio length around the 40 s Encoder-Window.
|
||||
|
||||
Compares three transcription paths on artificially trimmed audios:
|
||||
- path1_default: model.transcribe([wav]) — what main.py uses for ≤ 25 s
|
||||
- path1_bs1: model.transcribe([wav], batch_size=1) — NeMo ≥ 2.5 auto-chunking
|
||||
- path2: FrameBatchMultiTaskAED with chunk_len=40 — what main.py
|
||||
uses for > 25 s
|
||||
|
||||
Audios are 20, 30, 38 and 45 s slices of the 80.9 s source. The 25 s
|
||||
threshold in main.py sits in the middle, so this sweep tells us:
|
||||
- Does Path 1 truncate already at 30/38 s, or only > 40 s?
|
||||
- Does the threshold need to move?
|
||||
"""
|
||||
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"
|
||||
OUT_DIR = "/sweep"
|
||||
RESULTS_PATH = f"{OUT_DIR}/threshold_results.jsonl"
|
||||
LOG_PATH = f"{OUT_DIR}/threshold_log.txt"
|
||||
|
||||
|
||||
def log(msg: str) -> None:
|
||||
line = f"{time.strftime('%H:%M:%S')} {msg}"
|
||||
print(line, flush=True)
|
||||
with open(LOG_PATH, "a") as f:
|
||||
f.write(line + "\n")
|
||||
|
||||
|
||||
def transcode(src: str) -> str:
|
||||
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_duration(p: str) -> float:
|
||||
with wave.open(p, "rb") as w:
|
||||
return w.getnframes() / float(w.getframerate())
|
||||
|
||||
|
||||
def vram_mb() -> float:
|
||||
return torch.cuda.max_memory_allocated() / (1024 * 1024) if torch.cuda.is_available() else 0
|
||||
|
||||
|
||||
def write_manifest(wav_path, duration, pnc="yes"):
|
||||
f = tempfile.NamedTemporaryFile(suffix=".json", mode="w", delete=False)
|
||||
json.dump({"audio_filepath": wav_path, "duration": duration,
|
||||
"taskname": "asr", "source_lang": "de", "target_lang": "de",
|
||||
"pnc": pnc}, f)
|
||||
f.write("\n")
|
||||
f.close()
|
||||
return f.name
|
||||
|
||||
|
||||
def run_path1(model, wav, batch_size):
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
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=batch_size,
|
||||
)
|
||||
return (getattr(hyps[0], "text", None) or str(hyps[0])).strip(), \
|
||||
time.monotonic() - t0, vram_mb()
|
||||
|
||||
|
||||
def run_path2(model, model_cfg, wav, dur, chunk_len=40.0):
|
||||
manifest = write_manifest(wav, dur)
|
||||
feature_stride = model_cfg.preprocessor["window_stride"]
|
||||
model_stride_secs = feature_stride * 8
|
||||
frame_asr = FrameBatchMultiTaskAED(
|
||||
asr_model=model, frame_len=chunk_len,
|
||||
total_buffer=chunk_len, batch_size=8,
|
||||
)
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
t0 = time.monotonic()
|
||||
with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
|
||||
with torch.no_grad():
|
||||
hyps = get_buffered_pred_feat_multitaskAED(
|
||||
frame_asr, model_cfg.preprocessor, model_stride_secs,
|
||||
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, vram_mb()
|
||||
|
||||
|
||||
def main():
|
||||
open(LOG_PATH, "w").close()
|
||||
open(RESULTS_PATH, "w").close()
|
||||
|
||||
log("loading model bf16")
|
||||
model = ASRModel.from_pretrained("nvidia/canary-1b-v2")
|
||||
model = model.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)
|
||||
log(f"model loaded, vram={vram_mb():.0f}MB")
|
||||
|
||||
audios = sorted(os.listdir(AUDIO_DIR))
|
||||
wavs = {}
|
||||
for name in audios:
|
||||
if not name.endswith(".m4a"):
|
||||
continue
|
||||
src = os.path.join(AUDIO_DIR, name)
|
||||
wav = transcode(src)
|
||||
dur = wav_duration(wav)
|
||||
wavs[name] = (wav, dur)
|
||||
log(f" {name} → {dur:.1f}s")
|
||||
|
||||
def record(d):
|
||||
with open(RESULTS_PATH, "a") as f:
|
||||
f.write(json.dumps(d, ensure_ascii=False) + "\n")
|
||||
|
||||
for name, (wav, dur) in wavs.items():
|
||||
# path 1 default (no batch_size override → triggers no auto-chunking
|
||||
# because input is a list, not a single string)
|
||||
log(f"path1_default on {name}")
|
||||
try:
|
||||
text, t, vram = run_path1(model, wav, batch_size=4)
|
||||
record({"audio": name, "duration_audio": dur, "path": "path1_default",
|
||||
"batch_size": 4, "duration_infer": round(t, 2),
|
||||
"vram_peak_mb": round(vram, 0), "chars": len(text), "text": text})
|
||||
except Exception as e:
|
||||
log(f" FAIL: {e}")
|
||||
record({"audio": name, "path": "path1_default", "error": str(e)})
|
||||
|
||||
# path 1 with batch_size=1 (auto-chunking trigger per NeMo docs)
|
||||
log(f"path1_bs1 on {name}")
|
||||
try:
|
||||
text, t, vram = run_path1(model, wav, batch_size=1)
|
||||
record({"audio": name, "duration_audio": dur, "path": "path1_bs1",
|
||||
"batch_size": 1, "duration_infer": round(t, 2),
|
||||
"vram_peak_mb": round(vram, 0), "chars": len(text), "text": text})
|
||||
except Exception as e:
|
||||
log(f" FAIL: {e}")
|
||||
record({"audio": name, "path": "path1_bs1", "error": str(e)})
|
||||
|
||||
# path 2 (always chunked, chunk_len=40)
|
||||
log(f"path2 on {name}")
|
||||
try:
|
||||
text, t, vram = run_path2(model, cfg, wav, dur, chunk_len=40.0)
|
||||
record({"audio": name, "duration_audio": dur, "path": "path2",
|
||||
"chunk_len": 40.0, "batch_size": 8,
|
||||
"duration_infer": round(t, 2),
|
||||
"vram_peak_mb": round(vram, 0), "chars": len(text), "text": text})
|
||||
except Exception as e:
|
||||
log(f" FAIL: {e}")
|
||||
record({"audio": name, "path": "path2", "error": str(e)})
|
||||
|
||||
log("done")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user