Create model.py

src/model.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+from transformers import Wav2Vec2Model
+
+class Wav2Vec2_LSTM_MultiTask(nn.Module):
+    def __init__(self, num_emotions):
+        super().__init__()
+
+        self.wav2vec = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base")
+
+        self.lstm = nn.LSTM(
+            input_size=768,
+            hidden_size=256,
+            num_layers=2,
+            batch_first=True,
+            bidirectional=True
+        )
+
+        self.shared_fc = nn.Linear(512, 256)
+        self.emotion_head = nn.Linear(256, num_emotions)
+        self.stress_head = nn.Linear(256, 1)
+
+    def forward(self, input_values):
+        outputs = self.wav2vec(input_values)
+        x = outputs.last_hidden_state
+
+        lstm_out, _ = self.lstm(x)
+        pooled = torch.mean(lstm_out, dim=1)
+
+        shared = torch.relu(self.shared_fc(pooled))
+        return self.emotion_head(shared), self.stress_head(shared)