Update tasks/text.py

tasks/text.py

@@ -30,26 +30,6 @@ models_descriptions = {
     "sbert_distilroberta": "Fine-tuned sentence transformer DistilRoBERTa"
 }
 
-class SentenceBERTMultiClass(nn.Module, PyTorchModelHubMixin):
-    def __init__(self, model_name, num_labels=8):
-        super().__init__()
-        self.sbert = AutoModel.from_pretrained(model_name)
-        self.config = self.sbert.config
-        self.dropout = nn.Dropout(0.05)
-        self.classifier = nn.Linear(self.sbert.config.hidden_size, num_labels)
-
-    def forward(self, input_ids, attention_mask):
-        outputs = self.sbert(input_ids=input_ids, attention_mask=attention_mask)
-        if hasattr(outputs, "pooler_output"):
-            pooled_output = outputs.pooler_output
-        else:
-            pooled_output = outputs.last_hidden_state.mean(dim=1)
-
-        dropout_output = self.dropout(pooled_output)
-        logits = self.classifier(dropout_output)
-
-        return logits
-
 
 def baseline_model(dataset_length: int):
     # Make random predictions (placeholder for actual model inference)
@@ -106,12 +86,7 @@ def bert_classifier(test_dataset: dict, model: str):
 
     tokenizer = AutoTokenizer.from_pretrained(model_repo)
 
-    if model in ['bert_base_pruned']:
-        model = AutoModelForSequenceClassification.from_pretrained(model_repo)
-    elif model in ['sbert_distilroberta']:
-        model = SentenceBERTMultiClass.from_pretrained(model_repo)
-    else:
-        raise(ValueError)
+    model = AutoModelForSequenceClassification.from_pretrained(model_repo)
 
     # Use CUDA if available 
     device, _, _ = get_backend()