Update app.py

app.py

@@ -4,17 +4,17 @@ import torch
 import duckdb
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# 1) Load data into DuckDB
+# Load data into DuckDB
 df = pd.read_csv('synthetic_profit.csv')
 con = duckdb.connect(':memory:')
 con.register('sap', df)
 
-# 2) Build a one-line schema for prompting
-schema = ", ".join(df.columns)  # e.g. "Region,Product,FiscalYear,...."
+# One-line schema for prompts
+schema = ", ".join(df.columns)
 
-# 3) Load TAPEX-WikiSQL as a text2text generator
+# Load TAPEX for SQL generation
 MODEL_ID = "microsoft/tapex-base-finetuned-wikisql"
-device   = 0 if torch.cuda.is_available() else -1
+device = 0 if torch.cuda.is_available() else -1
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model     = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
@@ -28,43 +28,17 @@ sql_gen = pipeline(
     max_length=128,
 )
 
-# 4) Core QA fn: NL → SQL → execute → return result
 def answer_profitability(question: str) -> str:
-    # a) Prompt TAPEX to generate SQL
+    # 1) Ask TAPEX to write SQL
     prompt = (
-        f"-- Translate to SQL over table `sap` with columns ({schema})\n"
+        f"-- Translate to SQL for table `sap` ({schema})\n"
         f"Question: {question}\n"
         "SQL:"
     )
     sql = sql_gen(prompt)[0]['generated_text'].strip()
 
-    # b) Execute the generated SQL
+    # 2) Try to run it
     try:
-        result_df = con.execute(sql).df()
+        df_out = con.execute(sql).df()
     except Exception as e:
-        return f"❌ SQL Error: {e}\n\nGenerated SQL:\n{sql}"
-
-    # c) Format the output
-    if result_df.empty:
-        return f"No rows returned.\n\nGenerated SQL:\n{sql}"
-
-    # If it's a single cell result, just return that number
-    if result_df.shape == (1,1):
-        return str(result_df.iat[0,0])
-    # Otherwise pretty-print the DataFrame
-    return result_df.to_string(index=False)
-
-# 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 (SQL-Generation)",
-    description=(
-        "TAPEX converts your natural-language query into SQL,\n"
-        "then runs it via DuckDB—no hard-coded fallbacks."
-    )
-)
-
-if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+        return f"""**❌ SQL Error**