Update app.py

app.py

@@ -1,26 +1,20 @@
-import os
 import gradio as gr
 import pandas as pd
 import duckdb
 import torch
 from transformers import T5Tokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# 1) Load your synthetic data into DuckDB
+# Load data
 df = pd.read_csv("synthetic_profit.csv")
 conn = duckdb.connect(":memory:")
 conn.register("sap", df)
-
-# 2) Build a one-line schema description
 schema = ", ".join(df.columns)
-# e.g. "Region,Product,FiscalYear,FiscalQuarter,Revenue,Profit,ProfitMargin"
 
-# 3) Prepare the T5-WikiSQL model & tokenizer (slow, SentencePiece)
+# Model & tokenizer
 MODEL_ID = "mrm8488/t5-base-finetuned-wikisql"
 device   = 0 if torch.cuda.is_available() else -1
-
 tokenizer = T5Tokenizer.from_pretrained(MODEL_ID)
 model     = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
-
 sql_generator = pipeline(
     "text2text-generation",
     model=model,
@@ -30,59 +24,75 @@ sql_generator = pipeline(
     max_length=128,
 )
 
-# 4) NL → SQL with schema + example few-shot
+# Prompt→SQL with few-shot
 def generate_sql(question: str) -> str:
     prompt = f"""
 -- DuckDB table `sap` columns: {schema}
 
--- EXAMPLE
+-- EXAMPLE 1
 -- Q: What is the total profit by region?
--- SQL: SELECT Region, SUM(Profit) AS total_profit
---      FROM sap
---      GROUP BY Region;
+-- SQL:
+SELECT
+  Region,
+  SUM(Profit) AS total_profit
+FROM sap
+GROUP BY Region;
+
+-- EXAMPLE 2
+-- Q: What is the total revenue for Product A in EMEA in Q1 2024?
+-- SQL:
+SELECT
+  SUM(Revenue) AS total_revenue
+FROM sap
+WHERE
+  Product = 'Product A'
+  AND Region = 'EMEA'
+  AND FiscalYear = 2024
+  AND FiscalQuarter = 'Q1';
 
 -- NOW
 Q: {question}
 SQL:
-"""
+""".strip()
+
     out = sql_generator(prompt)[0]["generated_text"].strip()
-    # strip ``` if the model wraps it
-    if out.startswith("```") and out.endswith("```"):
-        out = "\n".join(out.splitlines()[1:-1])
-    return out
+    if "SELECT" in out.upper():
+        sql = out[out.upper().index("SELECT"):]
+        if ";" in sql:
+            sql = sql[: sql.rindex(";") + 1]
+        return sql
+    raise ValueError(f"Did not generate a SELECT; got:\n{out}")
 
-# 5) Core QA function: generate SQL, run it, format result
+# NL→SQL→DuckDB→Result
 def answer_profitability(question: str) -> str:
-    sql = generate_sql(question)
+    try:
+        sql = generate_sql(question)
+    except Exception as e:
+        return f"❌ Prompt/SQL error:\n{e}"
 
-    # run the SQL
     try:
-        result_df = conn.execute(sql).df()
+        rel = conn.execute(sql)
+        if rel is None:
+            return f"❌ No relation returned for SQL:\n```sql\n{sql}\n```"
+        df_out = rel.df()
     except Exception as e:
-        return (
-            f"❌ SQL execution error:\n{e}\n\n"
-            f"Generated SQL:\n```sql\n{sql}\n```"
-        )
+        return f"❌ SQL execution error:\n{e}\n\nGenerated SQL:\n```sql\n{sql}\n```"
 
-    # format output
-    if result_df.empty:
-        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)
+    if df_out.empty:
+        return f"No rows.\n\n```sql\n{sql}\n```"
+    if df_out.shape == (1,1):
+        return str(df_out.iat[0,0])
+    return df_out.to_markdown(index=False)
 
-# 6) Gradio UI
+# Gradio UI
 iface = gr.Interface(
     fn=answer_profitability,
-    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-Only SQL Generation)",
-    description=(
-        "Uses a T5-WikiSQL model with schema+example prompting to\n"
-        "translate your question into valid SQL, then runs it in DuckDB."
-    ),
+    inputs=gr.Textbox(lines=2, label="Question"),
+    outputs=gr.Textbox(lines=8, label="Answer"),
+    title="SAP Profitability Q&A",
+    description="Translate English → SQL → DuckDB → Answer",
     allow_flagging="never",
 )
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch(server_name="0.0.0.0", server_port=7860)