proper sbert model load

tasks/text.py

@@ -1,18 +1,21 @@
-from fastapi import APIRouter
 from datetime import datetime
-from datasets import load_dataset
-from sklearn.metrics import accuracy_score
 import random
-from skops.io import load
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
+
+import numpy as np
 import torch
+from torch import nn
 from torch.utils.data import DataLoader, Dataset
-import numpy as np
-from accelerate.test_utils.testing import get_backend
+from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer
+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
 
 router = APIRouter()
 
@@ -27,6 +30,26 @@ 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)
@@ -81,9 +104,15 @@ def bert_classifier(test_dataset: dict, model: str):
 
     model_repo = f"theterryzhang/frugal_ai_{model}"
 
-    model = AutoModelForSequenceClassification.from_pretrained(model_repo)
     tokenizer = AutoTokenizer.from_pretrained(model_repo)
 
+    if model.isin(['bert_base_pruned']):
+        model = AutoModelForSequenceClassification.from_pretrained(model_repo)
+    elif model.isin(['sbert_distilroberta']):
+        model = SentenceBERTMultiClass.from_pretrained(model_repo)
+    else:
+        raise(ValueError)
+
     # Use CUDA if available 
     device, _, _ = get_backend()