app.py

"""
TableMind AI — Gradio app for Hugging Face Spaces (CPU Space).

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: CPU basic).

NOTE: this is the CPU-only version - no ZeroGPU/@spaces.GPU code. Generation will be
noticeably slower than on a GPU (expect several seconds to tens of seconds per answer
depending on model size and question length), but it sidesteps ZeroGPU's quota limits,
worker allocation failures, and PRO-subscription requirement entirely.
"""

import os
import traceback

import gradio as gr
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.
# ---------------------------------------------------------------------------
# HF_TOKEN: set this as a Space secret (Settings -> Repository secrets) if your
# model repo is private. Harmless to leave unset if the repo is public.
HF_TOKEN = os.environ.get("HF_TOKEN")

print(f"Loading {MODEL_ID} ...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)

# On CPU, use float32 - NOT float16/bfloat16. Unlike CUDA, PyTorch's CPU backend
# has patchy/slow support for half-precision ops on many kernels, so forcing fp16
# or bf16 here would risk both errors and worse performance, not better. bf16 is
# only the right choice when there's an actual CUDA device with Ampere+ hardware
# support (kept here as a conditional for portability, e.g. if you later switch
# this Space back to GPU hardware - on CPU it will always fall through to float32).
if torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 8:
    COMPUTE_DTYPE = torch.bfloat16
elif torch.cuda.is_available():
    COMPUTE_DTYPE = torch.float16
else:
    COMPUTE_DTYPE = torch.float32

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=COMPUTE_DTYPE,
    device_map="auto" if torch.cuda.is_available() else None,
    token=HF_TOKEN,
)
if torch.cuda.is_available():
    model = model.to("cuda")
model.eval()
print(f"Model loaded on {'cuda' if torch.cuda.is_available() else 'cpu'} with dtype {COMPUTE_DTYPE}.")


# ---------------------------------------------------------------------------
# Generation function - plain function, no GPU decorator needed on a CPU Space.
# ---------------------------------------------------------------------------
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()
            print(f"[DEBUG] numeric question detected. Generated code: {code_line!r}")
            result, err = rag_utils.run_pandas_query(code_line, dataframes)
            print(f"[DEBUG] execution result={result!r} err={err!r}")
            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)
        print(f"[DEBUG] retrieved {len(retrieved)} chunks (index has {len(chunks)} total)")
        if retrieved:
            context_parts.append("Retrieved context:\n" + "\n---\n".join(retrieved))
    else:
        print("[DEBUG] no index in session state - was a file uploaded this session?")

    if not context_parts:
        print("[DEBUG] context_parts is EMPTY - model will answer with no grounding at all")
        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}"},
    ]
    print(f"[DEBUG] final prompt context (first 500 chars): {context_parts[0][:500] if context_parts else 'NONE'}")
    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-1.5B + RAG retrieval + a pandas execution layer for accuracy.*
        *Running on CPU — answers may take a while, especially longer ones.*
        """
    )

    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()