upload py

modeling_multitask.py

@@ -0,0 +1,79 @@
+from transformers import PretrainedConfig, AutoConfig
+
+class MultiTaskConfig(PretrainedConfig):
+    model_type = "barthez-multitask"
+
+    def __init__(
+        self,
+        base_model_name="moussaKam/barthez",
+        num_labels_type=3,
+        num_labels_priorite=3,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        self.base_model_name = base_model_name
+        self.num_labels_type = num_labels_type
+        self.num_labels_priorite = num_labels_priorite
+
+        # on charge la config du modèle de base
+        self.base_config = AutoConfig.from_pretrained(base_model_name)
+
+import torch
+import torch.nn as nn
+from transformers import BartPreTrainedModel, BartModel
+from dataclasses import dataclass
+from transformers.modeling_outputs import ModelOutput
+from typing import Optional
+
+
+@dataclass
+class MultiTaskOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits_type: torch.FloatTensor = None
+    logits_priorite: torch.FloatTensor = None
+
+
+class MultiTaskModel(BartPreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.model = BartModel(config)
+        hidden_size = config.d_model
+
+        self.classifier_type = nn.Linear(hidden_size, config.num_labels_type)
+        self.classifier_priorite = nn.Linear(hidden_size, config.num_labels_priorite)
+
+        self.loss_fct = nn.CrossEntropyLoss()
+
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels_type=None,
+        labels_priorite=None,
+    ):
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+        )
+
+        pooled_output = outputs.last_hidden_state.mean(dim=1)
+
+        logits_type = self.classifier_type(pooled_output)
+        logits_priorite = self.classifier_priorite(pooled_output)
+
+        loss = None
+        if labels_type is not None and labels_priorite is not None:
+            loss_type = self.loss_fct(logits_type, labels_type)
+            loss_priorite = self.loss_fct(logits_priorite, labels_priorite)
+            loss = loss_type + loss_priorite
+
+        return MultiTaskOutput(
+            loss=loss,
+            logits_type=logits_type,
+            logits_priorite=logits_priorite,
+        )