Praneethaneelapareddigari/visual-rag

1# app.py2import os3import json4import time5import tempfile6import streamlit as st7# --- Hugging Face Spaces / Linux runtime helpers ---8import shutil9import pytesseract10 11# Ensure Tesseract and Poppler are discoverable in the Space12pytesseract.pytesseract.tesseract_cmd = shutil.which("tesseract") or "/usr/bin/tesseract"13os.environ["PATH"] = os.environ.get("PATH", "") + ":/usr/bin:/usr/local/bin"14 15# Prefer a public, lightweight embedding model to avoid auth/gated downloads16# (If your rag_pipeline.py already sets this, you can remove the next 4 lines.)17try:18    from sentence_transformers import SentenceTransformer  # noqa: F40119    os.environ.setdefault("EMB_MODEL_NAME", "sentence-transformers/all-MiniLM-L6-v2")20except Exception:21    pass22 23 24 25from doc_loader import load_document26from figure_extractor import extract_figures27from rag_pipeline import (28    build_rag_pipeline,29    query_rag_full,30    evaluate_rag,31)32 33st.set_page_config(page_title="📑 Visual Document RAG", layout="wide")34st.title("📑 Visual Document RAG ")35 36domain = st.selectbox(37    "Domain focus",38    ["Finance", "Healthcare", "Law", "Education", "Multimodal"],39    index=0,40    help="Used to lightly steer retrieval/answers",41)42 43uploaded_file = st.file_uploader("📂 Upload a PDF or Image", type=["pdf", "png", "jpg"])44query = st.text_input("🔎 Ask a question about the document:")45summarize = st.checkbox("📝 Summarize the whole document")46 47if uploaded_file:48    # persist to temp path for libs49    with tempfile.NamedTemporaryFile(50        delete=False, suffix=f".{uploaded_file.name.split('.')[-1]}"51    ) as tmp:52        tmp.write(uploaded_file.read())53        tmp_path = tmp.name54 55    with st.spinner("Processing document..."):56        docs_text, sections = load_document(tmp_path, return_sections=True)57 58    # Optional figure extraction (PDF only)59    figures_meta = []60    extra_docs = []61    if uploaded_file.type == "application/pdf":62        try:63            figures_meta = extract_figures(tmp_path, out_dir="figures", lang="eng") or []64            if figures_meta:65                sections["Figures (OCR+captions)"] = "\n\n".join(66                    [67                        f"p.{f['page']} — {f.get('caption') or '(no caption)'}\n"68                        f"{(f.get('ocr_text') or '')[:200]}"69                        for f in figures_meta70                    ]71                )72                # vectorizable figure docs73                for f in figures_meta:74                    content = (75                        f"FIGURE p.{f['page']}: {f.get('caption') or ''}\n"76                        f"OCR: {f.get('ocr_text') or ''}\n"77                        f"TAGS: {' '.join(f.get('tags', []))}"78                    ).strip()79                    metadata = {80                        "type": "figure",81                        "page": f.get("page"),82                        "path": f.get("path"),83                        "caption": f.get("caption"),84                        "tags": f.get("tags", []),85                    }86                    extra_docs.append({"content": content, "metadata": metadata})87        except Exception as e:88            st.warning(f"Figure extraction skipped: {e}")89 90    # Build vector index (now indexes figures too)91    if docs_text.strip():92        db = build_rag_pipeline(docs_text, extra_docs=extra_docs or None)93        st.success(f"✅ Document indexed! (Domain: {domain})")94 95        # Show extracted sections96        with st.expander("📂 Extracted Document Content"):97            tab_names = list(sections.keys())98            tabs = st.tabs(tab_names)99            for i, name in enumerate(tab_names):100                with tabs[i]:101                    st.text_area(f"{name}", sections[name], height=230)102 103        # Actions (Summarize or Q&A)104        answer_text, retrieved_docs, latency = None, None, None105        colA, colB = st.columns([1, 1])106 107        with colA:108            if summarize:109                if st.button("📝 Summarize Document"):110                    start = time.time()111                    q = f"[Domain: {domain}] Summarize this document briefly with key points and numbers."112                    answer_text, _, retrieved_docs = query_rag_full(113                        db, q, domain=domain114                    )115                    latency = round(time.time() - start, 3)116            else:117                if query and st.button("💡 Get Answer"):118                    start = time.time()119                    q = f"[Domain: {domain}] {query}"120                    answer_text, _, retrieved_docs = query_rag_full(121                        db, q, domain=domain122                    )123                    latency = round(time.time() - start, 3)124 125        # Show results126        if answer_text is not None:127            st.subheader("💡 Answer")128            st.write(answer_text)129            st.caption(f"⏱️ Latency: {latency}s")130 131            with st.expander("🔍 Retrieved Contexts"):132                if retrieved_docs:133                    for i, d in enumerate(retrieved_docs, 1):134                        meta = getattr(d, "metadata", {}) or {}135                        if meta.get("type") == "figure" and meta.get("path"):136                            st.write(137                                f"Figure (page {meta.get('page')}): "138                                f"{meta.get('caption') or '(no caption)'}"139                            )140                            st.image(meta["path"], use_container_width=True)141                        else:142                            st.info(f"Chunk {i}:\n\n{d.page_content}")143 144            # Evaluation145            st.markdown("---")146            st.subheader("📊 Evaluation")147            if st.button("Evaluate (LLM-based: Faithfulness & Relevancy)"):148                try:149                    raw = evaluate_rag(150                        answer_text,151                        [d.page_content for d in (retrieved_docs or [])],152                        f"[Domain: {domain}] {query or 'Summary'}",153                    )154                    try:155                        payload = json.loads(raw) if isinstance(raw, str) else raw156                    except Exception:157                        payload = {"raw": raw}158                    st.json(payload)159                    st.download_button(160                        "⬇️ Download Evaluation JSON",161                        data=json.dumps(payload, indent=2),162                        file_name="evaluation.json",163                        mime="application/json",164                    )165                except Exception as e:166                    st.warning(f"Evaluation unavailable: {e}")167    else:168        st.error("❌ No text could be extracted from the uploaded file.")169