Update tasks/text.py

tasks/text.py

@@ -158,10 +158,15 @@ async def evaluate_text(request: TextEvaluationRequest):
     model.eval()
     predictions = []
     for batch in tqdm(test_dataloader):
-
+        
         with torch.no_grad():
-            logits = model(**batch)    
-
+            if MODEL =="mlp":
+                b_texts = batch
+                logits = model(b_texts)
+            elif MODEL == "ct":
+                b_input_ids, b_input_mask, b_token_type_ids, b_labels = batch
+                logits = model(b_input_ids, b_token_type_ids, b_input_mask)
+                
         logits = logits.detach().cpu().numpy()
         predictions.extend(logits.argmax(1))