Create modeling_multi_task_unixcoder.py

modeling_multi_task_unixcoder.py

@@ -0,0 +1,34 @@
+from transformers import PreTrainedModel, AutoModel, PretrainedConfig
+import torch
+from torch import nn
+
+class MultiTaskUnixCoderConfig(PretrainedConfig):
+    model_type = "multi_task_unixcoder"
+    def __init__(self, num_cwe_classes=106, **kwargs):
+        super().__init__(**kwargs)
+        self.num_cwe_classes = num_cwe_classes
+
+class MultiTaskUnixCoder(PreTrainedModel):
+    config_class = MultiTaskUnixCoderConfig
+    base_model_prefix = "base"
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.base = AutoModel.from_pretrained("microsoft/unixcoder-base")
+        self.vul_head = nn.Linear(768, 2)
+        self.cwe_head = nn.Linear(768, config.num_cwe_classes + 1)
+    
+    def forward(self, input_ids, attention_mask=None, labels_vul=None, labels_cwe=None):
+        outputs = self.base(input_ids=input_ids, attention_mask=attention_mask)
+        hidden_state = outputs.last_hidden_state[:, 0, :]
+        
+        vul_logits = self.vul_head(hidden_state)
+        cwe_logits = self.cwe_head(hidden_state)
+        
+        loss = None
+        if labels_vul is not None and labels_cwe is not None:
+            vul_loss = nn.CrossEntropyLoss()(vul_logits, labels_vul)
+            cwe_loss = nn.CrossEntropyLoss()(cwe_logits, labels_cwe + 1)
+            loss = vul_loss + cwe_loss
+        
+        return {"loss": loss, "vul_logits": vul_logits, "cwe_logits": cwe_logits} if loss is not None else {"vul_logits": vul_logits, "cwe_logits": cwe_logits}