Update tasks/text.py

tasks/text.py

@@ -62,7 +62,7 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     from transformers import AutoModelForSequenceClassification, AutoTokenizer
     import torch
-    from torch.utils.data import DataLoader, Dataset
+    from torch.utils.data import DataLoader, TensorDataset
   
     # Load model and tokenizer from Hugging Face Hub
     MODEL_REPO = "ClimateDebunk/FineTunedDistilBert4SeqClass"
@@ -71,14 +71,16 @@ async def evaluate_text(request: TextEvaluationRequest):
     MAX_LENGTH = 365
 
     model = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO)
+    model.to(device)
     model.eval()  # Set to evaluation mode
 
 
     # tokenize texts
     test_encodings = tokenizer(test_dataset["quote"], padding='max_length', truncation=True, max_length=MAX_LENGTH, return_tensors="pt")
-
-    test_dataset = TensorDataset(test_encodings["input_ids"], test_encodings["attention_mask"], test_labels)
-    test_loader = DataLoader(test_dataset, batch_size=16)
+    test_labels = torch.tensor(test_dataset["label"])
+    
+    test_dataset_0 = TensorDataset(test_encodings["input_ids"], test_encodings["attention_mask"], test_labels)
+    test_loader = DataLoader(test_dataset_0, batch_size=16)
     
     predictions = []
     with torch.no_grad():