dev

app.py

@@ -28,11 +28,44 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
         model = get_peft_model(model, lora_config)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
     
+        # Set training arguments
+        training_args = TrainingArguments(
+            output_dir='./results',
+            eval_strategy="steps",  # Change this to steps
+            save_strategy='steps',
+            learning_rate=lr*0.00001,
+            per_device_train_batch_size=int(batch_size),
+            per_device_eval_batch_size=int(batch_size), 
+            num_train_epochs=int(num_epochs),
+            weight_decay=0.01,
+            gradient_accumulation_steps=int(grad),
+            max_grad_norm = 1.0, 
+            load_best_model_at_end=True,
+            metric_for_best_model="accuracy",
+            greater_is_better=True,
+            logging_dir='./logs',
+            logging_steps=10,
+            #push_to_hub=True,
+            hub_model_id=hub_id.strip(),
+            fp16=True,
+            #lr_scheduler_type='cosine',
+            save_steps=200,  # Save checkpoint every 500 steps
+            save_total_limit=3, 
+        )
+    
         max_length = 128
         try:
             tokenized_train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
             tokenized_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_test_dataset')
-            tokenized_datasets = concatenate_datasets([tokenized_train_dataset, tokenized_test_dataset])
+            
+            # Create Trainer
+            trainer = Trainer(
+                model=model,
+                args=training_args,
+                train_dataset=tokenized_train_dataset,
+                eval_dataset=tokenized_test_dataset,
+                #callbacks=[LoggingCallback()], 
+            )            
         except:
             # Tokenize the dataset
             def tokenize_function(examples):
@@ -63,44 +96,22 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
             tokenized_datasets['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset')
             tokenized_datasets['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
         
+            # Create Trainer
+            trainer = Trainer(
+                model=model,
+                args=training_args,
+                train_dataset=tokenized_datasets['train'],
+                eval_dataset=tokenized_datasets['test'],
+                #callbacks=[LoggingCallback()], 
+            )            
     
-        # Set training arguments
-        training_args = TrainingArguments(
-            output_dir='./results',
-            eval_strategy="steps",  # Change this to steps
-            save_strategy='steps',
-            learning_rate=lr*0.00001,
-            per_device_train_batch_size=int(batch_size),
-            per_device_eval_batch_size=int(batch_size), 
-            num_train_epochs=int(num_epochs),
-            weight_decay=0.01,
-            gradient_accumulation_steps=int(grad),
-            max_grad_norm = 1.0, 
-            load_best_model_at_end=True,
-            metric_for_best_model="accuracy",
-            greater_is_better=True,
-            logging_dir='./logs',
-            logging_steps=10,
-            #push_to_hub=True,
-            hub_model_id=hub_id.strip(),
-            fp16=True,
-            #lr_scheduler_type='cosine',
-            save_steps=200,  # 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)
             
-        # Create Trainer
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=tokenized_datasets['train'],
-            eval_dataset=tokenized_datasets['test'],
-            #callbacks=[LoggingCallback()], 
-        )
+
 
         # Fine-tune the model
         trainer.train()