Update app.py

app.py

@@ -5,12 +5,29 @@ import pandas as pd
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# 1) Load and preprocess your synthetic profitability dataset
+# 1) Load your synthetic profitability dataset
 df = pd.read_csv('synthetic_profit.csv')
-# Ensure numeric Profit for aggregation
-df['Profit'] = pd.to_numeric(df['Profit'], errors='coerce')
 
-# 2) Model setup
+# 2) Ensure numeric columns for true aggregation (optional, but helps you verify sums)
+for col in ["Revenue", "Profit", "ProfitMargin"]:
+    df[col] = pd.to_numeric(df[col], errors='coerce')
+
+# 3) Build the schema description text
+schema_lines = [f"- {col}: {dtype.name}" for col, dtype in df.dtypes.iteritems()]
+schema_text = "Table schema:\n" + "\n".join(schema_lines)
+
+# 4) Few-shot examples teaching SUM and AVERAGE patterns
+example_block = """
+Example 1
+Q: Total profit by region?
+A: Group “Profit” by “Region” and sum → EMEA: 30172183.37; APAC: 32301788.32; Latin America: 27585378.50; North America: 25473893.34
+
+Example 2
+Q: Average profit margin for Product B in Americas?
+A: Filter Product=B & Region=Americas, take mean of “ProfitMargin” → 0.18
+""".strip()
+
+# 5) Model & pipeline setup
 MODEL_ID = "microsoft/tapex-base-finetuned-wikisql"
 device   = 0 if torch.cuda.is_available() else -1
 
@@ -25,40 +42,35 @@ table_qa = pipeline(
     device=device,
 )
 
-# 3) QA function with manual fallback for region‐based aggregations
+# 6) QA function with schema-aware prompting
 def answer_profitability(question: str) -> str:
-    q_lower = question.lower()
+    # 6a) cast all cells to string for safety
+    table = df.astype(str).to_dict(orient="records")
+
+    # 6b) assemble the full prompt
+    prompt = f"""{schema_text}
 
-    # Fallback: if user asks for total profit by region, do it in pandas
-    if "total profit" in q_lower and "region" in q_lower:
-        agg = df.groupby('Region', as_index=False)['Profit'].sum()
-        return "\n".join(f"{row.Region}: {row.Profit}" for row in agg.itertuples())
+{example_block}
 
-    # Otherwise, cast all cells to string and try TAPEX
-    df_str = df.astype(str)
-    table  = df_str.to_dict(orient="records")
+Q: {question}
+A:"""
 
+    # 6c) call TAPEX
     try:
-        out = table_qa(table=table, query=question)
+        out = table_qa(table=table, query=prompt)
         return out.get("answer", "No answer found.")
-    except IndexError as e:
-        # Catch the 'index out of range' and fallback to pandas for any region/group queries
-        if "index out of range" in str(e):
-            agg = df.groupby('Region', as_index=False)['Profit'].sum()
-            return "\n".join(f"{row.Region}: {row.Profit}" for row in agg.itertuples())
-        return f"Error: {e}"
     except Exception as e:
         return f"Error: {e}"
 
-# 4) Gradio interface
+# 7) Gradio interface
 iface = gr.Interface(
     fn=answer_profitability,
     inputs=gr.Textbox(lines=2, placeholder="Ask a question about profitability…"),
     outputs="text",
-    title="SAP Profitability Q&A (TAPEX-Base)",
+    title="SAP Profitability Q&A (Schema-Aware TAPEX)",
     description=(
-        "Free-form questions on the synthetic profitability dataset, "
-        "powered by microsoft/tapex-base-finetuned-wikisql with pandas fallbacks."
+        "Every query is prefixed with your table’s schema and two few-shot examples, "
+        "so the model learns to SUM, AVERAGE, FILTER, etc., without extra Python code."
     )
 )