Update app.py

app.py

@@ -1,51 +1,94 @@
-# app.py
-
+import re
 import gradio as gr
 import pandas as pd
 from transformers import pipeline
 
-# 1) Load your synthetic SAP data
+# 1) Load your data
 df = pd.read_csv("synthetic_profit.csv")
-table = df.astype(str).to_dict(orient="records")
 
-# 2) Use TAPEX fine-tuned on WikiTableQuestions
-qa = pipeline(
+# 2) Prepare the TAPAS fallback
+tapas = pipeline(
     "table-question-answering",
-    model="microsoft/tapex-base-finetuned-wtq",
-    tokenizer="microsoft/tapex-base-finetuned-wtq",
-    device=-1  # CPU; switch to 0 if you have GPU
+    model="google/tapas-base-finetuned-wtq",
+    tokenizer="google/tapas-base-finetuned-wtq",
+    device=-1
 )
+table = df.astype(str).to_dict(orient="records")
+
+# 3) Helpers for parsing
+OPERATIONS = {
+    "total": "sum",
+    "sum":   "sum",
+    "average": "mean",
+    "mean":    "mean"
+}
+COLUMNS = {
+    "revenue":     "Revenue",
+    "cost":        "Cost",
+    "profit":      "Profit",
+    "margin":      "ProfitMargin",
+    "profit margin":"ProfitMargin"
+}
 
-# 3) Few-shot examples teaching “filter + sum”
-EXAMPLES = """
-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
+def parse_and_compute(question: str):
+    q = question.lower()
+    # 1) detect op
+    op = next((OPERATIONS[k] for k in OPERATIONS if k in q), None)
+    # 2) detect column
+    col = next((COLUMNS[k] for k in COLUMNS if k in q), None)
+    # 3) detect product (assumes "Product X")
+    m = re.search(r"product\s*([A-Za-z0-9]+)", q)
+    prod = f"Product {m.group(1)}" if m else None
+    # 4) detect region from known values
+    region = next((r for r in df["Region"].unique() if r.lower() in q), None)
+    # 5) detect year
+    y = re.search(r"\b(20\d{2})\b", q)
+    year = int(y.group(1)) if y else None
+    # 6) detect quarter
+    qtr = next((fq for fq in df["FiscalQuarter"].unique() if fq.lower() in q), None)
+
+    if None in (op, col, prod, region, year, qtr):
+        return None  # fallback
+
+    # filter
+    sub = df[
+        (df["Product"] == prod) &
+        (df["Region"]  == region) &
+        (df["FiscalYear"]   == year) &
+        (df["FiscalQuarter"] == qtr)
+    ]
+
+    # compute
+    try:
+        val = getattr(sub[col], op)()
+    except Exception:
+        return None
 
-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
+    return f"{op.capitalize()} {col} for {prod} in {region}, {qtr} {year}: {val:.2f}"
 
-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
-"""
+# 4) Main answer fn
+def answer(question: str) -> str:
+    res = parse_and_compute(question)
+    if res is not None:
+        return res
 
-def answer_question(question: str) -> str:
-    prompt = EXAMPLES + f"\nQ: {question}\nA:"
+    # fallback to TAPAS
     try:
-        out = qa(table=table, query=prompt)
+        out = tapas(table=table, query=question)
         return out.get("answer", "No answer found.")
     except Exception as e:
         return f"❌ Error: {e}"
 
-# 4) Gradio UI
+# 5) Gradio UI
 iface = gr.Interface(
-    fn=answer_question,
-    inputs=gr.Textbox(lines=2, placeholder="e.g. What is the total revenue for Product A in EMEA in Q1 2024?", label="Your question"),
-    outputs=gr.Textbox(lines=3, label="Answer"),
-    title="SAP Profitability Q&A (TAPEX)",
-    description="Ask basic sum questions on SAP data, powered by TAPEX with few-shot examples.",
+    fn=answer,
+    inputs=gr.Textbox(lines=2, placeholder="e.g. What is the total revenue for Product A in EMEA in Q1 2024?"),
+    outputs=gr.Textbox(lines=3),
+    title="SAP Profitability Q&A",
+    description=(
+        "Supports any basic “total”/“average” question by parsing and computing via Pandas.  \n"
+        "Falls back to TAPAS for anything else."
+    ),
     allow_flagging="never",
 )