Update app.py

app.py

@@ -1,75 +1,82 @@
-## https://www.kaggle.com/code/unravel/fine-tuning-of-a-sql-model
-
-import spaces
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-import gradio as gr
+import streamlit as st
 import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from transformers.utils import logging
-from example_queries import small_query, long_query
 
+# Set up logging
 logging.set_verbosity_info()
 logger = logging.get_logger("transformers")
 
-model_name='t5-small'
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-original_model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
-
-ft_model_name="daljeetsingh/sql_ft_t5small_kag" #"cssupport/t5-small-awesome-text-to-sql"
-ft_model = AutoModelForSeq2SeqLM.from_pretrained(ft_model_name, torch_dtype=torch.bfloat16)
+# Model names
+original_model_name = 't5-small'
+fine_tuned_model_name = 'daljeetsingh/sql_ft_t5small_kag'
 
-original_model.to('cuda')
-ft_model.to('cuda')
+# Load models and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(original_model_name)
+original_model = AutoModelForSeq2SeqLM.from_pretrained(original_model_name, torch_dtype=torch.bfloat16)
+fine_tuned_model = AutoModelForSeq2SeqLM.from_pretrained(fine_tuned_model_name, torch_dtype=torch.bfloat16)
 
-@spaces.GPU
-def translate_text(text):
-   prompt = f"{text}"
-   inputs = tokenizer(prompt, return_tensors='pt')
-   inputs = inputs.to('cuda')
+# Move models to GPU
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+original_model.to(device)
+fine_tuned_model.to(device)
 
-   try:
-        output = tokenizer.decode(
-            original_model.generate(
-                inputs["input_ids"], 
-                max_new_tokens=200,
-            )[0], 
+def generate_sql_query(prompt):
+    """
+    Generate SQL queries using both the original and fine-tuned models.
+    """
+    inputs = tokenizer(prompt, return_tensors='pt').to(device)
+    try:
+        # Generate output from the original model
+        original_output = original_model.generate(
+            inputs["input_ids"], 
+            max_new_tokens=200,
+        )
+        original_sql = tokenizer.decode(
+            original_output[0], 
             skip_special_tokens=True
         )
-        ft_output = tokenizer.decode(
-            ft_model.generate(
-                inputs["input_ids"], 
-                max_new_tokens=200,
-            )[0], 
+
+        # Generate output from the fine-tuned model
+        fine_tuned_output = fine_tuned_model.generate(
+            inputs["input_ids"], 
+            max_new_tokens=200,
+        )
+        fine_tuned_sql = tokenizer.decode(
+            fine_tuned_output[0], 
             skip_special_tokens=True
         )
-        return [output, ft_output]
-   except Exception as e:
-       return f"Error: {str(e)}"
 
+        return original_sql, fine_tuned_sql
+    except Exception as e:
+        logger.error(f"Error: {str(e)}")
+        return f"Error: {str(e)}", None
 
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(
-                value=small_query,
-                lines=8,
-                placeholder="Enter prompt...",
-                label="Prompt"
-            )
-            submit_btn = gr.Button(value="Generate")
-        with gr.Column():
-            orig_output = gr.Textbox(label="OriginalModel", lines=2) 
-            ft_output = gr.Textbox(label="FTModel", lines=8)
+# Streamlit App Interface
+st.title("SQL Query Generation")
+st.markdown("This application generates SQL queries based on your input prompt.")
 
-    submit_btn.click(
-        translate_text, inputs=[prompt], outputs=[orig_output, ft_output], api_name=False
-    )
-    examples = gr.Examples(
-        examples=[
-            [small_query],
-            [long_query],
-        ],
-        inputs=[prompt],
-    )
+# Input prompt
+prompt = st.text_area(
+    "Enter your prompt here...",
+    value="Find all employees who joined after 2020.",
+    height=150
+)
 
-demo.launch(show_api=False, share=True, debug=True)
+# Generate button
+if st.button("Generate"):
+    if prompt:
+        original_sql, fine_tuned_sql = generate_sql_query(prompt)
+        st.subheader("Original Model Output")
+        st.text_area("Original SQL Query", value=original_sql, height=200)
+        st.subheader("Fine-Tuned Model Output")
+        st.text_area("Fine-Tuned SQL Query", value=fine_tuned_sql, height=200)
+    else:
+        st.warning("Please enter a prompt to generate SQL queries.")
 
+# Examples
+st.sidebar.title("Examples")
+st.sidebar.markdown("""
+- **Example 1**: Find all employees who joined after 2020.
+- **Example 2**: Retrieve the names of customers who purchased product X in the last month.
+""")