Update app.py

app.py

@@ -5,20 +5,18 @@ import pandas as pd
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# 1) Load your synthetic profitability dataset
+# 1) Load and preprocess 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) Choose the publicly available TAPEX WikiSQL model
+# 2) Model setup
 MODEL_ID = "microsoft/tapex-base-finetuned-wikisql"
+device   = 0 if torch.cuda.is_available() else -1
 
-# 3) Set device: GPU if available, else CPU
-device = 0 if torch.cuda.is_available() else -1
-
-# 4) Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model     = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
 
-# 5) Build the table-question-answering pipeline
 table_qa = pipeline(
     "table-question-answering",
     model=model,
@@ -27,18 +25,32 @@ table_qa = pipeline(
     device=device,
 )
 
-# 6) Define the QA function, casting all cells to strings to avoid float issues
+# 3) QA function with manual fallback for region‐based aggregations
 def answer_profitability(question: str) -> str:
-    # Cast entire DataFrame to string
+    q_lower = question.lower()
+
+    # 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())
+
+    # Otherwise, cast all cells to string and try TAPEX
     df_str = df.astype(str)
     table  = df_str.to_dict(orient="records")
+
     try:
         out = table_qa(table=table, query=question)
         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}"
 
-# 7) Define Gradio interface
+# 4) Gradio interface
 iface = gr.Interface(
     fn=answer_profitability,
     inputs=gr.Textbox(lines=2, placeholder="Ask a question about profitability…"),
@@ -46,10 +58,9 @@ iface = gr.Interface(
     title="SAP Profitability Q&A (TAPEX-Base)",
     description=(
         "Free-form questions on the synthetic profitability dataset, "
-        "powered end-to-end by microsoft/tapex-base-finetuned-wikisql."
+        "powered by microsoft/tapex-base-finetuned-wikisql with pandas fallbacks."
     )
 )
 
-# 8) Launch the app
 if __name__ == "__main__":
     iface.launch(server_name="0.0.0.0", server_port=7860)