Update model.py

model.py

@@ -2,12 +2,12 @@ import torch.nn as nn
 from transformers import AutoModel
 
 class MultitaskModel(nn.Module):
-    def __init__(self):
+    def __init__(self, model_name, num_sentiment_labels, num_emotion_labels):
         super().__init__()
-        self.encoder = AutoModel.from_pretrained("distilbert-base-uncased")
+        self.encoder = AutoModel.from_pretrained(model_name)
         self.dropout = nn.Dropout(0.1)
-        self.classifier_sent = nn.Linear(768, 2)  # num_sentiment_labels
-        self.classifier_emo = nn.Linear(768, 7)   # num_emotion_labels
+        self.classifier_sent = nn.Linear(self.encoder.config.hidden_size, num_sentiment_labels)
+        self.classifier_emo = nn.Linear(self.encoder.config.hidden_size, num_emotion_labels)
 
     def forward(self, input_ids=None, attention_mask=None):
         outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
@@ -15,4 +15,4 @@ class MultitaskModel(nn.Module):
         pooled = self.dropout(pooled)
         sent = self.classifier_sent(pooled)
         emo = self.classifier_emo(pooled)
-        return sent, emo
+        return sent, emo