Update tasks/text.py

tasks/text.py

@@ -92,6 +92,8 @@ async def evaluate_text(request: TextEvaluationRequest):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = AutoModelForSequenceClassification.from_pretrained(path_model).to(device).eval()
     tokenizer = AutoTokenizer.from_pretrained(path_tokenizer)
+
+    model.half()
     
     # Use optimized tokenization
     def preprocess_function(df):
@@ -106,10 +108,10 @@ async def evaluate_text(request: TextEvaluationRequest):
         return {"input_ids": input_ids, "attention_mask": attention_mask}
     
     # Optimized inference function
-    def predict(dataset):
+    def predict(dataset, batch_size=16):
         all_preds = []
         with torch.no_grad():  # No gradient computation (saves energy)
-            for batch in torch.utils.data.DataLoader(dataset, batch_size=len(dataset), collate_fn=collate_fn):
+            for batch in torch.utils.data.DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn):
                 outputs = model(**batch)
                 preds = torch.argmax(outputs.logits, dim=-1).cpu().numpy()
                 all_preds.extend(preds)
@@ -117,7 +119,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     
     # Run inference
     predictions = predict(tokenized_test)
-
+    print(predictions)
     # predictions = np.array([np.argmax(x) for x in preds[0]])
 
     #--------------------------------------------------------------------------------------------