Update app.py

app.py

@@ -1,24 +1,24 @@
-# app.py
-
+import os
 import gradio as gr
 import pandas as pd
 import duckdb
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import T5Tokenizer, AutoModelForSeq2SeqLM, pipeline
 
 # 1) Load your synthetic data into DuckDB
 df = pd.read_csv("synthetic_profit.csv")
 conn = duckdb.connect(":memory:")
 conn.register("sap", df)
 
-# 2) Build a one-line schema string for prompts
-schema = ", ".join(df.columns)  # e.g. "Region, Product, FiscalYear, ..."
+# 2) Build a one-line schema description
+schema = ", ".join(df.columns)
+# e.g. "Region,Product,FiscalYear,FiscalQuarter,Revenue,Profit,ProfitMargin"
 
-# 3) Load an open-source model fine-tuned on WikiSQL for SQL generation
+# 3) Prepare the T5-WikiSQL model & tokenizer (slow, SentencePiece)
 MODEL_ID = "mrm8488/t5-base-finetuned-wikisql"
 device   = 0 if torch.cuda.is_available() else -1
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
+tokenizer = T5Tokenizer.from_pretrained(MODEL_ID)
 model     = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
 
 sql_generator = pipeline(
@@ -30,25 +30,32 @@ sql_generator = pipeline(
     max_length=128,
 )
 
-# 4) Function to turn a natural-language question into SQL
+# 4) NL → SQL with schema + example few-shot
 def generate_sql(question: str) -> str:
-    prompt = (
-        f"Translate the following English question into SQL for a DuckDB table named `sap` "
-        f"with columns ({schema}):\n\n"
-        f"Question: {question}\nSQL:"
-    )
+    prompt = f"""
+-- DuckDB table `sap` columns: {schema}
+
+-- EXAMPLE
+-- Q: What is the total profit by region?
+-- SQL: SELECT Region, SUM(Profit) AS total_profit
+--      FROM sap
+--      GROUP BY Region;
+
+-- NOW
+Q: {question}
+SQL:
+"""
     out = sql_generator(prompt)[0]["generated_text"].strip()
-    # strip triple-backticks if present
+    # strip ``` if the model wraps it
     if out.startswith("```") and out.endswith("```"):
         out = "\n".join(out.splitlines()[1:-1])
     return out
 
-# 5) Core Q&A function: NL → SQL → execute → format
+# 5) Core QA function: generate SQL, run it, format result
 def answer_profitability(question: str) -> str:
-    # a) generate the SQL
     sql = generate_sql(question)
 
-    # b) execute it
+    # run the SQL
     try:
         result_df = conn.execute(sql).df()
     except Exception as e:
@@ -57,22 +64,22 @@ def answer_profitability(question: str) -> str:
             f"Generated SQL:\n```sql\n{sql}\n```"
         )
 
-    # c) format the result
+    # format output
     if result_df.empty:
-        return f"No results.\n\n```sql\n{sql}\n```"
+        return f"No rows returned.\n\nSQL was:\n```sql\n{sql}\n```"
     if result_df.shape == (1,1):
         return str(result_df.iat[0,0])
     return result_df.to_markdown(index=False)
 
-# 6) Gradio UI with explicit inputs & outputs
+# 6) Gradio UI
 iface = gr.Interface(
     fn=answer_profitability,
-    inputs=gr.Textbox(lines=2, placeholder="Ask a question about profitability…", label="Question"),
+    inputs=gr.Textbox(lines=2, placeholder="Ask a question…", label="Question"),
     outputs=gr.Textbox(lines=8, placeholder="Answer appears here", label="Answer"),
-    title="SAP Profitability Q&A (HF SQL Generation + DuckDB)",
+    title="SAP Profitability Q&A (HF-Only SQL Generation)",
     description=(
-        "Uses an open-source Hugging Face model fine-tuned on WikiSQL to translate your question into SQL, "
-        "executes it against the `sap` table in DuckDB, and returns the result."
+        "Uses a T5-WikiSQL model with schema+example prompting to\n"
+        "translate your question into valid SQL, then runs it in DuckDB."
     ),
     allow_flagging="never",
 )