Update app.py

app.py

@@ -4,24 +4,26 @@ from datasets import load_dataset
 from huggingface_hub import HfApi, login
 from transformers import AutoModelForCausalLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, pipeline
 
-hf_profile = "bstraehle"
+ACTION_1 = "Prompt base model"
+ACTION_2 = "Fine-tune base model"
+ACTION_3 = "Prompt fine-tuned model"
 
-action_1 = "Fine-tune pre-trained model"
-action_2 = "Prompt fine-tuned model"
+SYSTEM_PROMPT = "You are a text to SQL query translator. Given a question in English, generate a SQL query based on the provided SCHEMA. Do not generate any additional text. SCHEMA: {schema}"
+USER_PROMPT = "What is the total trade value and average price for each trader and stock in the trade_history table?"
+SQL_SCHEMA = "CREATE TABLE trade_history (id INT, trader_id INT, stock VARCHAR(255), price DECIMAL(5,2), quantity INT, trade_time TIMESTAMP);"
 
-system_prompt = "You are a text to SQL query translator. Given a question in English, generate a SQL query based on the provided SCHEMA. Do not generate any additional text. SCHEMA: {schema}"
-user_prompt = "What is the total trade value and average price for each trader and stock in the trade_history table?"
-sql_schema = "CREATE TABLE trade_history (id INT, trader_id INT, stock VARCHAR(255), price DECIMAL(5,2), quantity INT, trade_time TIMESTAMP);"
+BASE_MODEL_NAME = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+FT_MODEL_NAME = "bstraehle/Meta-Llama-3.1-8B-Instruct-text-to-sql"
+DATASET_NAME = "gretelai/synthetic_text_to_sql"
 
-model_name = "bstraehle/Meta-Llama-3.1-8B-Instruct-text-to-sql"
-dataset_name = "gretelai/synthetic_text_to_sql"
-
-def process(action, model_name, dataset_name, system_prompt, user_prompt, sql_schema):
+def process(action, base_model_name, ft_model_name, dataset_name, system_prompt, user_prompt, sql_schema):
     #raise gr.Error("Please clone and bring your own credentials.")
     if action == action_1:
-        result = fine_tune_model(model_name, dataset_name)
+        result = prompt_model(base_model_name, system_prompt, user_prompt, sql_schema)
+    elif action == action_2:
+        result = fine_tune_model(base_model_name, dataset_name)
     elif action == action_2:
-        result = prompt_model(model_name, system_prompt, user_prompt, sql_schema)
+        result = prompt_model(ft_model_name, system_prompt, user_prompt, sql_schema)
     return result
 
 def fine_tune_model(model_name, dataset_name):
@@ -104,7 +106,6 @@ def fine_tune_model(model_name, dataset_name):
 def prompt_model(model_name, system_prompt, user_prompt, sql_schema):
     pipe = pipeline("text-generation", 
                     model=model_name,
-                    model_kwargs={"torch_dtype": torch.bfloat16}, 
                     device_map="auto",
                     max_new_tokens=1000)
     
@@ -126,7 +127,8 @@ 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(model_name)
+    #tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer = model.tokenizer
 
     if not tokenizer.pad_token:
         tokenizer.pad_token = tokenizer.eos_token
@@ -135,10 +137,11 @@ def load_model(model_name):
     
 demo = gr.Interface(fn=process, 
                     inputs=[gr.Radio([action_1, action_2], label = "Action", value = action_2),
-                            gr.Textbox(label = "Model Name", value = 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),
-                            gr.Textbox(label = "SQL Schema", value = sql_schema, lines = 2)],
+                            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 = "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),
+                            gr.Textbox(label = "SQL Schema", value = SQL_SCHEMA, lines = 2)],
                     outputs=[gr.Textbox(label = "Prompt Completion", value = os.environ["OUTPUT"])])
 demo.launch()