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