Add multitask_model.py

multitask_model.py

@@ -0,0 +1,122 @@
+import torch
+import torch.nn as nn
+from transformers import BertModel, BertConfig, RobertaTokenizer, RobertaModel, RobertaConfig, PretrainedConfig
+from transformers.modeling_outputs import SequenceClassifierOutput
+from peft import LoraConfig, get_peft_model, LoraModel
+
+task_name_to_id = {"sentiment": 0, "hate": 1, "emotion": 2}
+
+# Number of classes for each task
+task_num_labels = {
+    "sentiment": 3,
+    "hate": 3,
+    "emotion": 4
+}
+
+
+
+class MultiTaskModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        model_name="roberta-base"
+        config = RobertaConfig.from_pretrained(model_name)
+        base_model  = RobertaModel.from_pretrained(model_name, config=config)
+
+        # self.task_weights = {
+        #     "sentiment": 1, 
+        #     "hate": 2,      
+        #     "emotion": 1       
+        # }
+
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            target_modules=['query', 'value'],
+            lora_dropout=0.05,
+            bias='none',
+            task_type='SEQ_CLS'
+        )
+
+        self.model = LoraModel(base_model, lora_config, adapter_name="shared")
+
+        # self.model.add_adapter(lora_config, adapter_name= 'hate')    # For hate
+        # self.model.set_adapter("shared")  # Set default
+
+
+        hidden_size = config.hidden_size
+        dropout_prob = 0.1
+        intermediate_size = 128 
+
+        self.sentiment_head = nn.Sequential(
+            nn.Linear(hidden_size, intermediate_size),
+            nn.ReLU(),
+            nn.Dropout(dropout_prob),
+            nn.Linear(intermediate_size, task_num_labels["sentiment"])
+        )
+
+        self.hate_head = nn.Sequential(
+            nn.Linear(hidden_size, intermediate_size),
+            nn.ReLU(),
+            nn.Dropout(dropout_prob),
+            nn.Linear(intermediate_size, task_num_labels["hate"])
+        )
+
+        self.emotion_head = nn.Sequential(
+            nn.Linear(hidden_size, intermediate_size),
+            nn.ReLU(),
+            nn.Dropout(dropout_prob),
+            nn.Linear(intermediate_size, task_num_labels["emotion"])
+        )
+        #self.bert.print_trainable_parameters()r
+        print(f"Trainable parameters (LoRA): {sum(p.numel() for p in self.model.parameters() if p.requires_grad)}")
+        print(f"Total parameters: {sum(p.numel() for p in self.parameters() if p.requires_grad)}")
+        self.loss_fct = nn.CrossEntropyLoss()
+
+    def forward(self, input_ids=None, attention_mask=None, task_id=None, labels=None):
+
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+        pooled = outputs.last_hidden_state[:, 0]  # Use first token (CLS)
+
+        sentiment_mask = task_id == task_name_to_id["sentiment"]
+        hate_mask = task_id == task_name_to_id["hate"]
+        emotion_mask = task_id == task_name_to_id["emotion"]
+
+        logits = {}
+        loss = 0
+
+    # Sentiment task
+        if sentiment_mask.any():
+            sentiment_pooled = pooled[sentiment_mask]
+            sentiment_logits = self.sentiment_head(sentiment_pooled)
+            logits["sentiment"] = sentiment_logits
+            if labels is not None:
+                sentiment_labels = labels[sentiment_mask]
+                loss += self.loss_fct(sentiment_logits, sentiment_labels)
+        else:
+            logits["sentiment"] = torch.empty(0, task_num_labels["sentiment"], device=input_ids.device)
+
+        # Hate task
+        if hate_mask.any():
+            hate_pooled = pooled[hate_mask]
+            hate_logits = self.hate_head(hate_pooled)
+            logits["hate"] = hate_logits
+            if labels is not None:
+                hate_labels = labels[hate_mask]
+                loss += self.loss_fct(hate_logits, hate_labels)
+        else:
+            logits["hate"] = torch.empty(0, task_num_labels["hate"], device=input_ids.device)
+
+        # Emotion task
+        if emotion_mask.any():
+            emotion_pooled = pooled[emotion_mask]
+            emotion_logits = self.emotion_head(emotion_pooled)
+            logits["emotion"] = emotion_logits
+            if labels is not None:
+                emotion_labels = labels[emotion_mask]
+                loss += self.loss_fct(emotion_logits, emotion_labels)
+        else:
+            logits["emotion"] = torch.empty(0, task_num_labels["emotion"], device=input_ids.device)
+
+        return {"loss": loss, "logits": logits} if labels is not None else {"logits": logits}
+
+