dev

app.py

@@ -58,7 +58,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
             hub_model_id=hub_id.strip(),
             fp16=True,
             #lr_scheduler_type='cosine',
-            save_steps=40,  # Save checkpoint every 500 steps
+            save_steps=100,  # Save checkpoint every 500 steps
             save_total_limit=3, 
         )
         # Check if a checkpoint exists and load it
@@ -69,7 +69,6 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
         max_length = 128
         try:
             tokenized_train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
-            print(tokenized_train_dataset[0])
             tokenized_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_test_dataset')
             
             # Create Trainer
@@ -123,11 +122,6 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
                 compute_metrics=compute_metrics,
                 #callbacks=[LoggingCallback()], 
             )            
-    
-
-        
-            
-
 
         # Fine-tune the model
         trainer.train()