Speech-Language Integration › Code Implementation: Building a Speech-LLM Interface

import torch

# Stub encoder and LLM objects to keep this example self-contained.
# In practice, load real models with:
#   from transformers import WhisperModel, GPT2Model, GPT2Tokenizer
#   whisper_encoder = WhisperModel.from_pretrained("openai/whisper-base").encoder
#   gpt2_model = GPT2Model.from_pretrained("gpt2")
#   gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

whisper_encoder = type(
    "obj",
    (object,),
    {
        "config": type("cfg", (object,), {"d_model": 512})(),
        "parameters": lambda self: [],
        "__call__": lambda self, x: type(
            "out", (object,), {"last_hidden_state": torch.randn(1, 1500, 512)}
        )(),
    },
)()

gpt2_model = type(
    "obj",
    (object,),
    {
        "config": type("cfg", (object,), {"n_embd": 768})(),
        "parameters": lambda self: [],
        "transformer": type(
            "t",
            (object,),
            {"wte": lambda self, x: torch.randn(1, x.shape[1], 768)},
        )(),
    },
)()

gpt2_tokenizer = type(
    "obj",
    (object,),
    {
        "encode": lambda self, text, return_tensors=None: torch.tensor(
            [[101, 102, 103]]
        ),
    },
)()

# Freeze base models so only the connector trains
for param in whisper_encoder.parameters():
    param.requires_grad = False
for param in gpt2_model.parameters():
    param.requires_grad = False

import torch.nn as nn


class SpeechLLMConnector(nn.Module):
    """
    Projects Whisper encoder outputs into GPT-2's embedding space.
    A linear projection maps audio_dim -> text_dim, and learned
    positional embeddings preserve temporal ordering information.
    """

    def __init__(self, audio_dim=512, text_dim=768, max_audio_len=1500):
        super().__init__()
        self.projection = nn.Linear(audio_dim, text_dim)
        self.layer_norm = nn.LayerNorm(text_dim)
        # Learned positional embeddings for audio tokens
        self.audio_pos_embed = nn.Parameter(
            torch.randn(1, max_audio_len, text_dim) * 0.02
        )

    def forward(self, audio_features):
        """
        Args:
            audio_features: [batch, seq_len, audio_dim] from Whisper encoder
        Returns:
            projected: [batch, seq_len, text_dim] ready for GPT-2
        """
        projected = self.projection(audio_features)
        seq_len = audio_features.size(1)
        projected = projected + self.audio_pos_embed[:, :seq_len, :]
        projected = self.layer_norm(projected)
        return projected


connector = SpeechLLMConnector(
    audio_dim=whisper_encoder.config.d_model,  # 512 for Whisper base
    text_dim=gpt2_model.config.n_embd,  # 768 for GPT-2
)

print(
    f"Connector parameters: {sum(p.numel() for p in connector.parameters()):,}"
)

import torch

# Simulated mel-spectrogram input (batch=1, n_mels=80, seq_len=3000)
# This represents approximately 30 seconds of audio at Whisper's frame rate.
batch_size = 1
n_mels = 80
sequence_length = 3000

dummy_audio_features = torch.randn(batch_size, n_mels, sequence_length)

# Step 1: Encode audio with Whisper
with torch.no_grad():
    encoder_outputs = whisper_encoder(dummy_audio_features)
    audio_hidden_states = encoder_outputs.last_hidden_state  # [1, 1500, 512]

# Step 2: Project to GPT-2 embedding space
audio_embeddings = connector(audio_hidden_states)  # [1, 1500, 768]

# Step 3: Prepare a text prefix to prepend before the audio context
prefix_text = "The user said:"
prefix_tokens = gpt2_tokenizer.encode(prefix_text, return_tensors="pt")
prefix_embeds = gpt2_model.transformer.wte(prefix_tokens)  # [1, N_tokens, 768]

# Step 4: Concatenate audio and text embeddings into a unified sequence
combined_embeds = torch.cat([audio_embeddings, prefix_embeds], dim=1)

encoder_shape = audio_hidden_states.shape
projected_shape = audio_embeddings.shape
prefix_shape = prefix_embeds.shape
combined_length = combined_embeds.shape[1]

class AlignmentLoss(nn.Module):
    """
    Training objective: Given audio features, predict the text transcript.
    This is essentially an ASR objective expressed through the LLM's vocabulary,
    teaching the connector to translate acoustic content into linguistic tokens.
    """

    def __init__(self):
        super().__init__()
        self.ce_loss = nn.CrossEntropyLoss()

    def forward(self, model_outputs, target_tokens):
        """
        model_outputs: [batch, seq_len, vocab_size]
        target_tokens: [batch, target_len]
        """
        # Shift logits and labels for next-token prediction
        shift_logits = model_outputs[..., :-1, :].contiguous()
        shift_labels = target_tokens[..., 1:].contiguous()

        loss = self.ce_loss(
            shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
        )
        return loss


alignment_criterion = AlignmentLoss()

# Simulate a target transcript
transcript = "What is the weather like today?"
target_tokens = gpt2_tokenizer.encode(transcript, return_tensors="pt")
target_ids_sample = target_tokens[0][:8]
target_shape = target_tokens.shape