Update tasks/text.py

tasks/text.py

@@ -63,6 +63,9 @@ async def evaluate_text(request: TextEvaluationRequest):
     from transformers import AutoModelForSequenceClassification, AutoTokenizer
     import torch
     from torch.utils.data import DataLoader, TensorDataset
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
   
     # Load model and tokenizer from Hugging Face Hub
     MODEL_REPO = "ClimateDebunk/FineTunedDistilBert4SeqClass"
@@ -71,7 +74,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     MAX_LENGTH = 365
 
     model = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO)
-    #model.to(device)
+    model.to(device)
     model.eval()  # Set to evaluation mode
 
 
@@ -85,7 +88,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     predictions = []
     with torch.no_grad():
         for batch in test_loader:
-    #        input_ids, attention_mask, labels = [x.to(device) for x in batch]
+            input_ids, attention_mask, labels = [x.to(device) for x in batch]
             outputs = model(input_ids, attention_mask=attention_mask)
             preds = torch.argmax(outputs.logits, dim=1)
             predictions.extend(preds.cpu().numpy())