Update tasks/text.py

tasks/text.py

@@ -8,7 +8,7 @@ from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 #additional imports 
-from transformers import Trainer, TrainingArguments, DistilBertForSequenceClassification, DistilBertTokenizerFast
+from transformers import Trainer, TrainingArguments, DistilBertForSequenceClassification, DistilBertTokenizerFast, AutoModelForSequenceClassification,DataCollatorWithPadding
 import logging
 
 router = APIRouter()
@@ -62,11 +62,29 @@ async def evaluate_text(request: TextEvaluationRequest):
     
     # Make random predictions (placeholder for actual model inference)
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
 
-    print("hello_world")
+    
+    model_name = "seronk/distillbert-frugal-ai"
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    def preprocess_function(df):
+        return tokenizer(df["quote"], truncation=True)
+    tokenized_test = test_dataset.map(preprocess_function, batched=True)
+
+    training_args = torch.load("./tasks/utils/training_args.bin")
+    training_args.eval_strategy='no'
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        tokenizer=tokenizer
+    )
+
+    ## prediction
+    preds = trainer.predict(tokenized_test)
+    predictions = np.array([np.argmax(x) for x in preds[0]])
 
- 
 
 
     #--------------------------------------------------------------------------------------------