update

tasks/text.py

@@ -68,21 +68,21 @@ async def evaluate_text(request: TextEvaluationRequest):
     tokenizer = AutoTokenizer.from_pretrained("cococli/bert-base-uncased-frugalai")
     model = AutoModelForSequenceClassification.from_pretrained("cococli/bert-base-uncased-frugalai").to(device)
     
-    def tokenize_function(examples):
-        return tokenizer(examples["quote"], padding=True, truncation=True, return_tensors='pt')
-    print('BEFORE TOKENIZING')
-    # Tokenize the test dataset
-    tokenized_test = test_dataset.map(tokenize_function, batched=True)
-    print('AFTER TOKENIZING')
-    print(tokenized_test.column_names)  # Debugging step
-    print(tokenized_test['input_ids'][:5])  # Debugging step
-
-    # Create DataLoader
-    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
-    dataloader = DataLoader(tokenized_test, batch_size=16, shuffle=False, collate_fn=data_collator)
+    # def tokenize_function(examples):
+    #     return tokenizer(examples["quote"], padding=True, truncation=True, return_tensors='pt')
+    # print('BEFORE TOKENIZING')
+    # # Tokenize the test dataset
+    # tokenized_test = test_dataset.map(tokenize_function, batched=True)
+    # print('AFTER TOKENIZING')
+    # print(tokenized_test.column_names)  # Debugging step
+    # print(tokenized_test['input_ids'][:5])  # Debugging step
+
+    # # Create DataLoader
+    # data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+    # dataloader = DataLoader(tokenized_test, batch_size=16, shuffle=False, collate_fn=data_collator)
 
     print("Started prediction run")
-
+    tokenized_test = tokenizer(test_dataset['quote'], padding=True, truncation=True, return_tensors='pt')
     # Model inference
     model.eval()
     predictions = np.array([])
@@ -90,9 +90,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     with torch.no_grad():
         print('BEFORE PREDICTION')
         
-        test_input_ids = tokenized_test["input_ids"]
-        test_attention_mask = tokenized_test["attention_mask"]
-        outputs = model(test_input_ids, test_attention_mask)
+        outputs = model(**tokenized_test)
         p = torch.argmax(outputs.logits, dim=1)
         predictions = np.append(predictions, p.cpu().numpy())