add custom classifiers

tasks/custom_classifiers.py

@@ -0,0 +1,28 @@
+from transformers import RobertaModel, AutoTokenizer
+from transformers.modeling_outputs import SequenceClassifierOutput
+from huggingface_hub import PyTorchModelHubMixin
+from torch.nn import CrossEntropyLoss
+import torch.nn as nn
+import torch
+
+class SentenceBERTClassifier(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, model_name="sentence-transformers/all-distilroberta-v1", num_labels=8):
+        super().__init__()
+        self.sbert = RobertaModel.from_pretrained(model_name)
+        self.config = self.sbert.config
+        self.config.num_labels = num_labels
+        self.dropout = nn.Dropout(0.05)
+        self.config.classifier_dropout = 0.05
+        self.classifier = nn.Linear(self.config.hidden_size, self.config.num_labels)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.sbert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.pooler_output
+        dropout_output = self.dropout(pooled_output)
+        logits = self.classifier(dropout_output)
+
+        return SequenceClassifierOutput(
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

tasks/text.py

@@ -10,12 +10,12 @@ from fastapi import APIRouter
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
 from skops.io import load
-from huggingface_hub import PyTorchModelHubMixin
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 from .utils.text_preprocessor import preprocess
 from accelerate.test_utils.testing import get_backend
+from custom_classifiers import SentenceBERTClassifier
 
 router = APIRouter()
 
@@ -86,7 +86,12 @@ def bert_classifier(test_dataset: dict, model: str):
 
     tokenizer = AutoTokenizer.from_pretrained(model_repo)
 
-    model = AutoModelForSequenceClassification.from_pretrained(model_repo)
+    if model in ["bert_base_pruned"]:
+        model = AutoModelForSequenceClassification.from_pretrained(model_repo)
+    elif model in ["sbert_distilroberta"]:
+        model = SentenceBERTClassifier.from_pretrained(model_repo)
+    else:
+        raise(ValueError)
 
     # Use CUDA if available 
     device, _, _ = get_backend()