Update app.py

app.py

@@ -1,39 +1,41 @@
 import gradio as gr
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-
-# Load the GaussAlgo model
-model_path = "gaussalgo/T5-LM-Large-text2sql-spider"
-model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
-tokenizer = AutoTokenizer.from_pretrained(model_path)
+# Load FLAN-T5-small
+model_name = "google/flan-t5-small"
+tokenizer = T5Tokenizer.from_pretrained(model_name)
+model = T5ForConditionalGeneration.from_pretrained(model_name)
 
 def generate_sql(payload):
-    # Extract parts from the JSON payload
     question = payload.get("question", "")
     schema = payload.get("schema", "")
     sample_rows = payload.get("sample_rows", [])
 
+    # Convert sample rows into flat string
+    rows_text = " ".join([str(row) for row in sample_rows]) if sample_rows else ""
 
-    # Convert sample rows into a single string
-    sample_str = " ".join([str(row) for row in sample_rows]) if sample_rows else ""
-
-    # Build model input prompt
-    prompt = f"Question: {question} Schema: {schema} Sample Rows: {sample_str}"
+    # Construct prompt for instruction tuning
+    prompt = (
+        f"You are a SQL expert.\n"
+        f"Schema: {schema}\n"
+        f"Sample Rows: {rows_text}\n"
+        f"Question: {question}\n"
+        f"Generate SQL:"
+    )
 
-    # Tokenize and generate
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=512)
-    generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Tokenize and generate SQL
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    outputs = model.generate(input_ids, max_length=256, temperature=0.6)
+    sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    return generated_sql
+    return sql
 
-# Gradio interface
 demo = gr.Interface(
     fn=generate_sql,
-    inputs=gr.JSON(label="Input JSON (question, schema, sample_rows)"),
+    inputs=gr.JSON(label="JSON (question, schema, sample_rows)"),
     outputs="text",
-    title="Text-to-SQL Generator",
-    description="Enter a JSON object with 'question', 'schema', and optional 'sample_rows'. The model will generate SQL."
+    title="FLAN-T5 Text-to-SQL",
+    description="Using FLAN-T5 to generate SQL from natural language and tabular schema."
 )
 
-demo.launch()
+demo.launch()