File size: 6,319 Bytes
a00ab1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | 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)) |