Views
No views yet
1import torch
2from torch import nn
3from transformers import AutoModel
4
5class OptimizedCrossAttentionModel(nn.Module):
6 def __init__(self, n_labels):
7 super().__init__()
8 # 1. 사전 학습된 모델 로드 (WavLM, BERT)
9 self.audio_model = AutoModel.from_pretrained("microsoft/wavlm-base")
10 self.text_model = AutoModel.from_pretrained("klue/bert-base")
11
12 audio_hidden = 768
13 text_hidden = 768
14
15 # 2. V6: 각 도메인 특징 정규화를 위한 LayerNorm
16 self.audio_norm = nn.LayerNorm(audio_hidden)
17 self.text_norm = nn.LayerNorm(text_hidden)
18
19 # 3. 텍스트 특징을 오디오 공간으로 투영 (Projection)
20 self.text_proj = nn.Linear(text_hidden, audio_hidden)
21
22 # 4. 크로스 어텐션 레이어 (12개 헤드)
23 self.cross_attn = nn.MultiheadAttention(
24 embed_dim=audio_hidden,
25 num_heads=12,
26 batch_first=True,
27 dropout=0.1
28 )
29
30 # 5. V6: 최종 분류기 (Audio Mean + Audio Max + Text [CLS] = 2304차원)
31 self.classifier = nn.Sequential(
32 nn.Linear((audio_hidden * 2) + text_hidden, 512),
33 nn.LayerNorm(512),
34 nn.ReLU(),
35 nn.Dropout(0.4),
36 nn.Linear(512, 256),
37 nn.ReLU(),
38 nn.Linear(256, n_labels)
39 )
40
41 def forward(self, input_values, audio_mask, input_ids, text_mask, labels=None):
42 audio_feat = self.audio_model(input_values, attention_mask=audio_mask).last_hidden_state
43 text_feat = self.text_model(input_ids, attention_mask=text_mask).last_hidden_state
44
45 audio_feat = self.audio_norm(audio_feat)
46 text_feat = self.text_norm(text_feat)
47
48 text_feat_proj = self.text_proj(text_feat)
49 attn_output, _ = self.cross_attn(
50 query=audio_feat,
51 key=text_feat_proj,
52 value=text_feat_proj,
53 key_padding_mask=(text_mask == 0)
54 )
55
56 audio_mean = attn_output.mean(dim=1)
57 audio_max, _ = attn_output.max(dim=1)
58 audio_vec = torch.cat([audio_mean, audio_max], dim=1)
59
60 text_vec = text_feat[:, 0, :]
61
62 logits = self.classifier(torch.cat([audio_vec, text_vec], dim=1))
63
64 return {"logits": logits}
65
66
67## python 버전 3.9
68## pip install torch torchvision torchaudio transformers librosa safetensors huggingface_hub accelerate