missing files and rename

custom_pipeline.py

@@ -0,0 +1,29 @@
+from transformers import TextClassificationPipeline, AutoTokenizer
+
+class CustomTextClassificationPipeline(TextClassificationPipeline):
+    def __init__(self, model, tokenizer=None, **kwargs):
+        # Initialize tokenizer first
+        if tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained(model.config._name_or_path)
+        # Make sure we store the tokenizer before calling super().__init__
+        self.tokenizer = tokenizer
+        super().__init__(model=model, tokenizer=tokenizer, **kwargs)
+        
+    def _sanitize_parameters(self, **kwargs):
+        preprocess_kwargs = {}
+        return preprocess_kwargs, {}, {}
+
+    def preprocess(self, inputs):
+        return self.tokenizer(inputs, return_tensors='pt', truncation=False)
+
+    def _forward(self, model_inputs):
+        input_ids = model_inputs['input_ids']
+        attention_mask = (input_ids != 0).long()
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+        return outputs
+
+    def postprocess(self, model_outputs):
+        predictions = model_outputs.logits.argmax(dim=-1).squeeze().tolist()
+        categories = ["Race/Origin", "Gender/Sex", "Religion", "Ability", "Violence", "Other"]
+        return dict(zip(categories, predictions))
+

model.py

@@ -1,20 +0,0 @@
-import torch
-import torch.nn as nn
-from transformers import DebertaV2Model, DebertaV2PreTrainedModel
-
-class MultiHeadDebertaForSequenceClassification(DebertaV2PreTrainedModel):
-    def __init__(self, config, num_heads=5):
-        super().__init__(config)
-        self.num_heads = num_heads
-        self.deberta = DebertaV2Model(config)  
-        self.heads = nn.ModuleList([nn.Linear(config.hidden_size, 4) for _ in range(num_heads)])  
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.post_init()
-
-    def forward(self, input_ids=None, attention_mask=None):
-        outputs = self.deberta(input_ids=input_ids, attention_mask=attention_mask)
-        sequence_output = outputs[0]  
-        logits_list = [head(self.dropout(sequence_output[:, 0, :])) for head in self.heads]
-        logits = torch.stack(logits_list, dim=1)
-        return logits
-

modelling_deberta_multi.py

@@ -0,0 +1,31 @@
+import torch
+from torch import nn, Tensor
+from typing import Optional
+from transformers import DebertaV2PreTrainedModel, DebertaV2Model
+from .configuration_deberta_multi import MultiHeadDebertaV2Config
+
+class MultiHeadDebertaForSequenceClassificationModel(DebertaV2PreTrainedModel):
+
+    config_class = MultiHeadDebertaV2Config
+    def __init__(self, config):  # type: ignore
+        super().__init__(config)
+        self.deberta = DebertaV2Model(config)
+        self.heads = nn.ModuleList(
+            [nn.Linear(config.hidden_size, 4) for _ in range(config.num_heads)]
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional["Tensor"] = None,
+        attention_mask: Optional["Tensor"] = None,
+    ) -> "Tensor":
+        outputs = self.deberta(input_ids=input_ids, attention_mask=attention_mask)
+        sequence_output = outputs[0]
+        logits_list = [
+            head(self.dropout(sequence_output[:, 0, :])) for head in self.heads
+        ]
+        logits = torch.stack(logits_list, dim=1)
+        outputs["logits"] = logits
+        return outputs