Update app.py

app.py

@@ -82,8 +82,8 @@ def fine_tune_model(base_model_name, dataset_name):
     training_args = Seq2SeqTrainingArguments(
         output_dir=f"./{FT_MODEL_NAME}",
         num_train_epochs=3, # 37,500 steps
-        max_steps=1, # overwrites num_train_epochs
-        push_to_hub=True, # only pushes model, also need to push tokenizer (see below)
+        max_steps=2, # overwrites num_train_epochs
+        #push_to_hub=True, # only pushes model, also need to push tokenizer (see below)
         # TODO https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments
     )
 
@@ -105,8 +105,9 @@ def fine_tune_model(base_model_name, dataset_name):
     
     trainer.train()
 
-    # Push tokenizer to HF
+    # Push model and tokenizer to HF
 
+    model.push_to_hub(FT_MODEL_NAME)
     tokenizer.push_to_hub(FT_MODEL_NAME)
     
 def prompt_model(model_name, system_prompt, user_prompt, sql_context):