Update app.py

app.py

@@ -91,7 +91,7 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
         "lr": 5.701568055793089e-04,
         "lr_scheduler_type": "cosine",
         "max_grad_norm": 0.5,
-        "num_train_epochs": 2,  #3, jw 20240628
+        "num_train_epochs": 1,  #3, jw 20240628
         "per_device_train_batch_size": 12,
         "r": 2,
         "weight_decay": 0.2,
@@ -185,8 +185,7 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
         seed=8893,
         fp16=True,
         #report_to='wandb'
-        report_to=None,
-        class_weights=class_weights,  #jw, 20240628
+        report_to=None
     )
 
     # Initialize Trainer
@@ -197,7 +196,8 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
         eval_dataset=test_dataset,
         tokenizer=tokenizer,
         data_collator=DataCollatorForTokenClassification(tokenizer=tokenizer),
-        compute_metrics=compute_metrics
+        compute_metrics=compute_metrics,
+        class_weights=class_weights,  #add class_weights as input, jw 20240628
     )
 
     # Train and Save Model