app.py

"""
TableMind AI — Phase 7: Gradio app for Hugging Face Spaces (ZeroGPU-ready).

Run locally to test:  python app.py
Deploy: upload this file + rag_utils.py + requirements.txt + README.md to a new Space
(SDK: Gradio, Hardware: ZeroGPU if you have HF PRO, otherwise CPU basic).
"""

import os
import traceback

import gradio as gr
import spaces  # no-op outside a real ZeroGPU Space
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

import rag_utils

# ---------------------------------------------------------------------------
# Configuration — change MODEL_ID to your pushed merged model
# ---------------------------------------------------------------------------
MODEL_ID = "samandar1105/tablemind-qwen2.5-1.5b"  # <-- CHANGE THIS

SYSTEM_PROMPT = (
    "You are TableMind, a precise data analyst assistant. Answer strictly using the "
    "context provided below (retrieved table rows, document text, and/or a computed "
    "result). Give a complete, clearly-written, well-structured answer - not just a bare "
    "value - and reference the exact figures you used. If the answer is not contained in "
    "the context, say so plainly instead of guessing. Never invent numbers."
)

# ---------------------------------------------------------------------------
# Load model ONCE at module level (required by ZeroGPU: CUDA setup must
# happen outside any @spaces.GPU-decorated function).
# ---------------------------------------------------------------------------
print(f"Loading {MODEL_ID} ...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# Same reasoning as the fine-tuning script: bfloat16 needs Ampere+ (compute
# capability >= 8.0). ZeroGPU hardware (H200) supports it fine, but if you ever
# run this on a T4 (e.g. testing locally, or a non-ZeroGPU T4 Space), forcing
# bf16 is a real slowdown, not just a compatibility footnote.
if torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 8:
    COMPUTE_DTYPE = torch.bfloat16
else:
    COMPUTE_DTYPE = torch.float16

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=COMPUTE_DTYPE,
    device_map="auto" if torch.cuda.is_available() else None,
)
if torch.cuda.is_available():
    model = model.to("cuda")
print("Model loaded.")


# ---------------------------------------------------------------------------
# The only GPU-bound function — everything else (parsing, retrieval, pandas
# execution, chart drawing) stays on CPU so it doesn't consume ZeroGPU quota.
# ---------------------------------------------------------------------------
@spaces.GPU(duration=60)
def llm_generate(messages, max_new_tokens=500):
    inputs = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
    )
    inputs = inputs.to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            temperature=None,
            top_p=None,
            pad_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()


# ---------------------------------------------------------------------------
# Session state helpers
# ---------------------------------------------------------------------------

def process_upload(file, state):
    """Parse an uploaded file, build the RAG index, store in session state."""
    if file is None:
        return state, "No file uploaded yet - you can still ask general questions."
    try:
        parsed = rag_utils.parse_document(file.name)
        index, chunks = rag_utils.build_index(parsed["text_chunks"])
        state = {
            "dataframes": parsed["dataframes"],
            "index": index,
            "chunks": chunks,
        }
        n_tables = len(parsed["dataframes"])
        n_chunks = len(chunks)
        summary = f"Loaded **{os.path.basename(file.name)}** — {n_tables} table(s), {n_chunks} indexed chunk(s). Ask away!"
        return state, summary
    except Exception as e:
        traceback.print_exc()
        return state, f"Could not parse that file: {e}"


def answer_question(question: str, state: dict, include_chart: bool):
    state = state or {}
    dataframes = state.get("dataframes", {})
    index = state.get("index")
    chunks = state.get("chunks", [])

    context_parts = []
    computed_result = None

    # --- Numeric / aggregate path: LLM writes pandas, we execute it exactly ---
    if dataframes and rag_utils.is_numeric_question(question):
        schema = rag_utils.describe_dataframes(dataframes)
        code_prompt = [
            {"role": "system", "content": (
                "You write exactly one line of pandas code and nothing else. "
                "The dataframes are available as dfs['name']. Do not import anything. "
                "Do not explain. Output only the code line."
            )},
            {"role": "user", "content": f"{schema}\n\nQuestion: {question}\n\nPandas code:"},
        ]
        try:
            code_line = llm_generate(code_prompt, max_new_tokens=80)
            code_line = code_line.strip().strip("`").replace("python", "", 1).strip()
            result, err = rag_utils.run_pandas_query(code_line, dataframes)
            if err is None:
                computed_result = result
                context_parts.append(f"Computed result (via `{code_line}`): {result}")
        except Exception:
            traceback.print_exc()  # fall through to RAG path below

    # --- RAG path: retrieve relevant chunks ---
    if index is not None:
        retrieved = rag_utils.retrieve(question, index, chunks, k=6)
        if retrieved:
            context_parts.append("Retrieved context:\n" + "\n---\n".join(retrieved))

    if not context_parts:
        context_parts.append("No document has been uploaded, or nothing relevant was found. "
                              "Answer only if this is general knowledge you're confident about, "
                              "otherwise say you need a document to answer.")

    final_prompt = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"{chr(10).join(context_parts)}\n\nQuestion: {question}"},
    ]
    answer = llm_generate(final_prompt, max_new_tokens=500)

    chart_path = None
    if include_chart or rag_utils.wants_chart(question):
        try:
            import pandas as pd
            if isinstance(computed_result, (pd.Series, pd.DataFrame)):
                chart_path = rag_utils.make_chart(computed_result, title=question[:60])
        except Exception:
            traceback.print_exc()

    return answer, chart_path


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="TableMind AI — Table Q&A Chatbot", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 📊 TableMind AI
        Upload an Excel, CSV, PDF, or Word file (any size — it's chunked and indexed),
        then ask questions in plain English. Numeric/aggregate questions are answered with
        exact computed results, not guesses. Toggle the chart option for a visual.

        *Fine-tuned Qwen2.5-0.5B + RAG retrieval + a pandas execution layer for accuracy.*
        """
    )

    session_state = gr.State({})

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(label="📁 Upload a document", file_types=[".xlsx", ".xls", ".csv", ".pdf", ".docx", ".txt", ".md"])
            upload_status = gr.Markdown("No file uploaded yet.")
            chart_toggle = gr.Checkbox(label="Include a chart when relevant", value=True)
        with gr.Column(scale=2):
            question_box = gr.Textbox(label="Ask a question", lines=3, placeholder="e.g. What was the total revenue in Q3?")
            ask_btn = gr.Button("Ask TableMind", variant="primary")
            answer_box = gr.Markdown(label="Answer")
            chart_output = gr.Image(label="Chart", visible=True)

    file_input.change(fn=process_upload, inputs=[file_input, session_state], outputs=[session_state, upload_status])
    ask_btn.click(fn=answer_question, inputs=[question_box, session_state, chart_toggle], outputs=[answer_box, chart_output])
    question_box.submit(fn=answer_question, inputs=[question_box, session_state, chart_toggle], outputs=[answer_box, chart_output])

if __name__ == "__main__":
    demo.launch()