dev

app.py

@@ -20,13 +20,11 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
             lora_dropout=0.1,  # Dropout for LoRA layers
             bias="none"  # Bias handling
         )        
-        # Load the dataset
-        dataset = load_dataset(dataset_name.strip())
     
         # Load the model and tokenizer
         model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2)
         model = get_peft_model(model, lora_config)
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
     
         # Set training arguments
         training_args = TrainingArguments(
@@ -52,6 +50,10 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
             save_steps=40,  # Save checkpoint every 500 steps
             save_total_limit=3, 
         )
+        # Check if a checkpoint exists and load it
+        if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
+            print("Loading model from checkpoint...")
+            model = AutoModelForSeq2SeqLM.from_pretrained(training_args.output_dir)        
     
         max_length = 128
         try:
@@ -67,6 +69,9 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
                 #callbacks=[LoggingCallback()], 
             )            
         except:
+            # Load the dataset
+            dataset = load_dataset(dataset_name.strip())
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
             # Tokenize the dataset
             def tokenize_function(examples):
                 
@@ -106,10 +111,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
             )            
     
 
-        # Check if a checkpoint exists and load it
-        if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
-            print("Loading model from checkpoint...")
-            model = AutoModelForSeq2SeqLM.from_pretrained(training_args.output_dir)
+        
             
 
 
@@ -138,7 +140,7 @@ try:
             gr.Textbox(label="HF hub to push to after training"),
             gr.Textbox(label="HF API token"),
             gr.Slider(minimum=1, maximum=10, value=3, label="Number of Epochs", step=1),
-            gr.Slider(minimum=1, maximum=16, value=1, label="Batch Size", step=1),
+            gr.Slider(minimum=1, maximum=60, value=1, label="Batch Size", step=1),
             gr.Slider(minimum=1, maximum=1000, value=1, label="Learning Rate (e-5)", step=1),
             gr.Slider(minimum=1, maximum=100, value=1, label="Gradient accumulation", step=1), 
         ],