Update app.py

app.py

@@ -8,12 +8,14 @@ ACTION_1 = "Prompt base model"
 ACTION_2 = "Fine-tune base model"
 ACTION_3 = "Prompt fine-tuned model"
 
+HF_ACCOUNT = "bstraehle"
+
 SYSTEM_PROMPT = "You are a text to SQL query translator. Given a question in English, generate a SQL query based on the provided SQL_CONTEXT. Do not generate any additional text. SQL_CONTEXT: {sql_context}"
 USER_PROMPT = "How many new users joined from countries with stricter data privacy laws than the United States in the past month?"
 SQL_CONTEXT = "CREATE TABLE users (user_id INT, country VARCHAR(50), joined_date DATE); CREATE TABLE data_privacy_laws (country VARCHAR(50), privacy_level INT); INSERT INTO users (user_id, country, joined_date) VALUES (1, 'USA', '2023-02-15'), (2, 'Germany', '2023-02-27'); INSERT INTO data_privacy_laws (country, privacy_level) VALUES ('USA', 5), ('Germany', 8);"
 
 BASE_MODEL_NAME = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-FT_MODEL_NAME = "bstraehle/Meta-Llama-3.1-8B-Instruct-text-to-sql"
+FT_MODEL_NAME = "Meta-Llama-3.1-8B-Instruct-text-to-sql"
 DATASET_NAME = "gretelai/synthetic_text_to_sql"
 
 def process(action, base_model_name, ft_model_name, dataset_name, system_prompt, user_prompt, sql_context):
@@ -77,11 +79,11 @@ def fine_tune_model(base_model_name, dataset_name):
     # Configure training arguments
     
     training_args = Seq2SeqTrainingArguments(
-        output_dir="./Meta-Llama-3.1-8B-Instruct-text-to-sql",
+        output_dir=f"./{FT_MODEL_NAME}",
         logging_dir="./logs",
         num_train_epochs=1,
-        max_steps=2, # overwrites num_train_epochs
-        push_to_hub=True,
+        max_steps=1, # overwrites num_train_epochs
+        push_to_hub=True, # only model, also need to push tokenizer
         #per_device_train_batch_size=16,
         #per_device_eval_batch_size=64,
         #eval_strategy="steps",
@@ -114,9 +116,9 @@ def fine_tune_model(base_model_name, dataset_name):
     
     trainer.train()
 
-    # Save tokenizer to HF
+    # Push tokenizer to HF
 
-    tokenizer.push_to_hub("Meta-Llama-3.1-8B-Instruct-text-to-sql")
+    tokenizer.push_to_hub(FT_MODEL_NAME)
     
 def prompt_model(model_name, system_prompt, user_prompt, sql_context):
     pipe = pipeline("text-generation", 
@@ -150,7 +152,7 @@ def load_model(model_name):
 demo = gr.Interface(fn=process, 
                     inputs=[gr.Radio([ACTION_1, ACTION_2, ACTION_3], label = "Action", value = ACTION_3),
                             gr.Textbox(label = "Base Model Name", value = BASE_MODEL_NAME, lines = 1),
-                            gr.Textbox(label = "Fine-Tuned Model Name", value = FT_MODEL_NAME, lines = 1),
+                            gr.Textbox(label = "Fine-Tuned Model Name", value = f"{HF_ACCOUNT}/{FT_MODEL_NAME}", lines = 1),
                             gr.Textbox(label = "Dataset Name", value = DATASET_NAME, lines = 1),
                             gr.Textbox(label = "System Prompt", value = SYSTEM_PROMPT, lines = 2),
                             gr.Textbox(label = "User Prompt", value = USER_PROMPT, lines = 2),