Update app.py

app.py

@@ -3,8 +3,8 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from datasets import load_dataset
 
 # Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("hrshtsharma2012/NL2SQL-Picard-final")
-model = AutoModelForSeq2SeqLM.from_pretrained("hrshtsharma2012/NL2SQL-Picard-final")
+tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
+model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
 
 # Initialize the pipeline
 nl2sql_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
@@ -13,12 +13,14 @@ nl2sql_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokeni
 spider_dataset = load_dataset("spider", split='train[:5]')
 
 def generate_sql(query):
-    formatted_query = f"translate English to SQL: {query}"
-    results = nl2sql_pipeline(formatted_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']
     return sql_query
 
-
 # Use examples from the Spider dataset
 example_questions = [(question['question'],) for question in spider_dataset]
 
@@ -28,10 +30,10 @@ interface = gr.Interface(
     inputs=gr.Textbox(lines=2, placeholder="Enter your natural language query here..."),
     outputs="text",
     examples=example_questions,
-    title="NL to SQL with Picard",
-    description="This model converts natural language queries into SQL using the Spider dataset. Try one of the example questions or enter your own!"
+    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!"
 )
 
 # Launch the app
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()