Update app.py

app.py

@@ -1,38 +1,32 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-# Load model — stick with this one for now, later we can upgrade
-generator = pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-wikiSQL")
+# Load the GaussAlgo model
+model_path = "gaussalgo/T5-LM-Large-text2sql-spider"
+model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
 
-def convert_to_sql(payload):
+def generate_sql(payload):
+    # Extract components from payload
     question = payload.get("question", "")
     schema = payload.get("schema", "")
-    sample_rows = payload.get("sample_rows", [])
-
-    # Craft prompt
-    prompt = f"""
-You are an AI that converts natural language into SQL for a DuckDB database.
-Given a table with the following schema:
-{schema}
-
-Here are some sample rows:
-{sample_rows}
-
-Write a syntactically correct SQL query (DuckDB-compatible) to answer this question: "{question}"
-
-Only return the SQL query — no explanation, no markdown.
-    """.strip()
-
-    result = generator(prompt, max_length=256)[0]["generated_text"]
-    return result.strip()
-
-# Define inputs/outputs for interactive mode (not used by FastAPI)
+    
+    # Build model input
+    full_prompt = f"Question: {question} Schema: {schema}"
+    
+    inputs = tokenizer(full_prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_length=512)
+    generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return generated_sql
+
+# Define expected input as a JSON object (dict)
 demo = gr.Interface(
-    fn=convert_to_sql,
-    inputs=gr.JSON(label="question + schema + sample_rows"),
+    fn=generate_sql,
+    inputs=gr.JSON(label="Input JSON (with 'question' and 'schema')"),
     outputs="text",
-    title="Text-to-SQL Generator (DuckDB)",
-    description="Send a JSON payload with question, schema, and sample rows"
+    title="Text-to-SQL Generator",
+    description="Input a JSON with your natural language question and database schema. Output is SQL."
 )
 
 demo.launch()