update

tasks/text.py

@@ -65,7 +65,7 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
-    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+    tokenizer = AutoTokenizer.from_pretrained("cococli/bert-base-uncased-frugalai")
     model = AutoModelForSequenceClassification.from_pretrained("cococli/bert-base-uncased-frugalai").to(device)
     
     def tokenize_function(examples):
@@ -76,7 +76,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     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)
@@ -88,7 +88,9 @@ async def evaluate_text(request: TextEvaluationRequest):
     predictions = np.array([])
 
     with torch.no_grad():
+        print('BEFORE PREDICTION')
         for batch in dataloader:
+            print('INSIDE PREDICTION')
             test_input_ids = batch["input_ids"].to(device)
             test_attention_mask = batch["attention_mask"].to(device)
             outputs = model(test_input_ids, test_attention_mask)