Update app.py

app.py

@@ -1,98 +1,52 @@
 import gradio as gr
 import pandas as pd
-import duckdb
-import torch
-from transformers import T5Tokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import pipeline
 
-# Load data
+# 1) Load your synthetic SAP data
 df = pd.read_csv("synthetic_profit.csv")
-conn = duckdb.connect(":memory:")
-conn.register("sap", df)
-schema = ", ".join(df.columns)
-
-# 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,
-    tokenizer=tokenizer,
-    framework="pt",
-    device=device,
-    max_length=128,
+table = df.astype(str).to_dict(orient="records")
+
+# 2) Table‐QA pipeline with TAPEX
+qa = pipeline(
+    "table-question-answering",
+    model="microsoft/tapex-base-finetuned-wtq",
+    tokenizer="microsoft/tapex-base-finetuned-wtq",
+    device=-1
 )
 
-# Prompt→SQL with few-shot
-def generate_sql(question: str) -> str:
-    prompt = f"""
--- DuckDB table `sap` columns: {schema}
-
--- EXAMPLE 1
--- Q: What is the total profit 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()
-    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}")
+# 3) Three few‐shot examples
+EXAMPLE_PROMPT = """
+Example 1:
+Q: What is the total revenue for Product A in EMEA in Q1 2024?
+A: Filter Product=A & Region=EMEA & FiscalYear=2024 & FiscalQuarter=Q1, then sum Revenue → 3075162.49
 
-# NL→SQL→DuckDB→Result
-def answer_profitability(question: str) -> str:
-    try:
-        sql = generate_sql(question)
-    except Exception as e:
-        return f"❌ Prompt/SQL error:\n{e}"
+Example 2:
+Q: What is the total cost for Product A in EMEA in Q1 2024?
+A: Filter Product=A & Region=EMEA & FiscalYear=2024 & FiscalQuarter=Q1, then sum Cost → 2894321.75
 
-    try:
-        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\nGenerated SQL:\n```sql\n{sql}\n```"
+Example 3:
+Q: What is the total margin for Product A in EMEA in Q1 2024?
+A: Filter Product=A & Region=EMEA & FiscalYear=2024 & FiscalQuarter=Q1, then sum ProfitMargin → 0.18
+"""
 
-    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)
+def answer_question(question: str) -> str:
+    # Prepend the examples to teach the model your pattern
+    full_query = EXAMPLE_PROMPT + f"\nQ: {question}\nA:"
+    result = qa(table=table, query=full_query)
+    return result.get("answer", "No answer found.")
 
-# Gradio UI
+# 4) Gradio UI
 iface = gr.Interface(
-    fn=answer_profitability,
-    inputs=gr.Textbox(lines=2, label="Question"),
-    outputs=gr.Textbox(lines=8, label="Answer"),
+    fn=answer_question,
+    inputs=gr.Textbox(lines=2, placeholder="Ask a basic question…", label="Your question"),
+    outputs=gr.Textbox(lines=4, label="Answer"),
     title="SAP Profitability Q&A",
-    description="Translate English → SQL → DuckDB → Answer",
+    description=(
+        "Ask simple revenue/cost/margin questions on the synthetic SAP data. "
+        "Powered by microsoft/tapex-base-finetuned-wtq with three few‐shot examples."
+    ),
     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)