Update app.py

app.py

@@ -104,15 +104,15 @@ if (runModel=='1'):
     # Create an instance of the custom loss function
     training_args = TrainingArguments(
         output_dir='./results_' + modelNameToUse,
-        num_train_epochs=25,
-        per_device_train_batch_size=2,
-        per_device_eval_batch_size=2,
+        num_train_epochs=5,
+        per_device_train_batch_size=8,
+        per_device_eval_batch_size=8,
         warmup_steps=500,
         weight_decay=0.02,
         logging_dir='./logs_' + modelNameToUse,
         logging_steps=10,
         evaluation_strategy="epoch",
-        
+        load_best_model_at_end=True,  # Load the best model based on evaluation
     )
     
     trainer = Trainer(
@@ -167,7 +167,13 @@ if (runModel=='1'):
         else:
             print("\nNo incorrect predictions found.")
 
-    train_dataloader = DataLoader(train_dataset, batch_size=10, shuffle=True)
+    train_dataloader = DataLoader(
+        train_dataset, 
+        batch_size=10, 
+        shuffle=True, 
+        num_workers=4  # Increase workers for faster data loading
+    )
+    
     evaluate_and_report_errors(model,train_dataloader, tokenizer)
 
     model_path = './'  + modelNameToUse + '_model'