Coverage for transformer_lens/benchmarks/audio.py: 26%
145 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Audio benchmarks for TransformerBridge.
3Tests that audio encoder models (HuBERT, wav2vec2, etc.) correctly handle
4audio waveform inputs through forward(), run_with_cache(), and produce
5stable representations.
6"""
8from typing import Any, List, Optional
10import torch
12from transformer_lens.benchmarks.encoder_common import (
13 benchmark_encoder_cache,
14 benchmark_encoder_forward,
15 benchmark_encoder_representation_stability,
16)
17from transformer_lens.benchmarks.utils import (
18 BenchmarkResult,
19 BenchmarkSeverity,
20 build_modality_input,
21 is_tiny_test_model,
22)
23from transformer_lens.model_bridge import TransformerBridge
25# Component-mapping names whose hook_out must be cached when the architecture
26# declares them: waveform encoders (HuBERT, wav2vec2) have a conv feature
27# extractor, while spectrogram encoders (AST) patch-embed the spectrogram directly.
28_CRITICAL_AUDIO_COMPONENTS = ("audio_feature_extractor", "conv_pos_embed", "embed_ln", "embed")
31def _prepare_audio_encoder_input(
32 bridge: Any, test_audio: Optional[torch.Tensor] = None
33) -> torch.Tensor:
34 """Model-ready audio input via the bridge's feature extractor when available.
36 Non-wav2vec2-style architectures (e.g. AST) consume feature-extractor
37 outputs (spectrograms), not raw waveforms, and declare their own sampling
38 rate — so input prep must go through ``bridge.processor`` whenever the
39 boot attached one. Falls back to a raw 16 kHz waveform otherwise.
40 """
41 processor = getattr(bridge, "processor", None)
42 fe = getattr(processor, "feature_extractor", processor)
43 sampling_rate = int(getattr(fe, "sampling_rate", 16000) or 16000)
45 device = bridge.cfg.device
46 dtype = bridge.cfg.dtype
47 if test_audio is None:
48 test_audio = torch.randn(1, sampling_rate, device=device, dtype=dtype)
50 if fe is not None and callable(fe):
51 try:
52 waveforms = [w for w in test_audio.detach().cpu().float().numpy()]
53 out = fe(waveforms, sampling_rate=sampling_rate, return_tensors="pt")
54 prepared = out.get("input_values", out.get("input_features"))
55 if prepared is not None: 55 ↛ 59line 55 didn't jump to line 59 because the condition on line 55 was always true
56 return prepared.to(device=device, dtype=dtype)
57 except Exception:
58 pass # fall through to the raw waveform
59 return test_audio
62def _prepare_audio_text_inputs(bridge: TransformerBridge):
63 """Build audio-conditioned inputs (synthetic waveform + audio token) for an
64 audio-text decoder; ``(None, None)`` if the processor has no audio path."""
65 processor = getattr(bridge, "processor", None)
66 audio_token = getattr(processor, "audio_token", None) if processor is not None else None
67 if processor is None or audio_token is None: 67 ↛ 69line 67 didn't jump to line 69 because the condition on line 67 was always true
68 return None, None
69 import numpy as np
71 sr = 16000
72 t = np.linspace(0, 1.0, sr, endpoint=False, dtype=np.float32)
73 audio = (0.1 * np.sin(2 * np.pi * (200 + 400 * t) * t)).astype(np.float32)
74 prompt = f"{audio_token}\nTranscribe this audio."
75 try:
76 inputs = processor(text=prompt, audio=audio, sampling_rate=sr, return_tensors="pt")
77 input_ids = inputs["input_ids"].to(bridge.cfg.device)
78 extra = {
79 k: (v.to(bridge.cfg.device) if hasattr(v, "to") else v)
80 for k, v in inputs.items()
81 if k != "input_ids"
82 }
83 return input_ids, extra
84 except Exception:
85 return None, None
88def benchmark_audio_text_forward(bridge: TransformerBridge) -> BenchmarkResult:
89 """Benchmark the audio-conditioned forward (input_features -> finite logits) of
90 an audio-text decoder -- the audio path that the image and encoder benchmarks do not cover."""
91 if not getattr(bridge.cfg, "is_multimodal", False):
92 return BenchmarkResult(
93 name="audio_text_forward",
94 severity=BenchmarkSeverity.SKIPPED,
95 message="Skipped: model is not multimodal",
96 )
97 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""):
98 return BenchmarkResult(
99 name="audio_text_forward",
100 severity=BenchmarkSeverity.INFO,
101 message="Skipped for tiny/test model",
102 )
104 input_ids, extra = _prepare_audio_text_inputs(bridge)
105 if input_ids is None: 105 ↛ 112line 105 didn't jump to line 112 because the condition on line 105 was always true
106 return BenchmarkResult(
107 name="audio_text_forward",
108 severity=BenchmarkSeverity.SKIPPED,
109 message="Skipped: processor could not build audio inputs (no audio_token?)",
110 )
112 try:
113 with torch.no_grad():
114 out = bridge(input_ids, return_type="logits", **extra)
115 logits = out if isinstance(out, torch.Tensor) else getattr(out, "logits", None)
116 if logits is None:
117 return BenchmarkResult(
118 name="audio_text_forward",
119 severity=BenchmarkSeverity.DANGER,
120 message="Audio-conditioned forward returned no logits",
121 passed=False,
122 )
123 d_vocab = getattr(bridge.cfg, "d_vocab", None)
124 shape_ok = logits.ndim == 3 and (d_vocab is None or logits.shape[-1] == d_vocab)
125 finite = bool(torch.isfinite(logits).all())
126 if not finite or not shape_ok:
127 return BenchmarkResult(
128 name="audio_text_forward",
129 severity=BenchmarkSeverity.DANGER,
130 message=f"Audio-conditioned forward produced invalid logits (finite={finite}, shape={tuple(logits.shape)})",
131 details={"logits_shape": list(logits.shape), "all_finite": finite},
132 passed=False,
133 )
134 return BenchmarkResult(
135 name="audio_text_forward",
136 severity=BenchmarkSeverity.INFO,
137 message=f"Audio-conditioned forward OK: finite logits {tuple(logits.shape)}",
138 details={
139 "logits_shape": list(logits.shape),
140 "audio_feature_keys": [k for k in extra if "feature" in k or "audio" in k],
141 },
142 )
143 except Exception as e:
144 return BenchmarkResult(
145 name="audio_text_forward",
146 severity=BenchmarkSeverity.ERROR,
147 message=f"Audio-conditioned forward failed: {str(e)}",
148 passed=False,
149 )
152def benchmark_audio_forward(
153 bridge: TransformerBridge,
154 test_audio: torch.Tensor,
155 reference_model: Optional[torch.nn.Module] = None,
156) -> BenchmarkResult:
157 """Benchmark forward pass with audio input.
159 Compares bridge output against HF native model on the same waveform.
160 For bare encoder models, compares last_hidden_state. For CTC models,
161 compares logits.
163 Args:
164 bridge: TransformerBridge model to test
165 test_audio: Audio waveform tensor [batch, num_samples]
166 reference_model: Optional HF reference model for comparison
167 """
168 return benchmark_encoder_forward(
169 bridge,
170 test_audio,
171 name="audio_forward",
172 ref_input_key="input_values",
173 reference_model=reference_model,
174 )
177def benchmark_audio_cache(
178 bridge: TransformerBridge,
179 test_audio: torch.Tensor,
180) -> BenchmarkResult:
181 """Benchmark run_with_cache() for audio models.
183 Verifies that critical audio-specific hooks fire and produce valid tensors.
185 Args:
186 bridge: TransformerBridge model to test
187 test_audio: Audio waveform tensor [batch, num_samples]
188 """
189 return benchmark_encoder_cache(
190 bridge,
191 test_audio,
192 name="audio_cache",
193 critical_components=_CRITICAL_AUDIO_COMPONENTS,
194 )
197def benchmark_audio_representation_stability(
198 bridge: TransformerBridge,
199 test_audio: torch.Tensor,
200) -> BenchmarkResult:
201 """Benchmark representation stability under small input perturbations.
203 Verifies that the model produces stable representations: similar audio
204 inputs should produce similar hidden states. Skip for tiny-random models
205 (random weights won't produce stable representations).
207 Args:
208 bridge: TransformerBridge model to test
209 test_audio: Audio waveform tensor [batch, num_samples]
210 """
211 return benchmark_encoder_representation_stability(
212 bridge,
213 test_audio,
214 name="audio_representation_stability",
215 )
218def benchmark_audio_feature_extractor(
219 bridge: TransformerBridge,
220 test_audio: torch.Tensor,
221) -> BenchmarkResult:
222 """Verify CNN feature extractor hook outputs.
224 Checks that the audio_feature_extractor.hook_out produces tensors with
225 correct shape and non-degenerate values.
227 Args:
228 bridge: TransformerBridge model to test
229 test_audio: Audio waveform tensor [batch, num_samples]
230 """
231 try:
232 if "audio_feature_extractor" not in (bridge.adapter.component_mapping or {}):
233 return BenchmarkResult(
234 name="audio_feature_extractor",
235 severity=BenchmarkSeverity.SKIPPED,
236 message="Skipped: architecture has no conv feature extractor (spectrogram input)",
237 )
239 with torch.no_grad():
240 _, cache = bridge.run_with_cache(test_audio)
242 hook_key = "audio_feature_extractor.hook_out"
243 if hook_key not in cache:
244 return BenchmarkResult(
245 name="audio_feature_extractor",
246 severity=BenchmarkSeverity.DANGER,
247 message=f"Hook '{hook_key}' not found in cache",
248 passed=False,
249 )
251 features = cache[hook_key]
253 # Check shape: should be [batch, conv_dim, num_frames]
254 if features.dim() != 3:
255 return BenchmarkResult(
256 name="audio_feature_extractor",
257 severity=BenchmarkSeverity.DANGER,
258 message=f"Expected 3D tensor [batch, conv_dim, frames], got {features.dim()}D",
259 passed=False,
260 details={"shape": str(features.shape)},
261 )
263 # Check for degenerate values
264 is_all_zeros = features.abs().max().item() == 0
265 has_nan = torch.isnan(features).any().item()
266 has_inf = torch.isinf(features).any().item()
268 if is_all_zeros or has_nan or has_inf:
269 issues = []
270 if is_all_zeros:
271 issues.append("all zeros")
272 if has_nan:
273 issues.append("NaN")
274 if has_inf:
275 issues.append("Inf")
276 return BenchmarkResult(
277 name="audio_feature_extractor",
278 severity=BenchmarkSeverity.DANGER,
279 message=f"Degenerate feature values: {', '.join(issues)}",
280 passed=False,
281 details={"shape": str(features.shape), "issues": issues},
282 )
284 return BenchmarkResult(
285 name="audio_feature_extractor",
286 severity=BenchmarkSeverity.INFO,
287 message=f"Feature extractor OK: shape={features.shape}, "
288 f"mean={features.mean().item():.4f}, std={features.std().item():.4f}",
289 details={
290 "shape": str(features.shape),
291 "mean": features.mean().item(),
292 "std": features.std().item(),
293 },
294 )
296 except Exception as e:
297 return BenchmarkResult(
298 name="audio_feature_extractor",
299 severity=BenchmarkSeverity.ERROR,
300 message=f"Feature extractor check failed: {str(e)}",
301 passed=False,
302 )
305def benchmark_audio_ctc_decode(
306 bridge: TransformerBridge,
307) -> BenchmarkResult:
308 """Benchmark CTC decoding for HubertForCTC models.
310 Loads a small sample from librispeech_asr_dummy, decodes via greedy CTC,
311 and reports the decoded text. Skipped for bare encoder models (no CTC head)
312 and tiny-random models.
314 Args:
315 bridge: TransformerBridge model to test
316 """
317 model_name = getattr(bridge.cfg, "model_name", "")
318 if is_tiny_test_model(model_name):
319 return BenchmarkResult(
320 name="audio_ctc_decode",
321 severity=BenchmarkSeverity.SKIPPED,
322 message="Skipped for tiny-random model (untrained CTC head)",
323 )
325 try:
326 from datasets import load_dataset
328 ds = load_dataset(
329 "hf-internal-testing/librispeech_asr_dummy",
330 "clean",
331 split="validation",
332 trust_remote_code=True,
333 )
334 audio = ds[0]["audio"]
335 reference_text = ds[0]["text"]
336 waveform = torch.tensor(audio["array"], dtype=torch.float32).unsqueeze(0)
337 waveform = waveform.to(bridge.cfg.device)
339 with torch.no_grad():
340 output = bridge(waveform, return_type=None)
342 if not hasattr(output, "logits") or output.logits is None:
343 return BenchmarkResult(
344 name="audio_ctc_decode",
345 severity=BenchmarkSeverity.SKIPPED,
346 message="Skipped: model output has no logits (bare encoder)",
347 )
349 # Greedy CTC decode
350 predicted_ids = torch.argmax(output.logits, dim=-1)
352 # Try to decode with processor
353 processor = getattr(bridge, "processor", None)
354 if processor is not None and hasattr(processor, "decode"):
355 decoded_text = processor.decode(predicted_ids[0])
356 elif processor is not None and hasattr(processor, "batch_decode"):
357 decoded_text = processor.batch_decode(predicted_ids)[0]
358 else:
359 decoded_text = str(predicted_ids[0].tolist()[:20]) + "..."
361 return BenchmarkResult(
362 name="audio_ctc_decode",
363 severity=BenchmarkSeverity.INFO,
364 message=f"CTC decode successful",
365 details={
366 "decoded_text": decoded_text[:200],
367 "reference_text": reference_text[:200],
368 "logits_shape": str(output.logits.shape),
369 },
370 )
372 except ImportError:
373 return BenchmarkResult(
374 name="audio_ctc_decode",
375 severity=BenchmarkSeverity.SKIPPED,
376 message="Skipped: 'datasets' package not available",
377 )
378 except Exception as e:
379 return BenchmarkResult(
380 name="audio_ctc_decode",
381 severity=BenchmarkSeverity.ERROR,
382 message=f"CTC decode failed: {str(e)}",
383 passed=False,
384 )
387def run_audio_benchmarks(
388 bridge: TransformerBridge,
389 test_audio: Optional[torch.Tensor] = None,
390 verbose: bool = True,
391) -> List[BenchmarkResult]:
392 """Run all audio benchmarks.
394 Args:
395 bridge: TransformerBridge model to test
396 test_audio: Optional audio input tensor. If None, generates a synthetic input
397 shaped for this architecture (waveform or spectrogram).
398 verbose: Whether to print progress
400 Returns:
401 List of BenchmarkResult objects
402 """
403 if test_audio is None:
404 test_audio = build_modality_input(bridge, device=bridge.cfg.device, dtype=bridge.cfg.dtype)
405 if test_audio is None:
406 return [
407 BenchmarkResult(
408 name="audio_forward",
409 severity=BenchmarkSeverity.ERROR,
410 message="Could not build an audio input for this model",
411 passed=False,
412 )
413 ]
415 results = []
417 if verbose:
418 print("1. Audio Forward Pass")
419 results.append(benchmark_audio_forward(bridge, test_audio))
421 if verbose:
422 print("2. Audio Cache Verification")
423 results.append(benchmark_audio_cache(bridge, test_audio))
425 if verbose:
426 print("3. Representation Stability")
427 results.append(benchmark_audio_representation_stability(bridge, test_audio))
429 if verbose:
430 print("4. Feature Extractor Verification")
431 results.append(benchmark_audio_feature_extractor(bridge, test_audio))
433 if verbose:
434 print("5. CTC Decoding")
435 results.append(benchmark_audio_ctc_decode(bridge))
437 return results