| import torch |
| import torch.nn as nn |
| from transformers import WavLMModel, DistilBertModel, DistilBertTokenizer |
|
|
|
|
| class MultimodalFusionModel(nn.Module): |
| """Frozen WavLM (audio) + DistilBERT (text) encoders with a |
| BiLSTM + attention fusion network and multi-label classification heads. |
| |
| The encoders are only used to compute turn-level embeddings |
| (`compute_embeddings` / `get_embeddings`). The fusion network |
| operates on precomputed `(batch, num_turns, emb_dim)` tensors. |
| """ |
|
|
| def __init__(self, |
| audio_emb_dim=768, |
| text_emb_dim=768, |
| fusion_hidden=256, |
| lstm_layers=2, |
| num_heads=4, |
| num_classes=4, |
| freeze_encoders=True): |
| super().__init__() |
|
|
| self.freeze_encoders = freeze_encoders |
| self.audio_emb_dim = audio_emb_dim |
| self.text_emb_dim = text_emb_dim |
| self.fusion_hidden = fusion_hidden |
| self.num_classes = num_classes |
|
|
| self.audio_encoder = WavLMModel.from_pretrained("microsoft/wavlm-base") |
| self.text_encoder = DistilBertModel.from_pretrained("distilbert-base-uncased") |
| self.tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased") |
|
|
| self.audio_encoder.eval() |
| self.text_encoder.eval() |
| if freeze_encoders: |
| for encoder in (self.audio_encoder, self.text_encoder): |
| for p in encoder.parameters(): |
| p.requires_grad = False |
|
|
| self.audio_proj = nn.Linear(audio_emb_dim, fusion_hidden) |
| self.text_proj = nn.Linear(text_emb_dim, fusion_hidden) |
|
|
| self.bilstm = nn.LSTM(fusion_hidden * 2, |
| fusion_hidden, |
| num_layers=lstm_layers, |
| bidirectional=True, |
| batch_first=True, |
| dropout=0.2 if lstm_layers > 1 else 0) |
|
|
| self.attention = nn.MultiheadAttention(embed_dim=fusion_hidden * 2, |
| num_heads=num_heads, |
| batch_first=True) |
|
|
| self.classifiers = nn.ModuleList([ |
| nn.Linear(fusion_hidden * 2, 1) for _ in range(num_classes) |
| ]) |
|
|
| def forward(self, audio_embeds=None, text_embeds=None, lengths=None): |
| """Fuse precomputed turn embeddings. |
| |
| audio_embeds: (batch, num_turns, audio_emb_dim) |
| text_embeds: (batch, num_turns, text_emb_dim) |
| lengths: (batch,) ints, number of real turns (optional, used to |
| mask padding when mean-pooling). |
| Returns logits: (batch, num_classes) |
| """ |
| if audio_embeds is None or text_embeds is None: |
| raise ValueError("forward() expects precomputed embeddings " |
| "(audio_embeds, text_embeds).") |
|
|
| seq_len = audio_embeds.shape[1] |
| if lengths is None: |
| lengths = torch.full((audio_embeds.shape[0],), |
| seq_len, |
| dtype=torch.long, |
| device=audio_embeds.device) |
|
|
| audio_proj = self.audio_proj(audio_embeds) |
| text_proj = self.text_proj(text_embeds) |
| combined = torch.cat([audio_proj, text_proj], dim=-1) |
|
|
| lstm_out, _ = self.bilstm(combined) |
|
|
| attn_out, _ = self.attention(lstm_out, lstm_out, lstm_out) |
|
|
| mask = torch.arange(seq_len, device=attn_out.device).unsqueeze(0) < lengths.unsqueeze(1) |
| mask = mask.unsqueeze(-1).float() |
| pooled = (attn_out * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0) |
|
|
| logits = torch.cat([cls(pooled) for cls in self.classifiers], dim=1) |
| return logits |
|
|
| def get_embeddings(self, audio_waveforms, input_ids, attention_mask=None): |
| """Compute per-turn audio and text embeddings from raw inputs. |
| |
| audio_waveforms: (batch, num_samples) float waveform (or 1/2 channel). |
| input_ids: (batch, seq_len) token ids. |
| attention_mask: (batch, seq_len) optional. |
| Returns (audio_emb, text_emb): |
| audio_emb: (batch, audio_emb_dim) mean-pooled over time |
| text_emb: (batch, text_emb_dim) [CLS] embedding |
| """ |
| with torch.no_grad(): |
| audio_dev = next(self.audio_encoder.parameters()).device |
| text_dev = next(self.text_encoder.parameters()).device |
|
|
| audio_out = self.audio_encoder(audio_waveforms.to(audio_dev), |
| output_hidden_states=True) |
| audio_feat = audio_out.last_hidden_state |
| if audio_feat.dim() == 4: |
| audio_feat = audio_feat.squeeze(1) |
| audio_emb = audio_feat.mean(dim=1) |
|
|
| out = self.text_encoder(input_ids=input_ids.to(text_dev), |
| attention_mask=attention_mask.to(text_dev)) |
| text_emb = out.last_hidden_state[:, 0, :] |
| return audio_emb, text_emb |
|
|
| def compute_embeddings(self, waveforms, texts, max_text_len=128, device=None): |
| """Batch helper: raw waveforms + raw strings -> per-turn embeddings. |
| |
| waveforms: list/tensor of (num_samples,) waveforms. |
| texts: list of strings. |
| Returns (audio_embeds, text_embeds): (batch, emb_dim) tensors. |
| """ |
| self.eval() |
| enc = self.tokenizer(texts, truncation=True, padding=True, |
| max_length=max_text_len, return_tensors="pt") |
| input_ids = enc["input_ids"] |
| attention_mask = enc["attention_mask"] |
| device = device or next(self.parameters()).device |
|
|
| if isinstance(waveforms, (list, tuple)): |
| max_len = max(w.shape[0] for w in waveforms) if waveforms else 0 |
| if waveforms and waveforms[0].dim() == 1: |
| wav = torch.zeros(len(waveforms), max_len) |
| else: |
| wav = torch.cat([w.unsqueeze(0) for w in waveforms], dim=0) |
| for i, w in enumerate(waveforms): |
| wav[i, :w.shape[-1]] = w |
| else: |
| wav = waveforms |
|
|
| return self.get_embeddings(wav.to(device), |
| input_ids.to(device), |
| attention_mask.to(device)) |
