Update tasks/text.py

tasks/text.py

@@ -1,10 +1,12 @@
+import torch
+import random
 from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
-import random
-from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer
-import torch
+from transformers import AutoTokenizer, AutoModel, AutoConfig
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
@@ -14,6 +16,23 @@ router = APIRouter()
 DESCRIPTION = "GTE Architecture"
 ROUTE = "/text"
 
+class AutoBertClassifier(nn.Module):
+    def __init__(self, num_labels=8, model_path="Alibaba-NLP/gte-modernbert-base"):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.bert = AutoModel.from_pretrained(model_path)
+        self.config = AutoConfig.from_pretrained(model_path)
+        self.config.num_labels = num_labels
+        self.dropout = nn.Dropout(0.05)
+        self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.last_hidden_state[:, 0]  # Using [CLS] token representation
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+        return logits
+
 @router.post(ROUTE, tags=["Text Task"], 
              description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):
@@ -49,13 +68,14 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     true_labels = test_dataset["label"]
     texts = test_dataset["quote"]
-
+    
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    
     model_repo = "elucidator8918/frugal-ai-text"
-    config = AutoConfig.from_pretrained(model_repo)
-    model = AutoModelForSequenceClassification.from_pretrained(model_repo)
+    model = AutoBertClassifier(num_labels=8)
+    model.load_state_dict(load_file(hf_hub_download(repo_id=model_repo, filename="model.safetensors")))
     tokenizer = AutoTokenizer.from_pretrained(model_repo)
-
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    
     model = model.to(device)
     model.eval()