Create model.py

model.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+from transformers import Wav2Vec2Model, Wav2Vec2PreTrainedModel
+
+
+class LanguageIdentificationLayer(nn.Module):
+    def __init__(self, hidden_size, num_languages=3):
+        super().__init__()
+
+        self.lid_head = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(hidden_size, num_languages)
+        )
+
+    def forward(self, x):
+        return self.lid_head(x)
+
+
+class LanguageAwareEmotionHead(nn.Module):
+    def __init__(self, hidden_size, num_emotions=5, num_languages=3):
+        super().__init__()
+
+        self.lang_embeddings = nn.Embedding(num_languages, hidden_size)
+
+        self.pre_classifier = nn.Linear(hidden_size, hidden_size)
+
+        self.classifier = nn.Linear(hidden_size, num_emotions)
+
+        self.dropout = nn.Dropout(0.1)
+
+    def forward(self, features, language_logits):
+
+        language_ids = torch.argmax(language_logits, dim=-1)
+
+        lang_embed = self.lang_embeddings(language_ids)
+
+        features = features + lang_embed
+
+        features = torch.relu(self.pre_classifier(features))
+
+        features = self.dropout(features)
+
+        return self.classifier(features)
+
+
+class MMSForMultilingualSER(Wav2Vec2PreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.wav2vec2 = Wav2Vec2Model(config)
+
+        hidden_size = config.hidden_size
+
+        self.lid_layer = LanguageIdentificationLayer(hidden_size)
+
+        self.emotion_head = LanguageAwareEmotionHead(hidden_size)
+
+        self.dropout = nn.Dropout(0.1)
+
+        self.init_weights()
+
+    def forward(self, input_values, attention_mask=None):
+
+        outputs = self.wav2vec2(
+            input_values,
+            attention_mask=attention_mask
+        )
+
+        hidden_states = outputs.last_hidden_state
+
+        pooled = hidden_states.mean(dim=1)
+
+        pooled = self.dropout(pooled)
+
+        language_logits = self.lid_layer(pooled)
+
+        emotion_logits = self.emotion_head(pooled, language_logits)
+
+        return emotion_logits