Update app.py

app.py

@@ -1,27 +1,30 @@
-# app.py
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import pandas as pd
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# Load your synthetic profitability dataset
+# 1) Load data
 df = pd.read_csv('synthetic_profit.csv')
 
-# Initialize the TAPEX small model fine-tuned on WikiSQL
-MODEL_ID = "microsoft/tapex-small-finetuned-wikisql"
+# 2) Use the publicly available TAPEX base WikiSQL model
+MODEL_ID = "microsoft/tapex-base-finetuned-wikisql"
+
+# 3) Ensure backend is available
+device = 0 if torch.cuda.is_available() else -1
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model     = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
 
-# Build a table-QA pipeline
 table_qa = pipeline(
     "table-question-answering",
     model=model,
     tokenizer=tokenizer,
     framework="pt",
-    device=-1   # set to 0 if you enable GPU in your Space
+    device=device,
 )
 
-def answer_profitability(question):
+# 4) QA function
+def answer_profitability(question: str) -> str:
     table = df.to_dict(orient="records")
     try:
         out = table_qa(table=table, query=question)
@@ -29,16 +32,13 @@ def answer_profitability(question):
     except Exception as e:
         return f"Error: {e}"
 
-# Gradio interface
+# 5) Gradio UI
 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-Small)",
-    description="""
-    Ask free-form questions on the synthetic profitability dataset.
-    Powered end-to-end by microsoft/tapex-small-finetuned-wikisql.
-    """
+    title="SAP Profitability Q&A (TAPEX-Base)",
+    description="Free-form questions on synthetic profitability data using microsoft/tapex-base-finetuned-wikisql."
 )
 
 if __name__ == "__main__":