Update app.py

app.py

@@ -1,25 +1,42 @@
-# app.py
-
 import gradio as gr
 import pandas as pd
 from transformers import pipeline
 
-# 1) Load your synthetic SAP data (all as strings so TAPAS never sees floats)
+# 1) Load your data
 df = pd.read_csv("synthetic_profit.csv")
 table = df.astype(str).to_dict(orient="records")
 
-# 2) Instantiate TAPAS fine-tuned on WikiTableQuestions
+# 2) TAPAS pipeline
 qa = pipeline(
     "table-question-answering",
     model="google/tapas-base-finetuned-wtq",
     tokenizer="google/tapas-base-finetuned-wtq",
-    device=-1  # CPU; change to 0 if you enable GPU
+    device=-1
 )
 
-# 3) Simple QA function
+# 3) Four few-shot examples: sum revenue, sum cost, sum margin, **mean margin**
+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
+
+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
+
+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
+
+Example 4:
+Q: What is the average profit margin for Product A in EMEA in Q1 2024?
+A: Filter Product=A & Region=EMEA & FiscalYear=2024 & FiscalQuarter=Q1, then mean ProfitMargin → 0.18
+"""
+
 def answer_question(question: str) -> str:
+    prompt = EXAMPLE_PROMPT + f"\nQ: {question}\nA:"
     try:
-        out = qa(table=table, query=question)
+        out = qa(table=table, query=prompt)
         return out.get("answer", "No answer found.")
     except Exception as e:
         return f"❌ Error: {e}"
@@ -27,10 +44,13 @@ def answer_question(question: str) -> str:
 # 4) 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?"),
-    outputs=gr.Textbox(lines=4),
+    inputs=gr.Textbox(lines=2, placeholder="e.g. What is the average profit margin by region?", label="Your question"),
+    outputs=gr.Textbox(lines=4, label="Answer"),
     title="SAP Profitability Q&A (TAPAS)",
-    description="Ask simple sum/avg questions on your SAP data. Powered by google/tapas-base-finetuned-wtq.",
+    description=(
+        "Ask simple sum/mean questions on your SAP data.  \n"
+        "Powered by google/tapas-base-finetuned-wtq with four few-shot examples."
+    ),
     allow_flagging="never",
 )