Update tasks/text.py

tasks/text.py

@@ -51,7 +51,6 @@ def baseline_model(dataset_length: int):
 
 
 def bert_model(test_dataset: dict, model_type: str):
-    print("Starting my code block.")
     texts = test_dataset["quote"]
 
     model_repo = f"evgeniiarazum/{MODEL_TYPE}"
@@ -67,7 +66,7 @@ def bert_model(test_dataset: dict, model_type: str):
     print("Using device:", device)
     model = model.to(device)
     dataset = TextDataset(texts, tokenizer=tokenizer)
-    dataloader = DataLoader(dataset, batch_size=32, shuffle=False)
+    dataloader = DataLoader(dataset, batch_size=16, shuffle=False)
     model.eval()
     with torch.no_grad():
         print("Starting model run.")
@@ -80,7 +79,6 @@ def bert_model(test_dataset: dict, model_type: str):
             predictions = np.append(predictions, p.cpu().numpy())
         print("End of model run.")
 
-    print("End of my code block.")
     return predictions