Update app.py

app.py

@@ -1,4 +1,4 @@
-import os, io, json, math, pickle, textwrap, shutil, re
+import os, io, json, math, pickle, textwrap, shutil, re, zipfile, tempfile
 from typing import List, Dict, Any, Tuple
 import numpy as np, faiss, fitz  # pymupdf
 from tqdm import tqdm
@@ -6,17 +6,24 @@ import torch
 from sentence_transformers import SentenceTransformer
 import gradio as gr
 from groq import Groq
+from docx import Document
+from docx.shared import Pt
 
-# ---------- Config ----------
+# ---------- Branding ----------
+APP_NAME = "ScholarLens"
+TAGLINE = "Query your literature, get page-level proof"
+
+# ---------- Config (same engine, nicer UI) ----------
 EMBED_MODEL_NAME = "intfloat/multilingual-e5-small"
 CHUNK_SIZE = 1200
 CHUNK_OVERLAP = 200
-TOP_K_DEFAULT = 5
-MAX_CONTEXT_CHARS = 12000
+TOP_K_DEFAULT = 7
+MAX_CONTEXT_CHARS = 16000
 
 INDEX_PATH = "rag_index.faiss"
 STORE_PATH = "rag_store.pkl"
 
+# You can edit the default model here. All are selectable in the UI.
 MODEL_CHOICES = [
     "llama-3.3-70b-versatile",
     "llama-3.1-8b-instant",
@@ -96,6 +103,28 @@ def load_index() -> bool:
     return False
 
 # ---------- Ingest ----------
+def _collect_pdf_paths(upload_paths: List[str]) -> List[str]:
+    """Accept PDFs and ZIPs of PDFs."""
+    if not upload_paths:
+        return []
+    out = []
+    for p in upload_paths:
+        p = str(p)
+        if p.lower().endswith(".pdf"):
+            out.append(p)
+        elif p.lower().endswith(".zip"):
+            tmpdir = tempfile.mkdtemp(prefix="pdfs_")
+            with zipfile.ZipFile(p, "r") as z:
+                for name in z.namelist():
+                    if name.lower().endswith(".pdf"):
+                        z.extract(name, tmpdir)
+            # collect extracted PDFs
+            for root, _, files in os.walk(tmpdir):
+                for f in files:
+                    if f.lower().endswith(".pdf"):
+                        out.append(os.path.join(root, f))
+    return out
+
 def ingest_pdfs(paths: List[str]) -> Tuple[Any, List[Dict[str, Any]]]:
     entries: List[Dict[str, Any]] = []
     for pdf in tqdm(paths, total=len(paths), desc="Parsing PDFs"):
@@ -122,7 +151,7 @@ def ingest_pdfs(paths: List[str]) -> Tuple[Any, List[Dict[str, Any]]]:
     index = build_faiss(embs)
     return index, entries
 
-# ---------- Retrieval (supports required keywords) ----------
+# ---------- Retrieval with optional keyword filter ----------
 def retrieve(query: str, top_k=5, must_contain: str = ""):
     global faiss_index, docstore
     if faiss_index is None or not docstore:
@@ -153,10 +182,10 @@ def retrieve(query: str, top_k=5, must_contain: str = ""):
     return hits
 
 # ---------- Groq LLM ----------
-def groq_answer(query: str, contexts, model_name="llama-3.1-70b-versatile", temperature=0.2, max_tokens=1000):
+def groq_answer(query: str, contexts, model_name="llama-3.3-70b-versatile", temperature=0.2, max_tokens=1000):
     try:
         if not os.environ.get("GROQ_API_KEY"):
-            return "GROQ_API_KEY is not set. Add it in your host's environment/secrets."
+            return "GROQ_API_KEY is not set. Add it in your Space secrets or the key box."
         client = Groq(api_key=os.environ["GROQ_API_KEY"])
 
         packed, used = [], 0
@@ -189,99 +218,235 @@ def groq_answer(query: str, contexts, model_name="llama-3.1-70b-versatile", temp
         import traceback
         return f"Groq API error: {e}\n```\n{traceback.format_exc()}\n```"
 
-# ---------- Helpers for UI ----------
+# ---------- Export helpers ----------
+def export_answer_to_docx(question: str, answer_md: str, rows: List[List[str]]) -> str:
+    """
+    Save Q&A with sources table to a .docx and return path.
+    rows: [Source, Page, Score, Snippet]
+    """
+    doc = Document()
+    styles = doc.styles
+    try:
+        styles['Normal'].font.name = 'Calibri'
+        styles['Normal'].font.size = Pt(11)
+    except Exception:
+        pass
+
+    doc.add_heading(f"{APP_NAME} - Answer", level=1)
+    doc.add_paragraph(f"Question: {question}")
+
+    doc.add_heading("Answer", level=2)
+    # Write as plain text to keep it simple in Word
+    for line in answer_md.splitlines():
+        doc.add_paragraph(line)
+
+    doc.add_heading("References (Top Passages)", level=2)
+    table = doc.add_table(rows=1, cols=4)
+    hdr = table.rows[0].cells
+    hdr[0].text = "Source"
+    hdr[1].text = "Page"
+    hdr[2].text = "Score"
+    hdr[3].text = "Snippet"
+    for r in rows:
+        row = table.add_row().cells
+        for i, val in enumerate(r):
+            row[i].text = str(val)
+
+    path = "scholarlens_answer.docx"
+    doc.save(path)
+    return path
+
+# ---------- UI helpers ----------
 def build_index_from_uploads(paths: List[str]) -> str:
     global faiss_index, docstore
-    if not paths: return "Please upload at least one PDF."
-    if len(paths) > 120: return "Please limit to ~100 PDFs per build."
-
-    faiss_index, entries = ingest_pdfs(paths)
+    pdfs = _collect_pdf_paths(paths)
+    if not pdfs:
+        return "Please upload at least one PDF or ZIP of PDFs."
+    faiss_index, entries = ingest_pdfs(pdfs)
     save_index(faiss_index, entries)
     docstore = entries
-    return f"Index built with {len(entries)} chunks from {len(paths)} PDFs. Saved to disk."
+    return f"✅ Index built with {len(entries)} chunks from {len(pdfs)} files. You can start asking questions."
 
 def reload_index() -> str:
     ok = load_index()
-    return f"Index reloaded. Chunks: {len(docstore)}" if ok else "No saved index found."
+    return f"🔁 Index reloaded. Chunks ready: {len(docstore)}" if ok else "No saved index found yet."
 
-def ask_rag(query: str, top_k, model_name: str, temperature: float, must_contain: str):
+def ask_rag(question: str, top_k, model_name: str, temperature: float, must_contain: str):
     try:
-        if not query.strip():
-            return "Please enter a question.", []
-        ctx = retrieve(query, top_k=int(top_k) if top_k else TOP_K_DEFAULT, must_contain=must_contain)
-        ans = groq_answer(query, ctx, model_name=model_name, temperature=temperature)
+        if not question.strip():
+            return "Please enter a question.", [], "", gr.update(visible=False)
+        ctx = retrieve(question, top_k=int(top_k) if top_k else TOP_K_DEFAULT, must_contain=must_contain)
+        ans = groq_answer(question, ctx, model_name=model_name, temperature=temperature)
+
+        # sources table
         rows = []
         for c in ctx:
             preview = c["text"][:200].replace("\n"," ") + ("..." if len(c["text"])>200 else "")
             rows.append([c["source"], str(c["page_start"]), f"{c['score']:.3f}", preview])
-        return ans, rows
+
+        # snippets pretty print
+        details = []
+        for c in ctx:
+            details.append(f"**{c['source']} p.{c['page_start']}**\n> {c['text'].strip()[:1000]}")
+        snippets_md = "\n\n---\n\n".join(details)
+
+        download_btn = gr.update(visible=True)
+        return ans, rows, snippets_md, download_btn
     except Exception as e:
         import traceback
-        return f"**Error:** {e}\n```\n{traceback.format_exc()}\n```", []
+        err = f"**Error:** {e}\n```\n{traceback.format_exc()}\n```"
+        return err, [], "", gr.update(visible=False)
 
 def set_api_key(k: str):
     if k and k.strip():
         os.environ["GROQ_API_KEY"] = k.strip()
-        return "API key set in runtime."
+        return "🔑 API key set for this session."
     return "No key provided."
 
 def download_index_zip():
     if not (os.path.exists(INDEX_PATH) and os.path.exists(STORE_PATH)):
         return None
-    base = "rag_index_bundle"
-    zip_path = shutil.make_archive(base, "zip", ".", ".")
-    # workaround for shutil: package explicit files
-    with shutil.make_archive("rag_index", "zip"):
-        pass
-    # build our own zip containing only index files
-    import zipfile
     zp = "rag_index_bundle.zip"
     with zipfile.ZipFile(zp, "w", zipfile.ZIP_DEFLATED) as z:
         z.write(INDEX_PATH)
         z.write(STORE_PATH)
     return zp
 
+def do_export_docx(question, answer_md, sources_rows):
+    if not answer_md or not sources_rows:
+        return None
+    try:
+        path = export_answer_to_docx(question, answer_md, sources_rows)
+        return path
+    except Exception:
+        return None
+
+# ---------- Theme ----------
+theme = gr.themes.Soft(
+    primary_hue="indigo",
+    secondary_hue="blue",
+    neutral_hue="slate",
+).set(
+    body_background_fill="#0B1220",  # dark-friendly hero
+    block_background_fill="#0F172A",
+    block_shadow="*shadow-lg",
+    radius_size="8px",
+)
+
 # ---------- Gradio UI ----------
-with gr.Blocks(title="RAG over PDFs (Groq)") as demo:
-    gr.Markdown("## RAG over your PDFs using Groq\nUpload PDFs, build an index, then ask questions with cited answers.")
+with gr.Blocks(title=f"{APP_NAME} | RAG over PDFs", theme=theme, css="""
+#hero {
+  background: radial-gradient(1200px 600px at 20% -10%, rgba(99,102,241,.25), transparent),
+              radial-gradient(1000px 500px at 120% 10%, rgba(14,165,233,.20), transparent);
+  border: 1px solid rgba(99,102,241,.20);
+}
+.kpi {text-align:center;padding:12px;border-radius:10px;border:1px solid rgba(255,255,255,.08);}
+.footer {opacity:.8;}
+""") as demo:
+    # --- Header / Hero ---
+    with gr.Group(elem_id="hero"):
+        gr.Markdown(
+            f"""
+<div style="display:flex;align-items:center;gap:16px;">
+  <div style="font-size:36px">📚🔎 <b>{APP_NAME}</b></div>
+  <div style="opacity:.9;">{TAGLINE}</div>
+</div>
+<p style="opacity:.85;margin-top:6px;">
+Upload your papers, build an index, and ask research questions with verifiable, page-level citations.
+</p>
+""")
+
+    # --- KPI row ---
+    with gr.Row():
+        gr.Markdown("**Meaning-aware retrieval**<br><span class='kpi'>E5 + FAISS</span>", elem_classes=["kpi"])
+        gr.Markdown("**Cited answers**<br><span class='kpi'>Page-level proof</span>", elem_classes=["kpi"])
+        gr.Markdown("**Runs anywhere**<br><span class='kpi'>HF Spaces or Colab</span>", elem_classes=["kpi"])
+
+    # --- Key / Settings ---
     with gr.Row():
-        api_box = gr.Textbox(label="(Optional) Set GROQ_API_KEY for this session", type="password", placeholder="sk_...") 
+        api_box = gr.Textbox(label="(Optional) Set GROQ_API_KEY", type="password", placeholder="sk_...")
         set_btn = gr.Button("Set Key")
         set_out = gr.Markdown()
-        set_btn.click(set_api_key, inputs=[api_box], outputs=[set_out])
-
-    with gr.Tab("1) Build or Load Index"):
-        file_u = gr.Files(label="Upload PDFs", file_types=[".pdf"], type="filepath")
-        with gr.Row():
-            build_btn = gr.Button("Build Index")
-            reload_btn = gr.Button("Reload Saved Index")
-            download_btn = gr.Button("Download Index (.zip)")
-        build_out = gr.Markdown()
-
-        def on_build(paths, progress=gr.Progress(track_tqdm=True)):
-            try:
-                return build_index_from_uploads(paths)
-            except Exception as e:
-                import traceback
-                return f"**Error while building index:** {e}\n\n```\n{traceback.format_exc()}\n```"
-
-        build_btn.click(on_build, inputs=[file_u], outputs=[build_out])
-        reload_btn.click(fn=reload_index, outputs=[build_out])
-        zpath = gr.File(label="Index zip", interactive=False)
-        download_btn.click(fn=download_index_zip, outputs=[zpath])
-
-    with gr.Tab("2) Ask Questions"):
-        q = gr.Textbox(label="Your question", lines=2, placeholder="Ask something present in the uploaded papers…")
-        with gr.Row():
-            topk = gr.Slider(1, 15, value=TOP_K_DEFAULT, step=1, label="Top-K passages")
-            model_dd = gr.Dropdown(MODEL_CHOICES, value=MODEL_CHOICES[0], label="Groq model")
-            temp = gr.Slider(0.0, 1.0, value=0.2, step=0.05, label="Temperature")
-        must = gr.Textbox(label="Must contain (comma-separated keywords)", placeholder="camera, CMOS, frame rate")
-        ask_btn = gr.Button("Answer")
-        ans = gr.Markdown()
-        src = gr.Dataframe(headers=["Source","Page","Score","Snippet"], wrap=True)
-        ask_btn.click(ask_rag, inputs=[q, topk, model_dd, temp, must], outputs=[ans, src])
-
-demo.queue()  # keep it simple for broad Gradio versions
+    set_btn.click(set_api_key, inputs=[api_box], outputs=[set_out])
+
+    with gr.Tabs():
+        # ---------------- Tab 1: Build / Load ----------------
+        with gr.Tab("1) Build or Load Index"):
+            gr.Markdown("Upload PDFs or a ZIP of PDFs, then click **Build Index**.")
+            file_u = gr.Files(label="Upload PDFs or ZIP", file_types=[".pdf", ".zip"], type="filepath")
+            with gr.Row():
+                build_btn = gr.Button("Build Index", variant="primary")
+                reload_btn = gr.Button("Reload Saved Index")
+                download_btn = gr.Button("Download Index (.zip)")
+            build_out = gr.Markdown()
+
+            def on_build(paths, progress=gr.Progress(track_tqdm=True)):
+                try:
+                    return build_index_from_uploads(paths)
+                except Exception as e:
+                    import traceback
+                    return f"**Error while building index:** {e}\n\n```\n{traceback.format_exc()}\n```"
+
+            build_btn.click(on_build, inputs=[file_u], outputs=[build_out])
+            reload_btn.click(fn=reload_index, outputs=[build_out])
+            zpath = gr.File(label="Index bundle", interactive=False)
+            download_btn.click(fn=download_index_zip, outputs=[zpath])
+
+        # ---------------- Tab 2: Ask ----------------
+        with gr.Tab("2) Ask Questions"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    q = gr.Textbox(label="Your question", lines=3, placeholder="e.g., Compare GTAW experimental parameters with citations")
+                    must = gr.Textbox(label="Must contain (comma-separated keywords)", placeholder="camera, CMOS, frame rate")
+                    with gr.Accordion("Advanced settings", open=False):
+                        topk = gr.Slider(1, 20, value=TOP_K_DEFAULT, step=1, label="Top-K passages")
+                        model_dd = gr.Dropdown(MODEL_CHOICES, value=MODEL_CHOICES[0], label="Groq model")
+                        temp = gr.Slider(0.0, 1.0, value=0.2, step=0.05, label="Temperature")
+                    with gr.Row():
+                        ask_btn = gr.Button("Answer", variant="primary")
+                        clear_btn = gr.Button("Clear")
+
+                    gr.Examples(
+                        examples=[
+                            ["List camera model, sensor type, resolution, and FPS across studies. Cite pages.", "camera, fps, resolution"],
+                            ["Extract limitations and future work across the corpus, with page references.", ""],
+                            ["Compare GTAW setups: current range, travel speed, torch standoff, sensors.", "GTAW, current, speed, torch"],
+                            ["Summarize the main results tables with metrics and page citations.", "table, accuracy, mAP, F1"]
+                        ],
+                        inputs=[q, must],
+                        label="Quick examples",
+                    )
+                with gr.Column(scale=1.4):
+                    ans = gr.Markdown(label="Answer", show_label=False)
+                    src = gr.Dataframe(headers=["Source","Page","Score","Snippet"], wrap=True, label="Top passages")
+                    with gr.Accordion("Show retrieved snippets", open=False):
+                        snippets_md = gr.Markdown("")
+                    with gr.Row():
+                        export_btn = gr.Button("Export Answer to DOCX", visible=False)
+                        exported = gr.File(label="Download answer", visible=True)
+
+            # wire buttons
+            ask_btn.click(fn=ask_rag, inputs=[q, topk, model_dd, temp, must], outputs=[ans, src, snippets_md, export_btn])
+            export_btn.click(fn=do_export_docx, inputs=[q, ans, src], outputs=[exported])
+            clear_btn.click(lambda: ("", [], "", gr.update(visible=False)), outputs=[ans, src, snippets_md, export_btn])
+
+        # ---------------- Tab 3: About ----------------
+        with gr.Tab("About"):
+            gr.Markdown(
+                """
+**ScholarLens** helps researchers move from reading to results with answers grounded in the papers you upload.
+
+- Meaning-aware retrieval (E5 + FAISS)
+- Answers limited to your corpus, with page-level citations
+- Optional keyword filter to stay on topic
+- Runs on Hugging Face Spaces or Google Colab
+- Powered by Groq models
+
+*Privacy note:* your files stay on this Space. Only the Groq call is external.
+                """
+            )
+
+# broad compatibility for Spaces
+demo.queue()
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))