Update app.py

app.py

@@ -1,39 +1,23 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from datasets import load_dataset
+from transformers import pipeline
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
-model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
+# Load the NL2SQL model
+model_id = "hrshtsharma2012/NL2SQL-Picard-final"
+nl2sql = pipeline("text2text-generation", model=model_id)
 
-# Initialize the pipeline
-nl2sql_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
-
-# Load a part of the Spider dataset
-spider_dataset = load_dataset("spider", split='train[:5]')
-
-def generate_sql(query):
-    # Format the input for the model
-    input_text = f"translate English to SQL: {query}"
-    # Run the pipeline
-    results = nl2sql_pipeline(input_text)
-    # Extract the SQL query
-    sql_query = results[0]['generated_text']
+def generate_sql(nl_query):
+    # Generate the SQL query from the natural language input
+    sql_query = nl2sql(nl_query)[0]["generated_text"]
     return sql_query
 
-# Use examples from the Spider dataset
-example_questions = [(question['question'],) for question in spider_dataset]
-
-# Create a Gradio interface
+# Create the Gradio interface
 interface = gr.Interface(
     fn=generate_sql,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your natural language query here..."),
-    outputs="text",
-    examples=example_questions,
-    title="NL to SQL with T5",
-    description="This model converts natural language queries into SQL. Try one of the example questions or enter your own!"
+    inputs=gr.Textbox(label="Natural Language Query"),
+    outputs=gr.Textbox(label="SQL Query"),
+    title="NL to SQL Converter",
+    description="Enter a natural language query and get the corresponding SQL query.",
 )
 
-# Launch the app
 if __name__ == "__main__":
     interface.launch()