Update app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
-from transformers import T5Tokenizer, T5ForConditionalGeneration, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from datasets import load_dataset
 
 # Load tokenizer and model
-tokenizer = T5Tokenizer.from_pretrained("t5-base")
-model = T5ForConditionalGeneration.from_pretrained("t5-base")
+tokenizer = AutoTokenizer.from_pretrained("hrshtsharma2012/NL2SQL-Picard-final")
+model = AutoModelForSeq2SeqLM.from_pretrained("hrshtsharma2012/NL2SQL-Picard-final")
 
 # Initialize the pipeline
 nl2sql_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
@@ -13,11 +13,7 @@ nl2sql_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokeni
 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, max_length=512, num_return_sequences=1)
-    # Extract the SQL query
+    results = nl2sql_pipeline(query)
     sql_query = results[0]['generated_text']
     return sql_query
 
@@ -30,10 +26,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 T5",
+    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!"
 )
 
 # Launch the app
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()