Brummel
6b7a1cea49
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.
426 lines
16 KiB
Python
426 lines
16 KiB
Python
"""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()
|