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()
|
||||
Reference in New Issue
Block a user