Update app.py

app.py

@@ -1,47 +1,29 @@
 import gradio as gr
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# Load the tokenizer and model
-tokenizer = T5Tokenizer.from_pretrained('t5-small', legacy=False)
-model = T5ForConditionalGeneration.from_pretrained('t5-small')
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("hrshtsharma2012/NL2SQL-Picard-final")
+model = AutoModelForSeq2SeqLM.from_pretrained("hrshtsharma2012/NL2SQL-Picard-final")
 
-# dataset = load_dataset("b-mc2/sql-create-context")
-dataset = load_dataset("wikisql", split="train")
+# Initialize the pipeline
+nl2sql_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
-examples = []
-
-for i in range(3):  # Let's take the first 3 examples
-    item = dataset[i]
-    question = item['question']
-    examples.append([question])
-
-def generate_sql(question):
-    # Format the question for the model if needed. For example:
-    input_text = f"translate English to SQL: {question}"
-    # input_text = f"{question}"  # Directly use the question if the model is fine-tuned for SQL generation
-    
-    # Tokenize the input text
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
-    
-    # Generate the output sequence
-    output_ids = model.generate(input_ids, max_length=512, num_beams=5)[0]
-    
-    # Decode the generated ids to get the SQL query
-    sql_query = tokenizer.decode(output_ids, skip_special_tokens=True)
+def generate_sql(query):
+    # Use the model to generate SQL from the natural language query
+    results = nl2sql_pipeline(query)
+    # Extract the first result (highest likelihood)
+    sql_query = results[0]['generated_text']
     return sql_query
 
-    
-# Define the Gradio interface
-iface = gr.Interface(
+# Create a Gradio interface
+interface = gr.Interface(
     fn=generate_sql,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your question here..."),
-    outputs=gr.Textbox(),
-    title="Natural Language to SQL",
-    description="This app uses a Seq2Seq model to generate SQL queries from natural language questions.",
-    examples=examples
+    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your natural language query here..."),
+    outputs="text",
+    title="NL to SQL with Picard",
+    description="This model converts natural language queries into SQL. It's based on the Spider dataset. Enter a query to get started!"
 )
 
 # Launch the app
 if __name__ == "__main__":
-    iface.launch()
+    interface.launch()