Package model correctly

README.md

@@ -33,15 +33,14 @@ For more detailed code regarding generating the annotations in Toxic Commons, tr
 
 # How to Use
 
-```
-from transformers import AutoTokenizer
-from celadon.model import MultiHeadDebertaForSequenceClassification
+```py
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-tokenizer = AutoTokenizer.from_pretrained("celadon")
-model = MultiHeadDebertaForSequenceClassification.from_pretrained("celadon")
+model = AutoModelForSequenceClassification.from_pretrained("PleIAs/celadon", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("PleIAs/celadon", trust_remote_code=True)
 model.eval()
 
-sample_text = "This is an example of a normal sentence"
+sample_text = "A very gender inappropriate comment"
 
 inputs = tokenizer(sample_text, return_tensors="pt", padding=True, truncation=True)
 outputs = model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
@@ -53,6 +52,22 @@ predictions = outputs.argmax(dim=-1).squeeze().tolist()
 print(f"Text: {sample_text}")
 for i, category in enumerate(categories):
     print(f"Prediction for Category {category}: {predictions[i]}")
+# Text: A very gender inappropriate comment
+# Prediction for Category Race/Origin: 0
+# Prediction for Category Gender/Sex: 3
+# Prediction for Category Religion: 0
+# Prediction for Category Ability: 0
+# Prediction for Category Violence: 0
+```
+
+you can also use transformers pipelines to get a more streamlined experience
+
+```py
+from transformers import pipeline
+pipe = pipeline("text-classification", model="PleIAs/celadon", trust_remote_code=True)
+result = pipe("This is an example of a normal sentence")
+print(result)
+# [{'Race/Origin': 0, 'Gender/Sex': 3, 'Religion': 0, 'Ability': 0, 'Violence': 0}]
 ```
 
 # How to Cite

config.json

@@ -1,8 +1,22 @@
 {
+  "_name_or_path": "./celadon",
   "architectures": [
-    "MultiHeadDebertaForSequenceClassification"
+    "MultiHeadDebertaForSequenceClassificationModel"
   ],
   "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoConfig": "configuration_deberta_multi.MultiHeadDebertaV2Config",
+    "AutoModelForSequenceClassification": "modelling_deberta_multi.MultiHeadDebertaForSequenceClassificationModel"
+  },
+  "custom_pipelines": {
+    "text-classification": {
+      "impl": "custom_pipeline.CustomTextClassificationPipeline",
+      "pt": [
+        "AutoModelForSequenceClassification"
+      ],
+      "tf": []
+    }
+  },
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,
   "hidden_size": 768,
@@ -11,9 +25,10 @@
   "layer_norm_eps": 1e-07,
   "max_position_embeddings": 512,
   "max_relative_positions": -1,
-  "model_type": "deberta-v2",
+  "model_type": "multi-head-deberta-for-sequence-classification",
   "norm_rel_ebd": "layer_norm",
   "num_attention_heads": 12,
+  "num_heads": 5,
   "num_hidden_layers": 6,
   "pad_token_id": 0,
   "pooler_dropout": 0,
@@ -28,7 +43,7 @@
   "relative_attention": true,
   "share_att_key": true,
   "torch_dtype": "float32",
-  "transformers_version": "4.45.2",
+  "transformers_version": "4.46.2",
   "type_vocab_size": 0,
   "vocab_size": 128100
 }

configuration_deberta_multi.py

@@ -0,0 +1,7 @@
+from transformers import DebertaV2Config
+
+class MultiHeadDebertaV2Config(DebertaV2Config):
+    model_type = "multi-head-deberta-for-sequence-classification"
+    def __init__(self, num_heads=5, **kwargs):
+        self.num_heads = num_heads
+        super().__init__(**kwargs)

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

special_tokens_map.json

@@ -1,10 +1,46 @@
 {
-  "bos_token": "[CLS]",
-  "cls_token": "[CLS]",
-  "eos_token": "[SEP]",
-  "mask_token": "[MASK]",
-  "pad_token": "[PAD]",
-  "sep_token": "[SEP]",
+  "bos_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
   "unk_token": {
     "content": "[UNK]",
     "lstrip": false,

tokenizer_config.json

@@ -47,12 +47,19 @@
   "do_lower_case": false,
   "eos_token": "[SEP]",
   "mask_token": "[MASK]",
-  "model_max_length": 1000000000000000019884624838656,
+  "max_length": 512,
+  "model_max_length": 512,
+  "pad_to_multiple_of": null,
   "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
   "sep_token": "[SEP]",
   "sp_model_kwargs": {},
   "split_by_punct": false,
+  "stride": 0,
   "tokenizer_class": "DebertaV2Tokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
   "unk_token": "[UNK]",
   "vocab_type": "spm"
 }