Update app.py

app.py

@@ -60,10 +60,10 @@ def fine_tune_model(base_model_name, dataset_name):
     print("###")
     
     # Split dataset into training and validation sets
-    train_dataset = dataset["train"]
-    test_dataset = dataset["test"]
-    #train_dataset = dataset["train"].shuffle(seed=42).select(range(1000))
-    #test_dataset = dataset["test"].shuffle(seed=42).select(range(100))
+    #train_dataset = dataset["train"]
+    #test_dataset = dataset["test"]
+    train_dataset = dataset["train"].shuffle(seed=42).select(range(1000))
+    test_dataset = dataset["test"].shuffle(seed=42).select(range(100))
 
     print("### Training dataset")
     print(train_dataset)
@@ -77,7 +77,7 @@ def fine_tune_model(base_model_name, dataset_name):
         logging_dir="./logging",
         num_train_epochs=1,
         max_steps=1, # overwrites num_train_epochs
-        push_to_hub=True,
+        #push_to_hub=True,
         #per_device_train_batch_size=16,
         #per_device_eval_batch_size=64,
         #eval_strategy="steps",
@@ -107,6 +107,9 @@ def fine_tune_model(base_model_name, dataset_name):
 
     # Train model
     trainer.train()
+
+    # Save model to HF
+    trainer.push_to_hub(FT_MODEL_NAME, use_auth_token=True)
     
 def prompt_model(model_name, system_prompt, user_prompt, sql_context):
     pipe = pipeline("text-generation",