Update tasks/text.py

tasks/text.py

@@ -91,7 +91,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     model.eval()
 
     predictions = []
-    for text in tqdm(dataset_test_df["quote"]):
+    for text in tqdm(test_dataset["quote"]):
         with torch.no_grad():        
             tokenized_text = tokenizer(text, truncation=True, padding='max_length', return_tensors = "pt")
             inputt = {k:v.to(device) for k,v in tokenized_text.items()}
@@ -111,7 +111,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     model.to(device)
     model.eval()
     
-    for i,text in tqdm(enumerate(dataset_test_df["quote"])):
+    for i,text in tqdm(enumerate(test_dataset["quote"])):
         if isinstance(predictions[i], str):
             continue
         with torch.no_grad():        
@@ -123,7 +123,8 @@ async def evaluate_text(request: TextEvaluationRequest):
             
             prediction = ID2LABEL[taxonomy_prediction]
             predictions[i] = prediction
-    
+
+    predictions = [LABEL_MAPPING[pred] for pred in predictions]
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------