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