transformer_lens.model_bridge.supported_architectures.whisper module

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.

class transformer_lens.model_bridge.supported_architectures.whisper.WhisperArchitectureAdapter(cfg: Any)

Bases: 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.

Transcription: bridge.generate(input_features, return_type="str") decodes from the same prefix (or decoder_input), applying the checkpoint’s suppress_tokens and begin_suppress_tokens as HF’s generate does.

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.

__init__(cfg: Any) → None

Set the audio seq2seq flags on top of the BART-family setup.

applicable_phases: list[int] = [1]
audio_input_name: str = 'input_features'
component_mapping: ComponentMapping | None
force_scale_embedding: bool = False
has_final_stack_norm: bool = True
n_layers_from: str = 'decoder'
require_symmetric_layers: bool = False
residual_layout() → ResidualLayout

Norm readers, residual writers and value/output pairs of both stacks.

Each block’s norms feed its own sublayers. The encoder’s final norm feeds the cross-attention keys and values of every decoder block, and the decoder’s final norm feeds the unembedding. Encoder writes are the attention and MLP outputs plus the sinusoidal positions; the GELU’d conv output is not linear in any weight, so it stays uncentered (centering is exact either way, since every read is through a LayerNorm).

unembed_name: str = 'proj_out'
uses_split_attention: bool
weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None