Update utilities/modeling.py

utilities/modeling.py

@@ -37,7 +37,9 @@ def get_peft(model, peft, max_seq_length, random_seed):
     return model
 
 
-def get_trainer(model, tokenizer, dataset, sft, data_field, max_seq_length, random_seed):
+def get_trainer(model, tokenizer, dataset, sft, 
+                data_field, max_seq_length, random_seed,
+                num_epochs, max_steps):
     
     trainer = SFTTrainer(
         model = model,
@@ -68,6 +70,7 @@ def get_trainer(model, tokenizer, dataset, sft, data_field, max_seq_length, rand
 
 
 def prepare_trainer(model_name, max_seq_length, random_seed,
+                    num_epochs, max_steps,
                     peft, sft, dataset, data_field):
 
     print("Loading Model")
@@ -77,7 +80,8 @@ def prepare_trainer(model_name, max_seq_length, random_seed,
     model = get_peft(model, peft, max_seq_length, random_seed)
 
     print("Getting Trainer Model")
-    trainer = get_trainer(model, tokenizer, dataset, data_field, max_seq_length, random_seed)
+    trainer = get_trainer(model, tokenizer, dataset, data_field, max_seq_length, random_seed,
+                          num_epochs, max_steps)
 
     return trainer