Update app.py

app.py

@@ -79,7 +79,8 @@ def fine_tune_model(base_model_name, dataset_name):
     training_args = Seq2SeqTrainingArguments(
         output_dir="./output",
         logging_dir="./logging",
-        num_train_epochs=1,
+        num_train_epochs=1
+        max_steps=2, ###
         #per_device_train_batch_size=16,
         #per_device_eval_batch_size=64,
         #eval_strategy="steps",
@@ -109,8 +110,8 @@ def fine_tune_model(base_model_name, dataset_name):
     )
 
     # Train and save model
-    #trainer.train()
-    #trainer.save_model()
+    trainer.train()
+    trainer.save_model()
     
 def prompt_model(model_name, system_prompt, user_prompt, sql_schema):
     pipe = pipeline("text-generation", 
@@ -136,7 +137,7 @@ def prompt_model(model_name, system_prompt, user_prompt, sql_schema):
 
 def load_model(model_name):
     model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
     #tokenizer.pad_token = tokenizer.eos_token
     
     return model, tokenizer