Update tasks/text.py

tasks/text.py

@@ -7,6 +7,12 @@ import random
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
+from transformers import AutoTokenizer,BertForSequenceClassification,AutoModelForSequenceClassification,Trainer, TrainingArguments,DataCollatorWithPadding
+from datasets import Dataset
+import torch
+import numpy as np
+
+
 router = APIRouter()
 
 DESCRIPTION = "modernBERT"
@@ -57,12 +63,13 @@ async def evaluate_text(request: TextEvaluationRequest):
     #--------------------------------------------------------------------------------------------   
     
     # Make random predictions (placeholder for actual model inference)
-    # true_labels = test_dataset["label"]
+    true_labels = test_dataset["label"]
     # predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
     path_model = 'MatthiasPi/CARDS_ModernBert_no_overfitting'
     path_tokenizer = "answerdotai/ModernBERT-base"
     
     model = AutoModelForSequenceClassification.from_pretrained(path_model)
+    tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
     
     def preprocess_function(df):
         return tokenizer(df["quote"], truncation=True)
@@ -77,7 +84,8 @@ async def evaluate_text(request: TextEvaluationRequest):
         tokenizer=tokenizer
     )
     
-    predictions = trainer.predict(tokenized_test)
+    preds = trainer.predict(tokenized_test)
+    predictions = np.array([np.argmax(x) for x in preds[0]])
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE