Update app.py

app.py

@@ -1,37 +1,37 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from datasets import load_dataset
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-base-multi-summarization-sql-en")
-model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5-base-multi-summarization-sql-en")
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("hrshtsharma2012/NL2SQL-Picard-final")
+model = AutoModelForSeq2SeqLM.from_pretrained("hrshtsharma2012/NL2SQL-Picard-final")
 
-def generate_sql(natural_language_query):
-    # Tokenize the input query
-    input_ids = tokenizer(natural_language_query, return_tensors="pt").input_ids
+# Initialize the pipeline
+nl2sql_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
-    # Generate the SQL query
-    output_ids = model.generate(input_ids, max_length=512)[0]
+# Load a part of the WikiSQL dataset
+wikisql_dataset = load_dataset("wikisql", split='train[:5]')
 
-    # Decode the generated SQL query
-    sql_query = tokenizer.decode(output_ids, skip_special_tokens=True)
+def generate_sql(query):
+    results = nl2sql_pipeline(query)
+    sql_query = results[0]['generated_text']
+    # Post-process the output to ensure it's a valid SQL query
+    sql_query = sql_query.replace('<pad>', '').replace('</s>', '').strip()
     return sql_query
 
-# Example questions for the interface
-example_questions = [
-    "What is the average salary of employees?",
-    "List the names of employees who work in the IT department.",
-    "Count the number of employees who joined after 2015."
-]
+# Use examples from the WikiSQL dataset
+example_questions = [(question['question'],) for question in wikisql_dataset]
 
-# Create the Gradio interface
+# Create a 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 CodeT5",
+    title="NL to SQL with Picard",
     description="This model converts natural language queries into SQL using the WikiSQL dataset. Try one of the example questions or enter your own!"
 )
 
-# Launch the interface
-interface.launch()
+# Launch the app
+if __name__ == "__main__":
+    interface.launch()