Coverage for transformer_lens/benchmarks/audio.py: 22%
127 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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 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_text_inputs(bridge: TransformerBridge):
32 """Build audio-conditioned inputs (synthetic waveform + audio token) for an
33 audio-text decoder; ``(None, None)`` if the processor has no audio path."""
34 processor = getattr(bridge, "processor", None)
35 audio_token = getattr(processor, "audio_token", None) if processor is not None else None
36 if processor is None or audio_token is None: 36 ↛ 38line 36 didn't jump to line 38 because the condition on line 36 was always true
37 return None, None
38 import numpy as np
40 sr = 16000
41 t = np.linspace(0, 1.0, sr, endpoint=False, dtype=np.float32)
42 audio = (0.1 * np.sin(2 * np.pi * (200 + 400 * t) * t)).astype(np.float32)
43 prompt = f"{audio_token}\nTranscribe this audio."
44 try:
45 inputs = processor(text=prompt, audio=audio, sampling_rate=sr, return_tensors="pt")
46 input_ids = inputs["input_ids"].to(bridge.cfg.device)
47 extra = {
48 k: (v.to(bridge.cfg.device) if hasattr(v, "to") else v)
49 for k, v in inputs.items()
50 if k != "input_ids"
51 }
52 return input_ids, extra
53 except Exception:
54 return None, None
57def benchmark_audio_text_forward(bridge: TransformerBridge) -> BenchmarkResult:
58 """Benchmark the audio-conditioned forward (input_features -> finite logits) of
59 an audio-text decoder -- the audio path that the image and encoder benchmarks do not cover."""
60 if not getattr(bridge.cfg, "is_multimodal", False):
61 return BenchmarkResult(
62 name="audio_text_forward",
63 severity=BenchmarkSeverity.SKIPPED,
64 message="Skipped: model is not multimodal",
65 )
66 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""):
67 return BenchmarkResult(
68 name="audio_text_forward",
69 severity=BenchmarkSeverity.INFO,
70 message="Skipped for tiny/test model",
71 )
73 input_ids, extra = _prepare_audio_text_inputs(bridge)
74 if input_ids is None: 74 ↛ 81line 74 didn't jump to line 81 because the condition on line 74 was always true
75 return BenchmarkResult(
76 name="audio_text_forward",
77 severity=BenchmarkSeverity.SKIPPED,
78 message="Skipped: processor could not build audio inputs (no audio_token?)",
79 )
81 try:
82 with torch.no_grad():
83 out = bridge(input_ids, return_type="logits", **extra)
84 logits = out if isinstance(out, torch.Tensor) else getattr(out, "logits", None)
85 if logits is None:
86 return BenchmarkResult(
87 name="audio_text_forward",
88 severity=BenchmarkSeverity.DANGER,
89 message="Audio-conditioned forward returned no logits",
90 passed=False,
91 )
92 d_vocab = getattr(bridge.cfg, "d_vocab", None)
93 shape_ok = logits.ndim == 3 and (d_vocab is None or logits.shape[-1] == d_vocab)
94 finite = bool(torch.isfinite(logits).all())
95 if not finite or not shape_ok:
96 return BenchmarkResult(
97 name="audio_text_forward",
98 severity=BenchmarkSeverity.DANGER,
99 message=f"Audio-conditioned forward produced invalid logits (finite={finite}, shape={tuple(logits.shape)})",
100 details={"logits_shape": list(logits.shape), "all_finite": finite},
101 passed=False,
102 )
103 return BenchmarkResult(
104 name="audio_text_forward",
105 severity=BenchmarkSeverity.INFO,
106 message=f"Audio-conditioned forward OK: finite logits {tuple(logits.shape)}",
107 details={
108 "logits_shape": list(logits.shape),
109 "audio_feature_keys": [k for k in extra if "feature" in k or "audio" in k],
110 },
111 )
112 except Exception as e:
113 return BenchmarkResult(
114 name="audio_text_forward",
115 severity=BenchmarkSeverity.ERROR,
116 message=f"Audio-conditioned forward failed: {str(e)}",
117 passed=False,
118 )
121def benchmark_audio_forward(
122 bridge: TransformerBridge,
123 test_audio: torch.Tensor,
124 reference_model: Optional[torch.nn.Module] = None,
125) -> BenchmarkResult:
126 """Benchmark forward pass with audio input.
128 Compares bridge output against HF native model on the same waveform.
129 For bare encoder models, compares last_hidden_state. For CTC models,
130 compares logits.
132 Args:
133 bridge: TransformerBridge model to test
134 test_audio: Audio waveform tensor [batch, num_samples]
135 reference_model: Optional HF reference model for comparison
136 """
137 return benchmark_encoder_forward(
138 bridge,
139 test_audio,
140 name="audio_forward",
141 ref_input_key="input_values",
142 reference_model=reference_model,
143 )
146def benchmark_audio_cache(
147 bridge: TransformerBridge,
148 test_audio: torch.Tensor,
149) -> BenchmarkResult:
150 """Benchmark run_with_cache() for audio models.
152 Verifies that critical audio-specific hooks fire and produce valid tensors.
154 Args:
155 bridge: TransformerBridge model to test
156 test_audio: Audio waveform tensor [batch, num_samples]
157 """
158 return benchmark_encoder_cache(
159 bridge,
160 test_audio,
161 name="audio_cache",
162 critical_components=_CRITICAL_AUDIO_COMPONENTS,
163 )
166def benchmark_audio_representation_stability(
167 bridge: TransformerBridge,
168 test_audio: torch.Tensor,
169) -> BenchmarkResult:
170 """Benchmark representation stability under small input perturbations.
172 Verifies that the model produces stable representations: similar audio
173 inputs should produce similar hidden states. Skip for tiny-random models
174 (random weights won't produce stable representations).
176 Args:
177 bridge: TransformerBridge model to test
178 test_audio: Audio waveform tensor [batch, num_samples]
179 """
180 return benchmark_encoder_representation_stability(
181 bridge,
182 test_audio,
183 name="audio_representation_stability",
184 )
187def benchmark_audio_feature_extractor(
188 bridge: TransformerBridge,
189 test_audio: torch.Tensor,
190) -> BenchmarkResult:
191 """Verify CNN feature extractor hook outputs.
193 Checks that the audio_feature_extractor.hook_out produces tensors with
194 correct shape and non-degenerate values.
196 Args:
197 bridge: TransformerBridge model to test
198 test_audio: Audio waveform tensor [batch, num_samples]
199 """
200 try:
201 if "audio_feature_extractor" not in (bridge.adapter.component_mapping or {}):
202 return BenchmarkResult(
203 name="audio_feature_extractor",
204 severity=BenchmarkSeverity.SKIPPED,
205 message="Skipped: architecture has no conv feature extractor (spectrogram input)",
206 )
208 with torch.no_grad():
209 _, cache = bridge.run_with_cache(test_audio)
211 hook_key = "audio_feature_extractor.hook_out"
212 if hook_key not in cache:
213 return BenchmarkResult(
214 name="audio_feature_extractor",
215 severity=BenchmarkSeverity.DANGER,
216 message=f"Hook '{hook_key}' not found in cache",
217 passed=False,
218 )
220 features = cache[hook_key]
222 # Check shape: should be [batch, conv_dim, num_frames]
223 if features.dim() != 3:
224 return BenchmarkResult(
225 name="audio_feature_extractor",
226 severity=BenchmarkSeverity.DANGER,
227 message=f"Expected 3D tensor [batch, conv_dim, frames], got {features.dim()}D",
228 passed=False,
229 details={"shape": str(features.shape)},
230 )
232 # Check for degenerate values
233 is_all_zeros = features.abs().max().item() == 0
234 has_nan = torch.isnan(features).any().item()
235 has_inf = torch.isinf(features).any().item()
237 if is_all_zeros or has_nan or has_inf:
238 issues = []
239 if is_all_zeros:
240 issues.append("all zeros")
241 if has_nan:
242 issues.append("NaN")
243 if has_inf:
244 issues.append("Inf")
245 return BenchmarkResult(
246 name="audio_feature_extractor",
247 severity=BenchmarkSeverity.DANGER,
248 message=f"Degenerate feature values: {', '.join(issues)}",
249 passed=False,
250 details={"shape": str(features.shape), "issues": issues},
251 )
253 return BenchmarkResult(
254 name="audio_feature_extractor",
255 severity=BenchmarkSeverity.INFO,
256 message=f"Feature extractor OK: shape={features.shape}, "
257 f"mean={features.mean().item():.4f}, std={features.std().item():.4f}",
258 details={
259 "shape": str(features.shape),
260 "mean": features.mean().item(),
261 "std": features.std().item(),
262 },
263 )
265 except Exception as e:
266 return BenchmarkResult(
267 name="audio_feature_extractor",
268 severity=BenchmarkSeverity.ERROR,
269 message=f"Feature extractor check failed: {str(e)}",
270 passed=False,
271 )
274def benchmark_audio_ctc_decode(
275 bridge: TransformerBridge,
276) -> BenchmarkResult:
277 """Benchmark CTC decoding for HubertForCTC models.
279 Loads a small sample from librispeech_asr_dummy, decodes via greedy CTC,
280 and reports the decoded text. Skipped for bare encoder models (no CTC head)
281 and tiny-random models.
283 Args:
284 bridge: TransformerBridge model to test
285 """
286 model_name = getattr(bridge.cfg, "model_name", "")
287 if is_tiny_test_model(model_name):
288 return BenchmarkResult(
289 name="audio_ctc_decode",
290 severity=BenchmarkSeverity.SKIPPED,
291 message="Skipped for tiny-random model (untrained CTC head)",
292 )
294 try:
295 from datasets import load_dataset
297 ds = load_dataset(
298 "hf-internal-testing/librispeech_asr_dummy",
299 "clean",
300 split="validation",
301 trust_remote_code=True,
302 )
303 audio = ds[0]["audio"]
304 reference_text = ds[0]["text"]
305 waveform = torch.tensor(audio["array"], dtype=torch.float32).unsqueeze(0)
306 waveform = waveform.to(bridge.cfg.device)
308 with torch.no_grad():
309 output = bridge(waveform, return_type=None)
311 if not hasattr(output, "logits") or output.logits is None:
312 return BenchmarkResult(
313 name="audio_ctc_decode",
314 severity=BenchmarkSeverity.SKIPPED,
315 message="Skipped: model output has no logits (bare encoder)",
316 )
318 # Greedy CTC decode
319 predicted_ids = torch.argmax(output.logits, dim=-1)
321 # Try to decode with processor
322 processor = getattr(bridge, "processor", None)
323 if processor is not None and hasattr(processor, "decode"):
324 decoded_text = processor.decode(predicted_ids[0])
325 elif processor is not None and hasattr(processor, "batch_decode"):
326 decoded_text = processor.batch_decode(predicted_ids)[0]
327 else:
328 decoded_text = str(predicted_ids[0].tolist()[:20]) + "..."
330 return BenchmarkResult(
331 name="audio_ctc_decode",
332 severity=BenchmarkSeverity.INFO,
333 message=f"CTC decode successful",
334 details={
335 "decoded_text": decoded_text[:200],
336 "reference_text": reference_text[:200],
337 "logits_shape": str(output.logits.shape),
338 },
339 )
341 except ImportError:
342 return BenchmarkResult(
343 name="audio_ctc_decode",
344 severity=BenchmarkSeverity.SKIPPED,
345 message="Skipped: 'datasets' package not available",
346 )
347 except Exception as e:
348 return BenchmarkResult(
349 name="audio_ctc_decode",
350 severity=BenchmarkSeverity.ERROR,
351 message=f"CTC decode failed: {str(e)}",
352 passed=False,
353 )
356def run_audio_benchmarks(
357 bridge: TransformerBridge,
358 test_audio: Optional[torch.Tensor] = None,
359 verbose: bool = True,
360) -> List[BenchmarkResult]:
361 """Run all audio benchmarks.
363 Args:
364 bridge: TransformerBridge model to test
365 test_audio: Optional audio input tensor. If None, generates a synthetic input
366 shaped for this architecture (waveform or spectrogram).
367 verbose: Whether to print progress
369 Returns:
370 List of BenchmarkResult objects
371 """
372 if test_audio is None:
373 test_audio = build_modality_input(bridge, device=bridge.cfg.device, dtype=bridge.cfg.dtype)
374 if test_audio is None:
375 return [
376 BenchmarkResult(
377 name="audio_forward",
378 severity=BenchmarkSeverity.ERROR,
379 message="Could not build an audio input for this model",
380 passed=False,
381 )
382 ]
384 results = []
386 if verbose:
387 print("1. Audio Forward Pass")
388 results.append(benchmark_audio_forward(bridge, test_audio))
390 if verbose:
391 print("2. Audio Cache Verification")
392 results.append(benchmark_audio_cache(bridge, test_audio))
394 if verbose:
395 print("3. Representation Stability")
396 results.append(benchmark_audio_representation_stability(bridge, test_audio))
398 if verbose:
399 print("4. Feature Extractor Verification")
400 results.append(benchmark_audio_feature_extractor(bridge, test_audio))
402 if verbose:
403 print("5. CTC Decoding")
404 results.append(benchmark_audio_ctc_decode(bridge))
406 return results