diff --git a/tests/QUARANTINES.md b/tests/QUARANTINES.md index b257424e0c..090e0675ca 100644 --- a/tests/QUARANTINES.md +++ b/tests/QUARANTINES.md @@ -99,6 +99,7 @@ Big-model adapter tests use `@pytest.mark.slow`, CI tier filters `-m "not slow"` | [`integration/model_bridge/test_raven_adapter.py`:42](integration/model_bridge/test_raven_adapter.py) (skipif + `slow`) | huginn-0125: ~14GB download + ~28GB RAM | | [`integration/model_bridge/test_rwkv7_adapter.py`:36](integration/model_bridge/test_rwkv7_adapter.py) (skipif + `slow`; `importorskip("fla")` at :30) | rwkv7-0.1B-g1: remote code + flash-linear-attention dep | | [`integration/model_bridge/test_bridge_vs_hf_roberta_parity.py`:23](integration/model_bridge/test_bridge_vs_hf_roberta_parity.py) | distilroberta-base 82M; not in the CI cache allowlist (test_roberta_tiny.py keeps the same surfaces in CI on random weights) | +| [`integration/model_bridge/test_whisper_adapter.py`:18](integration/model_bridge/test_whisper_adapter.py) | whisper-tiny 39M; not in the CI cache allowlist (test_whisper_tiny.py keeps the same surfaces in CI on random weights) | | [`integration/test_jacobian_lens_kurtosis.py`](integration/test_jacobian_lens_kurtosis.py) | Qwen3.5-0.8B (~1.6GB) + lens artifacts + wikitext validation split; CPU-feasible (#1539 Tier-1) | **Un-skip:** locally with `HF_TOKEN` sourced (slow-marked files: run the file directly or `-m slow`). diff --git a/tests/integration/model_bridge/test_whisper_adapter.py b/tests/integration/model_bridge/test_whisper_adapter.py new file mode 100644 index 0000000000..22da5f043f --- /dev/null +++ b/tests/integration/model_bridge/test_whisper_adapter.py @@ -0,0 +1,130 @@ +"""Whisper Bridge parity on openai/whisper-tiny against an independent eager HF reference. + +Slow: not in the CI cache (test_whisper_tiny.py keeps the same surfaces in CI on +random weights). The audio is a synthetic tone, so no dataset download. +""" + +import numpy as np +import pytest +import torch +from transformers import WhisperForConditionalGeneration, WhisperProcessor + +from transformer_lens.model_bridge.bridge import TransformerBridge + +MODEL = "openai/whisper-tiny" +PARITY_ATOL = 1e-4 +PROMPT = ["<|startoftranscript|>", "<|en|>", "<|transcribe|>", "<|notimestamps|>"] + +pytestmark = pytest.mark.slow + + +@pytest.fixture(scope="module") +def bridge(): + return TransformerBridge.boot_transformers(MODEL, device="cpu", dtype=torch.float32) + + +@pytest.fixture(scope="module") +def hf_eager(): + return WhisperForConditionalGeneration.from_pretrained( + MODEL, dtype=torch.float32, attn_implementation="eager" + ).eval() + + +@pytest.fixture(scope="module") +def features(bridge): + t = np.arange(16000) / 16000.0 + audio = (0.5 * np.sin(2 * np.pi * 440.0 * t)).astype(np.float32) + extractor = bridge.processor.feature_extractor + return extractor(audio, sampling_rate=16000, return_tensors="pt").input_features + + +@pytest.fixture(scope="module") +def decoder_ids(bridge): + tokenizer = bridge.processor.tokenizer + prompt = tokenizer.convert_tokens_to_ids(PROMPT) + text = tokenizer(" Hello world.", add_special_tokens=False).input_ids + return torch.tensor([prompt + text]) + + +class TestWhisperBridgeCreation: + def test_adapter_and_stacks(self, bridge): + from transformer_lens.model_bridge.supported_architectures.whisper import ( + WhisperArchitectureAdapter, + ) + + assert isinstance(bridge.adapter, WhisperArchitectureAdapter) + assert len(bridge.encoder_blocks) == 4 + assert len(bridge.decoder_blocks) == 4 + assert bridge.cfg.n_layers == 4 + assert bridge.cfg.n_ctx == 448 + + def test_tokenizer_and_processor_loaded(self, bridge): + assert bridge.tokenizer is not None + assert isinstance(bridge.processor, WhisperProcessor) + + def test_features_are_log_mel_window(self, features): + assert features.shape == (1, 80, 3000) + + +class TestWhisperForwardEquivalence: + def test_logits_match_independent_hf_eager(self, bridge, hf_eager, features, decoder_ids): + assert bridge.original_model is not hf_eager + with torch.no_grad(): + bridge_logits = bridge(features, decoder_input=decoder_ids) + hf_logits = hf_eager(input_features=features, decoder_input_ids=decoder_ids).logits + torch.testing.assert_close(bridge_logits, hf_logits, rtol=0.0, atol=PARITY_ATOL) + + def test_k_proj_has_no_bias(self, bridge, hf_eager): + """Whisper's k_proj is bias-free while q/v/out carry biases, on both stacks.""" + for i, block in enumerate(bridge.encoder_blocks): + assert block.attn.k.original_component.bias is None + assert hf_eager.model.encoder.layers[i].self_attn.k_proj.bias is None + for proj in (block.attn.q, block.attn.v, block.attn.o): + assert proj.original_component.bias is not None + for i, block in enumerate(bridge.decoder_blocks): + hf_layer = hf_eager.model.decoder.layers[i] + for attn, hf_attn in ( + (block.self_attn, hf_layer.self_attn), + (block.cross_attn, hf_layer.encoder_attn), + ): + assert attn.k.original_component.bias is None + assert hf_attn.k_proj.bias is None + for proj in (attn.q, attn.v, attn.o): + assert proj.original_component.bias is not None + + def test_shared_prompt_across_clips_matches_hf(self, bridge, hf_eager, features, decoder_ids): + t = np.arange(16000) / 16000.0 + audio = (0.5 * np.sin(2 * np.pi * 220.0 * t)).astype(np.float32) + extractor = bridge.processor.feature_extractor + second = extractor(audio, sampling_rate=16000, return_tensors="pt").input_features + clips = torch.cat([features, second]) + with torch.no_grad(): + logits = bridge(clips, decoder_input=decoder_ids) + per_clip = torch.cat( + [ + hf_eager(input_features=clip[None], decoder_input_ids=decoder_ids).logits + for clip in clips + ] + ) + torch.testing.assert_close(logits, per_clip, rtol=0.0, atol=PARITY_ATOL) + + +class TestWhisperDecoderPrompt: + def test_to_tokens_is_prefix_then_text(self, bridge): + tokenizer = bridge.tokenizer + tokens = bridge.to_tokens(" Hello world.")[0].tolist() + text_ids = tokenizer(" Hello world.", add_special_tokens=False).input_ids + assert tokens == tokenizer.prefix_tokens + text_ids + assert tokenizer.eos_token_id not in tokens + + +class TestWhisperHooks: + def test_encoder_and_decoder_hooks(self, bridge, features, decoder_ids): + with torch.no_grad(): + _, cache = bridge.run_with_cache(features, decoder_input=decoder_ids) + dec_len = decoder_ids.shape[1] + assert cache["conv1.hook_out"].shape == (1, 384, 3000) + assert cache["conv2.hook_out"].shape == (1, 384, 1500) + assert cache["encoder_blocks.3.hook_out"].shape == (1, 1500, 384) + assert cache["decoder_blocks.0.cross_attn.hook_pattern"].shape == (1, 6, dec_len, 1500) + assert cache["decoder_blocks.3.hook_out"].shape == (1, dec_len, 384) diff --git a/tests/integration/model_bridge/test_whisper_tiny.py b/tests/integration/model_bridge/test_whisper_tiny.py new file mode 100644 index 0000000000..8a8bde2697 --- /dev/null +++ b/tests/integration/model_bridge/test_whisper_tiny.py @@ -0,0 +1,374 @@ +"""Tiny from_config Whisper integration tests — CI coverage without a real checkpoint. + +The openai/whisper-tiny parity suite (test_whisper_adapter.py) is ``slow``-marked +since the checkpoint is not in the CI cache. This file builds a random-weight +WhisperForConditionalGeneration with asymmetric stacks (3 encoder / 2 decoder +layers, as on large-v3-turbo and distil-whisper) and exercises the same surfaces: +live wiring, teacher-forced logit and loss parity (including several clips sharing one +prompt), the decoder prompt from to_tokens, encoder and decoder hooks, +the channel-first conv hooks, and an encoder-side intervention. +""" +import copy + +import pytest +import torch +from transformers import AutoModelForSpeechSeq2Seq, WhisperConfig + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, +) +from transformer_lens.model_bridge.supported_architectures.whisper import ( + WhisperArchitectureAdapter, +) + +ARCH = "WhisperForConditionalGeneration" +N_MELS = 16 +N_FRAMES = 20 # max_source_positions; the mel input is twice this (conv2 has stride 2) +D_MODEL = 32 +ENC_LAYERS = 3 +DEC_LAYERS = 2 +DECODER_START = 3 + + +class _Tok: + pass + + +@pytest.fixture(scope="module") +def hf_model(): + torch.manual_seed(0) + cfg = WhisperConfig( + vocab_size=128, + num_mel_bins=N_MELS, + d_model=D_MODEL, + encoder_layers=ENC_LAYERS, + decoder_layers=DEC_LAYERS, + encoder_attention_heads=4, + decoder_attention_heads=4, + encoder_ffn_dim=64, + decoder_ffn_dim=64, + max_source_positions=N_FRAMES, + max_target_positions=24, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + decoder_start_token_id=DECODER_START, + ) + cfg.architectures = [ARCH] + model = AutoModelForSpeechSeq2Seq.from_config(cfg, attn_implementation="eager") + return model.to(torch.float32).eval() + + +@pytest.fixture(scope="module") +def reference_model(hf_model): + """Unwrapped copy: the bridge swaps its components into hf_model in place.""" + return copy.deepcopy(hf_model) + + +@pytest.fixture(scope="module") +def bridge(hf_model, reference_model): + bridge_cfg = build_bridge_config_from_hf(hf_model.config, ARCH, "whisper-tiny", torch.float32) + return TransformerBridge(hf_model, WhisperArchitectureAdapter(bridge_cfg), tokenizer=_Tok()) + + +@pytest.fixture(scope="module") +def features(): + torch.manual_seed(1) + return torch.randn(2, N_MELS, 2 * N_FRAMES) + + +@pytest.fixture(scope="module") +def decoder_ids(): + return torch.tensor([[DECODER_START, 5, 6, 7], [DECODER_START, 9, 10, 11]]) + + +@pytest.fixture(scope="module") +def reference_logits(reference_model, features, decoder_ids): + with torch.no_grad(): + return reference_model(input_features=features, decoder_input_ids=decoder_ids).logits + + +@pytest.fixture(scope="module") +def cache(bridge, features, decoder_ids): + with torch.no_grad(): + _, c = bridge.run_with_cache(features, decoder_input=decoder_ids) + return c + + +class TestWhisperTinyStructure: + def test_asymmetric_stacks(self, bridge): + assert len(bridge.encoder_blocks) == ENC_LAYERS + assert len(bridge.decoder_blocks) == DEC_LAYERS + assert bridge.cfg.n_layers == DEC_LAYERS + + def test_n_ctx_is_decoder_context(self, bridge, hf_model): + assert bridge.cfg.n_ctx == hf_model.config.max_target_positions + + def test_front_end_wraps_live_modules(self, bridge, hf_model): + """The bridge components sit in HF's module tree, so HF's forward fires them.""" + encoder = hf_model.model.encoder + assert encoder.conv1 is bridge.conv1 + assert encoder.conv2 is bridge.conv2 + assert encoder.embed_positions is bridge.pos_embed + assert hf_model.proj_out is bridge.unembed + assert isinstance(bridge.conv1.original_component, torch.nn.Conv1d) + assert isinstance(bridge.pos_embed.original_component, torch.nn.Embedding) + + def test_unembed_tied_to_decoder_embed(self, bridge): + unembed_w = bridge.unembed.original_component.weight + embed_w = bridge.decoder_embed.original_component.weight + assert unembed_w.data_ptr() == embed_w.data_ptr() + + def test_k_proj_has_no_bias(self, bridge): + """Whisper builds k_proj with bias=False; q/v/out keep theirs.""" + attns = [b.attn for b in bridge.encoder_blocks] + attns += [b.self_attn for b in bridge.decoder_blocks] + attns += [b.cross_attn for b in bridge.decoder_blocks] + for attn in attns: + assert attn.k.original_component.bias is None + for proj in (attn.q, attn.v, attn.o): + assert proj.original_component.bias is not None + + +class TestWhisperTinyForwardEquivalence: + def test_logits_match_hf(self, bridge, features, decoder_ids, reference_logits): + with torch.no_grad(): + logits = bridge(features, decoder_input=decoder_ids) + assert logits.shape == (2, decoder_ids.shape[1], 128) + max_diff = (logits - reference_logits).abs().max().item() + assert max_diff < 1e-5, f"Bridge vs HF max diff = {max_diff}" + + def test_run_with_cache_logits_match_hf(self, bridge, features, decoder_ids, reference_logits): + with torch.no_grad(): + logits, _ = bridge.run_with_cache(features, decoder_input=decoder_ids) + assert (logits - reference_logits).abs().max().item() < 1e-5 + + def test_one_dimensional_decoder_input_is_batched(self, bridge, features, decoder_ids): + with torch.no_grad(): + flat = bridge(features[:1], decoder_input=decoder_ids[0]) + batched = bridge(features[:1], decoder_input=decoder_ids[:1]) + assert torch.equal(flat, batched) + + @pytest.mark.parametrize("prompt_shape", ["1d", "single_row"]) + def test_shared_prompt_gives_each_clip_its_own_logits( + self, bridge, reference_model, features, decoder_ids, prompt_shape + ): + prompt = decoder_ids[0] if prompt_shape == "1d" else decoder_ids[:1] + with torch.no_grad(): + logits = bridge(features, decoder_input=prompt) + per_clip = torch.cat( + [ + reference_model( + input_features=features[i : i + 1], decoder_input_ids=decoder_ids[:1] + ).logits + for i in range(features.shape[0]) + ] + ) + assert logits.shape == per_clip.shape + assert (logits - per_clip).abs().max().item() < 1e-5 + assert not torch.allclose(logits[0], logits[1]) + + def test_mismatched_decoder_batch_rejected(self, bridge, features, decoder_ids): + three_rows = torch.cat([decoder_ids, decoder_ids[:1]]) + with pytest.raises(ValueError, match="batch size 3 but the audio features have 2"): + bridge(features, decoder_input=three_rows) + + def test_shared_labels_match_per_clip_logits_and_loss(self, bridge, reference_model, features): + labels = torch.tensor([[5, 6, 7, 2]]) + with torch.no_grad(): + expected = reference_model(input_features=features, labels=labels.repeat(2, 1)) + logits = bridge(features, labels=labels) + loss = bridge(features, labels=labels, return_type="loss") + assert logits.shape == expected.logits.shape + assert (logits - expected.logits).abs().max().item() < 1e-5 + assert torch.allclose(loss, expected.loss, atol=1e-6) + + def test_labels_loss_matches_hf(self, bridge, reference_model, features): + labels = torch.tensor([[5, 6, 7, 2], [9, 10, 11, 2]]) + with torch.no_grad(): + hf_loss = reference_model(input_features=features, labels=labels).loss + loss = bridge(features, labels=labels, return_type="loss") + assert torch.allclose(loss, hf_loss, atol=1e-6) + + +class TestWhisperTinyHooks: + def test_encoder_front_end_shapes(self, cache): + assert cache["conv1.hook_in"].shape == (2, N_MELS, 2 * N_FRAMES) + assert cache["conv1.hook_out"].shape == (2, D_MODEL, 2 * N_FRAMES) + assert cache["conv2.hook_out"].shape == (2, D_MODEL, N_FRAMES) + assert cache["pos_embed.hook_out"].shape == (N_FRAMES, D_MODEL) + + def test_both_stacks_fire(self, cache, decoder_ids): + dec_len = decoder_ids.shape[1] + assert cache[f"encoder_blocks.{ENC_LAYERS - 1}.hook_out"].shape == (2, N_FRAMES, D_MODEL) + assert cache["encoder_ln_final.hook_out"].shape == (2, N_FRAMES, D_MODEL) + assert cache["decoder_embed.hook_out"].shape == (2, dec_len, D_MODEL) + assert cache[f"decoder_blocks.{DEC_LAYERS - 1}.hook_out"].shape == (2, dec_len, D_MODEL) + assert cache["decoder_blocks.0.cross_attn.hook_pattern"].shape == (2, 4, dec_len, N_FRAMES) + assert cache["unembed.hook_out"].shape == (2, dec_len, 128) + + def test_front_end_matches_hf_math(self, cache): + """hook_out on the convs is pre-GELU; the first block sees conv2 + positions.""" + conv2_act = torch.nn.functional.gelu(cache["conv2.hook_out"]).permute(0, 2, 1) + expected = conv2_act + cache["pos_embed.hook_out"] + assert torch.allclose(cache["encoder_blocks.0.hook_in"], expected, atol=1e-6) + + def test_pos_slice_uses_declared_position_axes(self, bridge, features, decoder_ids): + with torch.no_grad(): + _, sliced = bridge.run_with_cache(features, decoder_input=decoder_ids, pos_slice=3) + assert sliced["conv1.hook_out"].shape == (2, D_MODEL, 1) + assert sliced["conv2.hook_out"].shape == (2, D_MODEL, 1) + assert sliced["pos_embed.hook_out"].shape == (1, D_MODEL) + assert sliced["encoder_blocks.0.hook_out"].shape == (2, 1, D_MODEL) + + def test_encoder_patch_matches_hf_with_patched_encoder_outputs( + self, bridge, reference_model, features, decoder_ids, cache + ): + """Overwriting the encoder output through a hook equals handing HF that output.""" + patched = cache["encoder_ln_final.hook_out"].clone() + patched[:, 5:10] = 0.0 + + def patch(tensor, hook): + return patched + + with torch.no_grad(): + logits = bridge.run_with_hooks( + features, + decoder_input=decoder_ids, + fwd_hooks=[("encoder_ln_final.hook_out", patch)], + ) + bridge.reset_hooks() + with torch.no_grad(): + expected = reference_model( + encoder_outputs=(patched,), decoder_input_ids=decoder_ids + ).logits + unpatched = reference_model(input_features=features, decoder_input_ids=decoder_ids) + assert (logits - expected).abs().max().item() < 1e-5 + assert (logits - unpatched.logits).abs().max().item() > 1e-3 + + +SPECIAL_TOKENS = [ + "<|endoftext|>", + "<|startoftranscript|>", + "<|en|>", + "<|translate|>", + "<|transcribe|>", + "<|startoflm|>", + "<|startofprev|>", + "<|nospeech|>", + "<|notimestamps|>", +] + + +@pytest.fixture +def prompt_bridge(tmp_path): + """Tiny Whisper with a real byte-level WhisperTokenizer, built without a download. + + Saved and reloaded so it carries the from_pretrained metadata the boot path relies on. + """ + from transformers import WhisperTokenizer + from transformers.convert_slow_tokenizer import bytes_to_unicode + + from transformer_lens.model_bridge.sources._bridge_builder import ( + configure_tokenizer, + ) + + vocab = {ch: i for i, ch in enumerate(bytes_to_unicode().values())} + for token in SPECIAL_TOKENS: + vocab[token] = len(vocab) + WhisperTokenizer( + vocab=vocab, merges=[], additional_special_tokens=SPECIAL_TOKENS[1:] + ).save_pretrained(tmp_path) + tokenizer = WhisperTokenizer.from_pretrained(tmp_path) + eot = vocab["<|endoftext|>"] + cfg = WhisperConfig( + vocab_size=len(vocab), + num_mel_bins=N_MELS, + d_model=D_MODEL, + encoder_layers=1, + decoder_layers=1, + encoder_attention_heads=4, + decoder_attention_heads=4, + encoder_ffn_dim=64, + decoder_ffn_dim=64, + max_source_positions=N_FRAMES, + max_target_positions=24, + pad_token_id=eot, + bos_token_id=eot, + eos_token_id=eot, + decoder_start_token_id=vocab["<|startoftranscript|>"], + ) + cfg.architectures = [ARCH] + model = AutoModelForSpeechSeq2Seq.from_config(cfg, attn_implementation="eager").eval() + bridge_cfg = build_bridge_config_from_hf(model.config, ARCH, "whisper-tok", torch.float32) + adapter = WhisperArchitectureAdapter(bridge_cfg) + return TransformerBridge(model, adapter, tokenizer=configure_tokenizer(tokenizer, adapter.cfg)) + + +class TestWhisperTinyDecoderPrompt: + def test_to_tokens_is_prefix_then_text(self, prompt_bridge): + """No <|endoftext|> lands inside the prompt, before or after the text.""" + tokenizer = prompt_bridge.tokenizer + tokens = prompt_bridge.to_tokens(" Hello")[0].tolist() + text_ids = tokenizer(" Hello", add_special_tokens=False).input_ids + assert tokens == tokenizer.prefix_tokens + text_ids + assert tokenizer.eos_token_id not in tokens + + def test_language_and_task_prefix(self, prompt_bridge): + tokenizer = prompt_bridge.tokenizer + tokenizer.set_prefix_tokens(language="english", task="transcribe") + tokens = prompt_bridge.to_tokens(" Hello")[0].tolist() + prefix = ["<|startoftranscript|>", "<|en|>", "<|transcribe|>", "<|notimestamps|>"] + assert tokens[:4] == tokenizer.convert_tokens_to_ids(prefix) + assert tokenizer.eos_token_id not in tokens + + def test_prompt_drives_a_multi_clip_forward(self, prompt_bridge): + prompt = prompt_bridge.to_tokens(" Hello") + with torch.no_grad(): + logits = prompt_bridge(torch.randn(3, N_MELS, 2 * N_FRAMES), decoder_input=prompt) + assert logits.shape == (3, prompt.shape[1], prompt_bridge.cfg.d_vocab) + + +class TestWhisperTinyErrors: + def test_decoder_input_required(self, bridge, features): + with pytest.raises(ValueError, match="decoder_input"): + bridge(features) + + def test_text_input_rejected(self, bridge, decoder_ids): + with pytest.raises(ValueError, match="input_features"): + bridge("hello", decoder_input=decoder_ids) + + def test_input_values_rejected(self, bridge, features, decoder_ids): + with pytest.raises(ValueError, match="input_values"): + bridge(features, input_values=features, decoder_input=decoder_ids) + + def test_generate_not_supported(self, bridge, features): + with pytest.raises(NotImplementedError): + bridge.generate(features, max_new_tokens=2, verbose=False) + + def test_compatibility_mode_not_supported(self, bridge): + with pytest.raises(NotImplementedError, match="audio seq2seq"): + bridge.enable_compatibility_mode() + + +class TestWhisperTinyPhase1Benchmarks: + """The verify_models Phase 1 checks, run on the tiny model.""" + + def test_forward_pass_benchmark(self, bridge, reference_model): + from transformer_lens.benchmarks.forward_pass import benchmark_forward_pass + from transformer_lens.benchmarks.utils import build_modality_input + + features = build_modality_input(bridge) + assert features is not None + assert features.shape == (1, N_MELS, 2 * N_FRAMES) + result = benchmark_forward_pass(bridge, features, reference_model=reference_model) + assert result.passed, result.message + + def test_component_benchmark(self, bridge, reference_model): + from transformer_lens.benchmarks.component_benchmark import ( + benchmark_all_components, + ) + + result = benchmark_all_components(bridge, reference_model) + assert result.passed, result.message diff --git a/tests/unit/benchmarks/test_build_modality_input.py b/tests/unit/benchmarks/test_build_modality_input.py index 9d8640f3d6..148f35f462 100644 --- a/tests/unit/benchmarks/test_build_modality_input.py +++ b/tests/unit/benchmarks/test_build_modality_input.py @@ -44,6 +44,18 @@ def test_spectrogram_dimensions_are_not_hardcoded() -> None: assert got.shape == (3, 512, 64) +def test_audio_seq2seq_uses_log_mel_window() -> None: + """Whisper takes channel-first [batch, n_mels, 2 * n_frames]; conv2 halves the frames.""" + bridge = _bridge( + hf_config=SimpleNamespace(num_mel_bins=80, max_source_positions=1500), + is_audio_model=True, + is_audio_seq2seq=True, + ) + got = build_modality_input(bridge, batch_size=2) + assert got is not None + assert got.shape == (2, 80, 3000) + + def test_waveform_audio_without_mel_bins() -> None: """num_mel_bins is what separates a spectrogram encoder from a waveform one.""" bridge = _bridge(hf_config=SimpleNamespace(), is_audio_model=True) diff --git a/tests/unit/model_bridge/sources/test_modality_helpers.py b/tests/unit/model_bridge/sources/test_modality_helpers.py index 522a0acd48..af4f70d6af 100644 --- a/tests/unit/model_bridge/sources/test_modality_helpers.py +++ b/tests/unit/model_bridge/sources/test_modality_helpers.py @@ -1,10 +1,11 @@ """Modality helpers shared by the boot sources. ``load_modality_processor`` must apply its loader table in declared order -(multimodal → audio → vision, last write wins on ``bridge.processor``), swallow -per-loader failures, and keep the torchvision-install retry an AutoProcessor-only -special case. ``skip_tokenizer_for_modality`` must skip exactly the audio/vision -flags (multimodal models keep their text tokenizer). The boot test mocks the +(multimodal → audio → audio seq2seq → vision, last write wins on +``bridge.processor``), swallow per-loader failures, and keep the +torchvision-install retry an AutoProcessor-only special case. +``skip_tokenizer_for_modality`` must skip exactly the audio/vision flags +(multimodal and audio seq2seq models keep their text tokenizer). The boot test mocks the external boundaries in the style of test_boot_architecture_resolution.py; the helpers themselves run real. """ @@ -96,6 +97,19 @@ def test_single_flag_selects_single_loader(self, monkeypatch, flag, loader_name) if name != loader_name: loader.from_pretrained.assert_not_called() + def test_audio_seq2seq_full_processor_replaces_feature_extractor(self, monkeypatch): + """Whisper sets both audio flags; the full processor (feature extractor + + tokenizer) loads after the bare feature extractor and wins.""" + calls: list = [] + _patch_all_loaders(monkeypatch, calls) + bridge = SimpleNamespace() + cfg = SimpleNamespace(is_audio_model=True, is_audio_seq2seq=True) + + load_modality_processor(bridge, cfg, "org/model", False, None) + + assert [c[0] for c in calls] == ["AutoFeatureExtractor", "AutoProcessor"] + assert bridge.processor == "AutoProcessor-processor" + def test_token_and_trust_remote_code_threaded_through(self, monkeypatch): """The token comes in as a parameter (single derivation at the call site), not from a per-block env re-read.""" @@ -198,6 +212,11 @@ def test_multimodal_alone_does_not_skip(self): cfg = SimpleNamespace(is_multimodal=True) assert skip_tokenizer_for_modality(cfg) is False + def test_audio_seq2seq_keeps_tokenizer(self): + """Whisper is an audio model whose decoder emits text tokens.""" + cfg = SimpleNamespace(is_audio_model=True, is_audio_seq2seq=True) + assert skip_tokenizer_for_modality(cfg) is False + class _FakeBridge(TransformerBridge): """TransformerBridge subclass (boot's return hint is beartype-enforced) that skips @@ -250,3 +269,51 @@ def test_visual_boot_gets_image_processor_and_no_tokenizer(self, monkeypatch): image_loader.from_pretrained.assert_called_once_with( "fake-org/tiny-vit", token="boot-token", trust_remote_code=False ) + + +class TestBootTransformersAudioSeq2Seq: + def test_audio_seq2seq_boot_keeps_tokenizer_and_gets_full_processor(self, monkeypatch): + """End-to-end wiring through boot(): an audio seq2seq model (Whisper) still + loads its text tokenizer, and bridge.processor ends on AutoProcessor.""" + calls: list = [] + _patch_all_loaders(monkeypatch, calls) + + cfg = boot_cfg( + architecture="WhisperForConditionalGeneration", + is_audio_model=True, + is_audio_seq2seq=True, + ) + adapter = stub_adapter(cfg) + monkeypatch.setattr(f"{_SOURCE}.build_bridge_config_from_hf", MagicMock(return_value=cfg)) + monkeypatch.setattr( + f"{_SOURCE}.architecture_adapter_factory.ArchitectureAdapterFactory.select_architecture_adapter", + lambda c: adapter, + ) + tokenizer_loader = MagicMock(name="autotokenizer", return_value="base-tokenizer") + monkeypatch.setattr( + "transformer_lens.utilities.hf_utils.autotokenizer_with_special_token_compat", + tokenizer_loader, + ) + monkeypatch.setattr(f"{_SOURCE}.configure_tokenizer", lambda tok, c: f"configured-{tok}") + monkeypatch.setattr(f"{_SOURCE}.TransformerBridge", _FakeBridge) + monkeypatch.setattr( + "transformer_lens.model_bridge.sources.transformers_driver.TransformersDriver", + MagicMock(name="TransformersDriver"), + ) + monkeypatch.setattr( + "transformer_lens.utilities.multi_gpu.find_embedding_device", lambda m: None + ) + hf_model = SimpleNamespace( + config=SimpleNamespace( + architectures=["WhisperForConditionalGeneration"], model_type="whisper" + ) + ) + + from transformer_lens.model_bridge.sources.transformers.source import boot + + bridge = boot("fake-org/tiny-whisper", device="cpu", hf_model=hf_model) + + assert tokenizer_loader.call_args.args[0] == "fake-org/tiny-whisper" + assert bridge.tokenizer == "configured-base-tokenizer" + assert [c[0] for c in calls] == ["AutoFeatureExtractor", "AutoProcessor"] + assert bridge.processor == "AutoProcessor-processor" diff --git a/tests/unit/model_bridge/supported_architectures/test_whisper_adapter.py b/tests/unit/model_bridge/supported_architectures/test_whisper_adapter.py new file mode 100644 index 0000000000..dfc6b8a400 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_whisper_adapter.py @@ -0,0 +1,208 @@ +"""Unit tests for the WhisperArchitectureAdapter. + +Download-free: tiny synthetic configs and structural assertions only. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.factories.architecture_adapter_factory import ( + SUPPORTED_ARCHITECTURES, +) +from transformer_lens.model_bridge.generalized_components import ( + EmbeddingBridge, + NormalizationBridge, + PosEmbedBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, + skip_tokenizer_for_modality, +) +from transformer_lens.model_bridge.supported_architectures.blenderbot import ( + BlenderbotArchitectureAdapter, +) +from transformer_lens.model_bridge.supported_architectures.whisper import ( + WhisperArchitectureAdapter, +) + +ARCH = "WhisperForConditionalGeneration" + + +def _base_cfg(encoder_layers: int = 4, decoder_layers: int = 2) -> TransformerBridgeConfig: + cfg = TransformerBridgeConfig( + d_model=64, + d_head=16, + n_layers=encoder_layers, + n_ctx=448, + n_heads=4, + d_mlp=256, + d_vocab=512, + architecture=ARCH, + ) + cfg.encoder_layers = encoder_layers + cfg.decoder_layers = decoder_layers + cfg.encoder_attention_heads = 4 + cfg.decoder_attention_heads = 4 + cfg.encoder_ffn_dim = 256 + cfg.decoder_ffn_dim = 256 + return cfg + + +@pytest.fixture(scope="class") +def adapter() -> WhisperArchitectureAdapter: + return WhisperArchitectureAdapter(_base_cfg()) + + +class TestWhisperConfig: + def test_asymmetric_stacks_follow_decoder(self, adapter): + """large-v3-turbo (32/4) and distil-whisper (n/2) are asymmetric.""" + assert adapter.cfg.n_layers == 2 + + def test_audio_seq2seq_flags(self, adapter): + assert adapter.cfg.is_audio_model is True + assert adapter.cfg.is_audio_seq2seq is True + assert adapter.cfg.is_multimodal is False + + def test_declares_input_features(self, adapter): + assert adapter.audio_input_name == "input_features" + + def test_generation_and_weight_processing_disabled(self, adapter): + assert adapter.supports_generation is False + assert adapter.supports_fold_ln is False + assert adapter.supports_center_writing_weights is False + + def test_scale_embedding_not_forced_on(self, adapter): + assert not adapter.cfg.scale_embedding + + def test_no_bos_prepended(self, adapter): + """The tokenizer's <|startoftranscript|> prefix starts the decoder prompt.""" + assert adapter.cfg.default_prepend_bos is False + + def test_rejects_asymmetric_heads(self): + cfg = _base_cfg() + cfg.decoder_attention_heads = 8 + with pytest.raises(ValueError, match="attention heads"): + WhisperArchitectureAdapter(cfg) + + def test_rejects_asymmetric_ffn(self): + cfg = _base_cfg() + cfg.decoder_ffn_dim = 512 + with pytest.raises(ValueError, match="FFN dims"): + WhisperArchitectureAdapter(cfg) + + def test_n_ctx_from_max_target_positions(self): + """With no other context field, n_ctx is the decoder's 448, not the 2048 fallback.""" + from transformers import WhisperConfig + + hf_cfg = WhisperConfig( + vocab_size=128, + num_mel_bins=16, + d_model=32, + encoder_layers=2, + decoder_layers=2, + encoder_attention_heads=4, + decoder_attention_heads=4, + max_target_positions=448, + ) + cfg = build_bridge_config_from_hf(hf_cfg, ARCH, "whisper-unit", torch.float32) + assert cfg.n_ctx == 448 + + +class TestWhisperComponentMapping: + def test_conv_front_end_replaces_token_embedding(self, adapter): + mapping = adapter.component_mapping + assert "embed" not in mapping + assert "embed_ln" not in mapping + assert type(mapping["conv1"]) is GeneralizedComponent + assert mapping["conv1"].name == "model.encoder.conv1" + assert type(mapping["conv2"]) is GeneralizedComponent + assert mapping["conv2"].name == "model.encoder.conv2" + assert isinstance(mapping["pos_embed"], PosEmbedBridge) + assert mapping["pos_embed"].name == "model.encoder.embed_positions" + + def test_encoder_front_end_runs_first(self, adapter): + keys = list(adapter.component_mapping) + assert keys[:4] == ["conv1", "conv2", "pos_embed", "encoder_blocks"] + + def test_channel_first_hooks_declare_position_axis(self, adapter): + mapping = adapter.component_mapping + for name in ("conv1", "conv2"): + assert mapping[name].hook_in.pos_dim == 2 + assert mapping[name].hook_out.pos_dim == 2 + # The sinusoidal table is read whole as [n_frames, d_model], no batch dim. + assert mapping["pos_embed"].hook_out.pos_dim == 0 + + def test_decoder_side_and_stack_norms(self, adapter): + mapping = adapter.component_mapping + assert isinstance(mapping["decoder_embed"], EmbeddingBridge) + assert mapping["decoder_embed"].name == "model.decoder.embed_tokens" + assert isinstance(mapping["decoder_pos_embed"], PosEmbedBridge) + assert isinstance(mapping["encoder_ln_final"], NormalizationBridge) + assert mapping["encoder_ln_final"].name == "model.encoder.layer_norm" + assert mapping["decoder_ln_final"].name == "model.decoder.layer_norm" + + def test_unembed_is_proj_out(self, adapter): + unembed = adapter.component_mapping["unembed"] + assert isinstance(unembed, UnembeddingBridge) + assert unembed.name == "proj_out" + + def test_blocks_match_bart_family_layout(self, adapter): + """Block layout is inherited unchanged; only the front end and head differ.""" + reference = BlenderbotArchitectureAdapter(_base_cfg()).component_mapping + mapping = adapter.component_mapping + for stack in ("encoder_blocks", "decoder_blocks"): + assert mapping[stack].name == reference[stack].name + assert set(mapping[stack].submodules) == set(reference[stack].submodules) + assert set(mapping["decoder_blocks"].submodules) == { + "ln1", + "self_attn", + "ln2", + "cross_attn", + "ln3", + "mlp", + } + + +class TestWhisperRegistration: + def test_factory_lookup(self): + assert SUPPORTED_ARCHITECTURES[ARCH] is WhisperArchitectureAdapter + + def test_model_type_detection(self): + from transformer_lens.model_bridge.sources.transformers import ( + determine_architecture_from_hf_config, + ) + + cfg = SimpleNamespace(model_type="whisper", architectures=[]) + assert determine_architecture_from_hf_config(cfg) == ARCH + + def test_loads_via_speech_seq2seq_auto_class(self): + from transformers import AutoModelForSpeechSeq2Seq + + from transformer_lens.model_bridge.sources.transformers.helpers import ( + get_hf_model_class_for_architecture, + ) + + assert get_hf_model_class_for_architecture(ARCH) is AutoModelForSpeechSeq2Seq + + def test_classified_as_seq2seq_not_audio_text(self): + from transformer_lens.utilities.architectures import ( + AUDIO_SEQ2SEQ_ARCHITECTURES, + AUDIO_TEXT_ARCHITECTURES, + SEQ2SEQ_ARCHITECTURES, + classify_architecture, + ) + + assert ARCH in AUDIO_SEQ2SEQ_ARCHITECTURES + assert ARCH not in AUDIO_TEXT_ARCHITECTURES + assert ARCH not in SEQ2SEQ_ARCHITECTURES + assert classify_architecture(ARCH) == "seq2seq" + + def test_tokenizer_kept_despite_audio_flag(self, adapter): + assert skip_tokenizer_for_modality(adapter.cfg) is False diff --git a/tests/unit/tools/model_registry/test_verify_gates.py b/tests/unit/tools/model_registry/test_verify_gates.py index 4054054009..a2637db72b 100644 --- a/tests/unit/tools/model_registry/test_verify_gates.py +++ b/tests/unit/tools/model_registry/test_verify_gates.py @@ -46,6 +46,10 @@ def test_vision_arch_runs_phases_1_and_9(self): # ViT/DeiT declare applicable_phases=[1]; phases 2-4 filter out, 1 and 9 stay. assert _phases_to_run("ViTModel", [1, 2, 3, 4, 9]) == [1, 9] + def test_whisper_runs_phase_1_only(self): + # P2/P4 need generate() and P3 weight processing; both are follow-ups. + assert _phases_to_run("WhisperForConditionalGeneration", [1, 2, 3, 4]) == [1] + def test_unknown_architecture_defaults_to_all_phases(self): assert _phases_to_run("NotARealArchitecture", [1, 2, 3, 4]) == [1, 2, 3, 4] diff --git a/transformer_lens/benchmarks/component_benchmark.py b/transformer_lens/benchmarks/component_benchmark.py index c8ca121f90..f10f8dd04d 100644 --- a/transformer_lens/benchmarks/component_benchmark.py +++ b/transformer_lens/benchmarks/component_benchmark.py @@ -61,6 +61,9 @@ def benchmark_all_components( if getattr(bridge.cfg, "is_audio_model", False): # Audio preprocessing needs waveform input; validated in Phase 8 skip_components.extend(["audio_feature_extractor", "feat_proj", "conv_pos_embed"]) + if getattr(bridge.cfg, "is_audio_seq2seq", False): + # Channel-first spectrogram convs; covered end to end by the forward pass + skip_components.extend(["conv1", "conv2"]) # Run comprehensive benchmark report = benchmarker.benchmark_all_components(skip_components=skip_components) diff --git a/transformer_lens/benchmarks/component_outputs.py b/transformer_lens/benchmarks/component_outputs.py index af0c0bef72..7f1df8a2f8 100644 --- a/transformer_lens/benchmarks/component_outputs.py +++ b/transformer_lens/benchmarks/component_outputs.py @@ -622,7 +622,9 @@ def _test_component( # For embedding components, generate the embedding input once shared_embed_input = None if component_path in ("embed", "encoder_embed", "decoder_embed"): - shared_embed_input = test_inputs.get("modality_input") + # Decoder embeddings take token ids even under an audio/vision encoder. + if component_path != "decoder_embed": + shared_embed_input = test_inputs.get("modality_input") if shared_embed_input is None: batch, seq_len, _ = test_input.shape shared_embed_input = torch.randint( diff --git a/transformer_lens/benchmarks/forward_pass.py b/transformer_lens/benchmarks/forward_pass.py index e285b2c5ef..e807b1cb11 100644 --- a/transformer_lens/benchmarks/forward_pass.py +++ b/transformer_lens/benchmarks/forward_pass.py @@ -87,13 +87,18 @@ def benchmark_forward_pass( decoder_input_ids = _get_decoder_input_ids(bridge.original_model, batch_size) decoder_input_ids = decoder_input_ids.to(tokens.device) extra_kwargs["decoder_input_ids"] = decoder_input_ids + elif getattr(bridge.cfg, "is_audio_seq2seq", False) and isinstance( + test_input, torch.Tensor + ): + decoder_input_ids = _get_decoder_input_ids(bridge.original_model, test_input.shape[0]) + extra_kwargs["decoder_input_ids"] = decoder_input_ids.to(test_input.device) # Run bridge forward pass (use no_grad to match HF reference context — # MPS SDPA can produce different results with vs without gradient tracking) with torch.no_grad(): if _is_audio and isinstance(test_input, torch.Tensor): # Audio models: pass waveform, extract tensor from output - bridge_output_raw = bridge(test_input, return_type="logits") + bridge_output_raw = bridge(test_input, return_type="logits", **extra_kwargs) if isinstance(bridge_output_raw, torch.Tensor): bridge_output = bridge_output_raw elif hasattr(bridge_output_raw, "logits") and bridge_output_raw.logits is not None: @@ -138,7 +143,7 @@ def benchmark_forward_pass( # (input_values for wav2vec2-style, input_features for AST-style) assert reference_model is not None with torch.no_grad(): - hf_output = reference_model(test_input) + hf_output = reference_model(test_input, **extra_kwargs) if hasattr(hf_output, "logits") and hf_output.logits is not None: reference_output = hf_output.logits else: diff --git a/transformer_lens/benchmarks/main_benchmark.py b/transformer_lens/benchmarks/main_benchmark.py index 4e65422345..3200f99570 100644 --- a/transformer_lens/benchmarks/main_benchmark.py +++ b/transformer_lens/benchmarks/main_benchmark.py @@ -37,6 +37,7 @@ from transformer_lens.benchmarks.component_benchmark import benchmark_all_components from transformer_lens.benchmarks.forward_pass import ( _compute_self_target_loss, + _get_decoder_input_ids, benchmark_forward_pass, benchmark_loss_equivalence, ) @@ -900,6 +901,9 @@ def cleanup_model(model, model_name_str: str): _is_visual = bridge_unprocessed is not None and getattr( bridge_unprocessed.cfg, "is_visual_model", False ) + _is_audio_seq2seq = bridge_unprocessed is not None and getattr( + bridge_unprocessed.cfg, "is_audio_seq2seq", False + ) # Shared non-text input (spectrogram, waveform, or pixels) — the same tensor is used # for the HF reference capture and the bridge forward so they stay comparable. _test_modality_input = ( @@ -943,6 +947,14 @@ def cleanup_model(model, model_name_str: str): with torch.no_grad(): if _is_visual: hf_out = hf_model(pixel_values=_test_modality_input) + elif _is_audio_seq2seq: + dec_ids = _get_decoder_input_ids( + hf_model, _test_modality_input.shape[0] + ).to(_test_modality_input.device) + audio_kwargs = { + bridge_unprocessed.adapter.audio_input_name: _test_modality_input + } + hf_out = hf_model(**audio_kwargs, decoder_input_ids=dec_ids) else: hf_out = hf_model(input_values=_test_modality_input) # Bare encoders output last_hidden_state, not logits @@ -1338,12 +1350,14 @@ def cleanup_model(model, model_name_str: str): ) # ======================================================================== - # Phase 8: Audio Tests (only for audio encoder models) + # Phase 8: Audio Tests (only for audio encoder models; audio seq2seq models + # need decoder ids on every forward, and Phase 1 covers their audio path) # Runs before Phase 3 so we can reuse bridge_unprocessed before cleanup. # ======================================================================== if ( bridge_unprocessed is not None and getattr(bridge_unprocessed.cfg, "is_audio_model", False) + and not _is_audio_seq2seq and should_run_phase(8) ): current_phase[0] = 8 diff --git a/transformer_lens/benchmarks/utils.py b/transformer_lens/benchmarks/utils.py index 70de76c9ae..9f192d2b9c 100644 --- a/transformer_lens/benchmarks/utils.py +++ b/transformer_lens/benchmarks/utils.py @@ -137,6 +137,12 @@ def _from_config(name: str, default: int) -> int: batch_size, num_channels, image_size, image_size, device=device, dtype=dtype ) + if getattr(cfg, "is_audio_seq2seq", False): + # Whisper-style log-mel window [batch, n_mels, 2 * n_frames]: conv2 has stride 2. + n_mels = _from_config("num_mel_bins", 80) + n_frames = _from_config("max_source_positions", 1500) + return torch.randn(batch_size, n_mels, 2 * n_frames, device=device, dtype=dtype) + if getattr(cfg, "is_audio_model", False): # num_mel_bins is what distinguishes a spectrogram encoder from a waveform one. num_mel_bins = getattr(hf_config, "num_mel_bins", None) if hf_config is not None else None diff --git a/transformer_lens/config/transformer_bridge_config.py b/transformer_lens/config/transformer_bridge_config.py index 1c85228a98..1cd5a5900e 100644 --- a/transformer_lens/config/transformer_bridge_config.py +++ b/transformer_lens/config/transformer_bridge_config.py @@ -110,6 +110,9 @@ def __init__( attn_implementation: Optional[str] = None, # Audio model configuration is_audio_model: bool = False, + # Audio encoder with a text decoder (Whisper): keeps the tokenizer and + # takes decoder token ids alongside the audio features. + is_audio_seq2seq: bool = False, # Vision model (ViT, DeiT) configuration is_visual_model: bool = False, # Stateful model configuration (e.g., Mamba SSMs use cache_params, @@ -211,6 +214,7 @@ def __init__( self.attn_implementation = attn_implementation # Audio model configuration self.is_audio_model = is_audio_model + self.is_audio_seq2seq = is_audio_seq2seq # Vision model (ViT, DeiT) configuration self.is_visual_model = is_visual_model # Stateful model configuration diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 02a44f9dd3..f64ad97ecf 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -152,6 +152,7 @@ VaultGemmaArchitectureAdapter, ViTArchitectureAdapter, Wav2Vec2ArchitectureAdapter, + WhisperArchitectureAdapter, XGLMArchitectureAdapter, YoutuArchitectureAdapter, Zamba2ArchitectureAdapter, @@ -242,6 +243,7 @@ "Wav2Vec2ForCTC": Wav2Vec2ArchitectureAdapter, "Wav2Vec2ForPreTraining": Wav2Vec2ArchitectureAdapter, "Wav2Vec2Model": Wav2Vec2ArchitectureAdapter, + "WhisperForConditionalGeneration": WhisperArchitectureAdapter, "HunYuanDenseV1ForCausalLM": HunYuanDenseV1ArchitectureAdapter, "Idefics3ForConditionalGeneration": Idefics3ArchitectureAdapter, "InternLM2ForCausalLM": InternLM2ArchitectureAdapter, diff --git a/transformer_lens/model_bridge/architecture_adapter.py b/transformer_lens/model_bridge/architecture_adapter.py index 4464c31caa..7655447442 100644 --- a/transformer_lens/model_bridge/architecture_adapter.py +++ b/transformer_lens/model_bridge/architecture_adapter.py @@ -96,6 +96,10 @@ def native_sampler_kwargs(self, max_new_tokens: int, prompt_len: int) -> Dict[st # Whether Bridge's shifted next-token cross-entropy is meaningful. supports_causal_loss: bool = True + # Keyword the wrapped HF forward takes audio under. Read by the audio seq2seq + # forward path (cfg.is_audio_seq2seq), e.g. Whisper's log-mel "input_features". + audio_input_name: str = "input_values" + # HF-style 4D boolean attention masks keep positions marked True. Override # for architectures whose forward path uses True to block attention. bool_4d_mask_is_keep: bool = True diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 1a4489070b..1cf0b3ef8b 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -262,7 +262,10 @@ def detect_tokenizer_bos_eos(tokenizer: Any) -> tuple[bool, bool]: def skip_tokenizer_for_modality(cfg: Any) -> bool: """True when a source must not auto-load a text tokenizer: audio/vision models use - feature extractors or image processors, not text tokenizers (their repos ship none).""" + feature extractors or image processors, not text tokenizers (their repos ship none). + Audio seq2seq models (Whisper) are the exception: their decoder emits text tokens.""" + if getattr(cfg, "is_audio_seq2seq", False): + return False return bool(getattr(cfg, "is_audio_model", False) or getattr(cfg, "is_visual_model", False)) diff --git a/transformer_lens/model_bridge/sources/_hf_format.py b/transformer_lens/model_bridge/sources/_hf_format.py index 3544bdf1ef..39e6ffe239 100644 --- a/transformer_lens/model_bridge/sources/_hf_format.py +++ b/transformer_lens/model_bridge/sources/_hf_format.py @@ -204,6 +204,9 @@ def legacy_per_layer(base_name: str, global_name: str) -> Any: tl_config.n_ctx = (img_h // patch_h) * (img_w // patch_w) + 1 elif hasattr(source_config, "max_sequence_length"): tl_config.n_ctx = source_config.max_sequence_length + elif hasattr(source_config, "max_target_positions"): + # Audio seq2seq (Whisper): the token context is the decoder's. + tl_config.n_ctx = source_config.max_target_positions else: # ALiBi models (Bloom) have no context length field; 2048 is a safe fallback. tl_config.n_ctx = 2048 @@ -329,6 +332,7 @@ def determine_architecture_from_hf_config(hf_config): "marian": "MarianMTModel", "mbart": "MBartForConditionalGeneration", "pegasus": "PegasusForConditionalGeneration", + "whisper": "WhisperForConditionalGeneration", "seed_oss": "SeedOssForCausalLM", "starcoder2": "Starcoder2ForCausalLM", "nemotron": "NemotronForCausalLM", diff --git a/transformer_lens/model_bridge/sources/transformers/helpers.py b/transformer_lens/model_bridge/sources/transformers/helpers.py index 4f95c337ce..3d091d50e4 100644 --- a/transformer_lens/model_bridge/sources/transformers/helpers.py +++ b/transformer_lens/model_bridge/sources/transformers/helpers.py @@ -18,6 +18,7 @@ def get_hf_model_class_for_architecture(architecture: str): from transformer_lens.utilities.architectures import ( AUDIO_ARCHITECTURES, AUDIO_CLASSIFICATION_ARCHITECTURES, + AUDIO_SEQ2SEQ_ARCHITECTURES, AUDIO_TEXT_ARCHITECTURES, BASE_AUTOMODEL_ARCHITECTURES, MASKED_LM_ARCHITECTURES, @@ -27,7 +28,11 @@ def get_hf_model_class_for_architecture(architecture: str): VISION_CLASSIFICATION_ARCHITECTURES, ) - if architecture in SEQ2SEQ_ARCHITECTURES or architecture in AUDIO_TEXT_ARCHITECTURES: + if architecture in AUDIO_SEQ2SEQ_ARCHITECTURES: + from transformers import AutoModelForSpeechSeq2Seq + + return AutoModelForSpeechSeq2Seq + elif architecture in SEQ2SEQ_ARCHITECTURES or architecture in AUDIO_TEXT_ARCHITECTURES: return AutoModelForSeq2SeqLM elif architecture in MASKED_LM_ARCHITECTURES: return AutoModelForMaskedLM @@ -64,12 +69,13 @@ def get_hf_model_class_for_architecture(architecture: str): # Modality flag → HF auto-loader class, applied in order. Last write wins on -# bridge.processor: multimodal → audio → vision. The flags are disjoint today -# (vit.py is the only is_visual_model adapter and it doesn't set is_multimodal), -# but the order must be preserved for future dual-flag adapters. +# bridge.processor: multimodal → audio → audio seq2seq → vision. Audio seq2seq +# (Whisper) sets is_audio_model too and must follow it, so the full processor +# (feature extractor + tokenizer) replaces the bare feature extractor. _MODALITY_PROCESSOR_LOADERS: list[tuple[str, str]] = [ ("is_multimodal", "AutoProcessor"), ("is_audio_model", "AutoFeatureExtractor"), + ("is_audio_seq2seq", "AutoProcessor"), ("is_visual_model", "AutoImageProcessor"), ] diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index 95ea53873f..e6fb7e4b54 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -318,6 +318,9 @@ from transformer_lens.model_bridge.supported_architectures.vaultgemma import ( VaultGemmaArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.whisper import ( + WhisperArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.xglm import XGLMArchitectureAdapter from transformer_lens.model_bridge.supported_architectures.youtu import YoutuArchitectureAdapter from transformer_lens.model_bridge.supported_architectures.zamba2 import Zamba2ArchitectureAdapter @@ -463,6 +466,7 @@ "T5GemmaArchitectureAdapter", "T5Gemma2ArchitectureAdapter", "ViTArchitectureAdapter", + "WhisperArchitectureAdapter", "XGLMArchitectureAdapter", "Zamba2ArchitectureAdapter", "HrmTextArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/bart.py b/transformer_lens/model_bridge/supported_architectures/bart.py index cb3c59fd91..7d553832fa 100644 --- a/transformer_lens/model_bridge/supported_architectures/bart.py +++ b/transformer_lens/model_bridge/supported_architectures/bart.py @@ -1,5 +1,5 @@ """BART adapter and the shared BART-family encoder-decoder base (BART, Marian, -MBart, Pegasus, Blenderbot, M2M100/NLLB); per-member differences are declarative.""" +MBart, Pegasus, Blenderbot, M2M100/NLLB, Whisper); per-member differences are declarative.""" from typing import Any, Dict @@ -31,8 +31,10 @@ class BartFamilyArchitectureAdapter(ArchitectureAdapter): # layernorm_embedding after the token+position embeds (BART, MBart). has_layernorm_embedding: bool = False # Trailing per-stack layer_norm — the pre-LN members (MBart, Pegasus, - # Blenderbot, M2M100). + # Blenderbot, M2M100, Whisper). has_final_stack_norm: bool = False + # Output projection; Whisper names it proj_out. + unembed_name: str = "lm_head" def __init__(self, cfg: Any) -> None: """Validate the config, set family flags, and build the mapping.""" @@ -156,13 +158,18 @@ def _decoder_block(self) -> BlockBridge: }, ) - def _build_component_mapping(self) -> Dict[str, GeneralizedComponent]: + def _encoder_embeddings(self) -> Dict[str, GeneralizedComponent]: + """Encoder input stage seam; Whisper swaps in its audio conv front end.""" mapping: Dict[str, GeneralizedComponent] = { "embed": EmbeddingBridge(name="model.encoder.embed_tokens"), "pos_embed": PosEmbedBridge(name="model.encoder.embed_positions"), } if self.has_layernorm_embedding: mapping["embed_ln"] = self._norm("model.encoder.layernorm_embedding") + return mapping + + def _build_component_mapping(self) -> Dict[str, GeneralizedComponent]: + mapping = self._encoder_embeddings() mapping["encoder_blocks"] = self._encoder_block() if self.has_final_stack_norm: mapping["encoder_ln_final"] = self._norm("model.encoder.layer_norm") @@ -173,7 +180,7 @@ def _build_component_mapping(self) -> Dict[str, GeneralizedComponent]: mapping["decoder_blocks"] = self._decoder_block() if self.has_final_stack_norm: mapping["decoder_ln_final"] = self._norm("model.decoder.layer_norm") - mapping["unembed"] = UnembeddingBridge(name="lm_head") + mapping["unembed"] = UnembeddingBridge(name=self.unembed_name) return mapping def setup_hook_compatibility(self, bridge: Any) -> None: diff --git a/transformer_lens/model_bridge/supported_architectures/whisper.py b/transformer_lens/model_bridge/supported_architectures/whisper.py new file mode 100644 index 0000000000..0236053340 --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/whisper.py @@ -0,0 +1,77 @@ +"""Whisper (``WhisperForConditionalGeneration``) adapter: a BART-family pre-norm +encoder-decoder whose encoder reads log-mel audio through two convolutions and a +fixed sinusoidal position table instead of token embeddings. The text decoder is +the standard BART-family stack with learned positions and a ``proj_out`` +unembedding tied to ``decoder_embed``; there is no ``final_logits_bias``.""" + +from typing import Any, Dict + +from transformer_lens.model_bridge.generalized_components import PosEmbedBridge +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) +from transformer_lens.model_bridge.supported_architectures.bart import ( + BartFamilyArchitectureAdapter, +) + + +class WhisperArchitectureAdapter(BartFamilyArchitectureAdapter): + """Architecture adapter for WhisperForConditionalGeneration models. + + Inputs: ``bridge(input_features, decoder_input=decoder_ids)``, where + ``input_features`` is the log-mel spectrogram ``[batch, n_mels, 2 * n_frames]`` + from ``bridge.processor`` and ``decoder_ids`` are decoder token ids (e.g. the + start-of-transcript prompt), one row per clip or one row shared by all clips. + ``bridge.to_tokens(text)`` builds that prompt from the tokenizer's prefix tokens + (set the language and task with ``bridge.tokenizer.set_prefix_tokens``). Returns + decoder logits under teacher forcing. + + Encoder front end (channel-first, so its hooks declare position on dim 2): + + - ``conv1`` / ``conv2``: the two ``nn.Conv1d`` layers. HF applies GELU outside + the modules, so ``hook_out`` is pre-activation; ``conv2`` halves the frames. + - ``pos_embed``: the fixed sinusoidal table, read whole as ``[n_frames, d_model]`` + (no batch dim) and added to the transposed ``conv2`` output. + + ``encoder_blocks.{i}.hook_out`` is ``[batch, n_frames, d_model]``: 1500 frames, + one per 20 ms of the 30 s window, on every released checkpoint. + + Attention biases: ``q_proj``, ``v_proj`` and ``out_proj`` carry a bias but + ``k_proj`` does not (HF builds it with ``bias=False``), so ``attn.k.bias`` is + ``None`` on every self- and cross-attention. + + Stacks may be asymmetric (large-v3-turbo: 32 encoder / 4 decoder layers; + distil-whisper: 2 decoder layers); ``n_layers`` follows the decoder, as for + Blenderbot. + """ + + require_symmetric_layers = False + n_layers_from = "decoder" + # Checkpoints ship scale_embedding=False, and HF's forward never applies it. + force_scale_embedding = False + has_final_stack_norm = True + unembed_name = "proj_out" + audio_input_name = "input_features" + # generate() is a follow-up; it needs the audio-conditioned decode loop. + supports_generation = False + # P1 (component + forward parity vs HF) only: P2/P4 need generate() and P3 + # needs weight processing, both follow-ups. + applicable_phases = [1] + + def __init__(self, cfg: Any) -> None: + """Set the audio seq2seq flags on top of the BART-family setup.""" + super().__init__(cfg) + self.cfg.is_audio_model = True + self.cfg.is_audio_seq2seq = True + # The tokenizer's own <|startoftranscript|> prefix starts the decoder prompt; a + # prepended BOS (<|endoftext|>) would land inside it. + self.cfg.default_prepend_bos = False + + def _encoder_embeddings(self) -> Dict[str, GeneralizedComponent]: + conv1 = GeneralizedComponent(name="model.encoder.conv1") + conv2 = GeneralizedComponent(name="model.encoder.conv2") + pos_embed = PosEmbedBridge(name="model.encoder.embed_positions") + conv1.mark_pos_dim(2) + conv2.mark_pos_dim(2) + pos_embed.mark_pos_dim(0) + return {"conv1": conv1, "conv2": conv2, "pos_embed": pos_embed} diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 081e700685..7e5ed1a536 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -613,6 +613,13 @@ def enable_compatibility_mode( "for the hook/component compatibility layer without the weight transforms." ) + if getattr(self.cfg, "is_audio_seq2seq", False): + raise NotImplementedError( + "enable_compatibility_mode() is not supported for audio seq2seq models " + "(e.g. Whisper) yet: weight processing for their audio encoder is not " + "implemented. Use the bridge's native hooks (run_with_cache / " + "run_with_hooks) directly." + ) if getattr(self.cfg, "is_audio_model", False): # Audio encoders have no text embed/unembed for the legacy weight # processing to operate on; without this guard the processing path @@ -2123,7 +2130,10 @@ def forward( Only valid when cfg.is_multimodal is True or cfg.is_visual_model is True. input_values: Optional audio waveform tensor for audio models (e.g., HuBERT). The tensor is passed directly to the underlying HuggingFace model. - Only valid when cfg.is_audio_model is True. + Only valid when cfg.is_audio_model is True. Audio seq2seq models + (cfg.is_audio_seq2seq, e.g. Whisper) instead take their features as + ``input`` and require ``decoder_input`` (decoder token ids) or ``labels``, + with one row per clip or a single row shared by every clip. **kwargs: Additional arguments passed to model Returns: @@ -2192,6 +2202,12 @@ def forward( try: if isinstance(input, (str, list)): + if getattr(self.cfg, "is_audio_seq2seq", False): + raise ValueError( + "Audio seq2seq models take the audio features tensor as `input` " + f"(passed to the model as {self.adapter.audio_input_name!r}); " + "decoder token ids go in decoder_input." + ) if getattr(self.cfg, "is_audio_model", False): raise ValueError( "Audio models require tensor input (raw waveform), not text. " @@ -2299,8 +2315,15 @@ def forward( kwargs["use_cache"] = True if kwargs.pop("use_past_kv_cache", False) or kwargs.get("use_cache", False): kwargs["use_cache"] = True - # Auto-generate decoder_input_ids for encoder-decoder models - if "decoder_input_ids" not in kwargs and labels is None and is_encoder_decoder: + is_audio_seq2seq = getattr(self.cfg, "is_audio_seq2seq", False) + # Auto-generate decoder_input_ids for encoder-decoder models. Audio + # seq2seq is excluded: its encoder input is audio, not decoder tokens. + if ( + "decoder_input_ids" not in kwargs + and labels is None + and is_encoder_decoder + and not is_audio_seq2seq + ): decoder_start_token_id = getattr( self.original_model.config, "decoder_start_token_id", None ) @@ -2341,8 +2364,32 @@ def forward( ) kwargs["input_values"] = input_values + if is_audio_seq2seq: + if input_values is not None: + raise ValueError( + "Audio seq2seq models take the audio features tensor as `input` " + f"(passed to the model as {self.adapter.audio_input_name!r}), " + "not input_values." + ) + if "decoder_input_ids" not in kwargs and labels is None: + raise ValueError( + "Audio seq2seq models require decoder_input (decoder token ids, " + "e.g. the start-of-transcript prompt) or labels." + ) + assert isinstance(input, torch.Tensor) # str/list input is rejected above + n_clips = input.shape[0] + decoder_input_ids = kwargs.get("decoder_input_ids") + if isinstance(decoder_input_ids, torch.Tensor): + kwargs["decoder_input_ids"] = self._align_decoder_batch( + decoder_input_ids, n_clips, "decoder_input" + ) + if labels is not None: + labels = self._align_decoder_batch(labels, n_clips, "labels") + kwargs["labels"] = labels + kwargs[self.adapter.audio_input_name] = input + result = self._driver.forward(**kwargs) # Audio models use input_values (waveform), not input_ids - if getattr(self.cfg, "is_audio_model", False): + elif getattr(self.cfg, "is_audio_model", False): if input_values is not None: result = self._driver.forward(**kwargs) elif isinstance(input, torch.Tensor): @@ -2438,6 +2485,24 @@ def forward( if start_at_layer is not None: self._teardown_start_at_layer() + @staticmethod + def _align_decoder_batch(ids: torch.Tensor, n_clips: int, name: str) -> torch.Tensor: + """Line decoder ids up with the audio batch: one row per clip, or one row shared by all. + + HF does not check this, and a single row against several clips returns one row + of logits that matches none of them. + """ + if ids.ndim == 1: + ids = ids.unsqueeze(0) + if ids.shape[0] == n_clips: + return ids + if ids.shape[0] == 1: + return ids.repeat(n_clips, 1) + raise ValueError( + f"{name} has batch size {ids.shape[0]} but the audio features have {n_clips}; " + "pass one row per clip, or a single row shared by every clip." + ) + def _setup_start_at_layer(self, input: Any, start_at_layer: int) -> torch.Tensor: """Arm residual-stream injection at block ``start_at_layer``. diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index 5388076be1..4caf1d9ab6 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -124,6 +124,7 @@ "Wav2Vec2ForCTC", "Wav2Vec2ForPreTraining", "Wav2Vec2Model", + "WhisperForConditionalGeneration", "HrmTextForCausalLM", "HunYuanDenseV1ForCausalLM", "Idefics3ForConditionalGeneration", @@ -288,6 +289,7 @@ "Wav2Vec2ForCTC": ["facebook"], "Wav2Vec2ForPreTraining": ["facebook"], "Wav2Vec2Model": ["facebook"], + "WhisperForConditionalGeneration": ["openai", "distil-whisper"], "HrmTextForCausalLM": ["sapientinc"], "HunYuanDenseV1ForCausalLM": ["tencent"], "Idefics3ForConditionalGeneration": ["HuggingFaceM4", "ibm-granite"], diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index 2f26a8f2d6..8e53507cb6 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/transformer_lens/tools/model_registry/data/supported_models.json @@ -6,9 +6,9 @@ "min_downloads": 500, "scan_duration_seconds": 0.1 }, - "total_architectures": 146, - "total_models": 15736, - "total_verified": 1198, + "total_architectures": 147, + "total_models": 15753, + "total_verified": 1200, "total_provisional": 7, "models": [ { @@ -221575,6 +221575,261 @@ "phase7_score": null, "phase8_score": null, "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-tiny", + "status": 1, + "verified_date": "2026-10-22", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-tiny.en", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-base.en", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-small", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-small.en", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-medium", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-medium.en", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-large", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-large-v2", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-large-v3", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "openai/whisper-large-v3-turbo", + "status": 1, + "verified_date": "2026-10-22", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "distil-whisper/distil-small.en", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "distil-whisper/distil-medium.en", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "distil-whisper/distil-large-v2", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "distil-whisper/distil-large-v3", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "WhisperForConditionalGeneration", + "model_id": "distil-whisper/distil-large-v3.5", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null } ] } diff --git a/transformer_lens/tools/model_registry/data/verification_history.json b/transformer_lens/tools/model_registry/data/verification_history.json index 8d0c82b8ce..fbd3926708 100644 --- a/transformer_lens/tools/model_registry/data/verification_history.json +++ b/transformer_lens/tools/model_registry/data/verification_history.json @@ -1,5 +1,5 @@ { - "last_updated": "2026-10-10T23:22:28.618265", + "last_updated": "2026-10-22T23:49:22.714765", "records": [ { "model_id": "Macropodus/macbert4mdcspell_v1", @@ -41448,6 +41448,30 @@ "notes": "Full verification completed", "invalidated": false, "invalidation_reason": null + }, + { + "model_id": "openai/whisper-tiny", + "architecture_id": "WhisperForConditionalGeneration", + "verified_date": "2026-10-22", + "verified_by": "verify_models", + "transformerlens_version": "0.0.0", + "prompt_profile": null, + "p4_scoring_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "openai/whisper-large-v3-turbo", + "architecture_id": "WhisperForConditionalGeneration", + "verified_date": "2026-10-22", + "verified_by": "verify_models", + "transformerlens_version": "0.0.0", + "prompt_profile": null, + "p4_scoring_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null } ] } diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index fa25690f63..22b72c8c14 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -137,6 +137,7 @@ "DeiTForImageClassification": "DeiT with a single CLS-token classification head", "Wav2Vec2Model": "Facebook's Wav2Vec 2.0 for speech", "HubertModel": "Facebook's HuBERT for speech", + "WhisperForConditionalGeneration": "OpenAI's Whisper speech recognition encoder-decoder (log-mel audio encoder, text decoder)", "Cohere2ForCausalLM": "Cohere's Command-A architecture with interleaved sliding-window RoPE and full-attention NoPE layers", "OuroForCausalLM": "ByteDance's Ouro looped language model (LoopLM) with weight-shared iterated depth", "RavenForCausalLM": "tomg-group-umd's Huginn depth-recurrent decoder (prelude / weight-tied recurrent core / coda) with runtime recurrence count", diff --git a/transformer_lens/utilities/architectures.py b/transformer_lens/utilities/architectures.py index 3f23cfa963..5b1f0c2bb3 100644 --- a/transformer_lens/utilities/architectures.py +++ b/transformer_lens/utilities/architectures.py @@ -73,6 +73,12 @@ "MusicFlamingoForConditionalGeneration", } +# Audio encoder + text decoder (Whisper); load via AutoModelForSpeechSeq2Seq and +# take log-mel input_features plus decoder token ids. +AUDIO_SEQ2SEQ_ARCHITECTURES: set[str] = { + "WhisperForConditionalGeneration", +} + # Audio spectrogram models for classification AUDIO_CLASSIFICATION_ARCHITECTURES: set[str] = { "ASTForAudioClassification", @@ -115,7 +121,7 @@ def classify_architecture(architecture: str) -> str: Returns one of: "seq2seq", "masked_lm", "multimodal", "audio", "vision", "causal_lm" """ - if architecture in SEQ2SEQ_ARCHITECTURES: + if architecture in SEQ2SEQ_ARCHITECTURES or architecture in AUDIO_SEQ2SEQ_ARCHITECTURES: return "seq2seq" if architecture in MASKED_LM_ARCHITECTURES: return "masked_lm"