Publish MedDiscover-HF app

NOTE.md

@@ -0,0 +1,15 @@
+# MedDiscover-HF status (Nov 20)
+
+- Added Gradio app scaffold (`app.py`) for Hugging Face Spaces (ZeroGPU):
+  - MedCPT embeddings (`ncbi/MedCPT-*`) + FAISS IP index stored in `/data`.
+  - PDF upload → text extraction → chunking (500/50) → embedding → index build/load.
+  - Generator dropdown using open-weight HF models: `openai/gpt-oss-20b`, `google/gemma-3-12b-it`, `deepseek-ai/DeepSeek-VL2-small`, `ibm-granite/granite-3.1-8b-instruct`, `ibm-granite/granite-docling-258M`.
+  - Streaming answers; prompt forces context-grounded responses.
+  - Uses `@spaces.GPU()` on heavy steps; caches/models under `/data/.cache` (`HF_HOME`).
+  - No reranker; straight FAISS retrieval.
+- Dependencies captured in `requirements.txt`; Spaces front matter in `README.md`.
+- Pushed to HF Space `VatsalPatel18/MedDisover-space` (commit f6d2d7d).
+- To do next when resuming:
+  1) Set `HF_TOKEN` secret in the Space if Granite/Docling needs auth.
+  2) Let the Space rebuild/install; test with a light model first (granite docling/3.1-8b) on ZeroGPU.
+  3) Adjust model IDs if your tested demos use different ones; tune `max_new_tokens` if OOM. 

PLAN.md

@@ -0,0 +1,70 @@
+# MedDiscover-HF Rollout Plan (Hugging Face Spaces, ZeroGPU)
+
+## 1) Goals & Constraints
+- Host MedDiscover as a Gradio Space using ZeroGPU; no external API models (OpenAI, etc.).
+- Provide a model dropdown with OSS models tested in other Spaces: `openai/gpt-oss-20b`, `google/gemma-3-12b-it`, `deepseek-vl2-small`, `ibm-granite/granite-vision` family, `ibm-granite/granite-docling-258M`, plus room for additions.
+- Keep existing pipeline: PDF ingest → chunking (MedCPT tokenizer/encoder) → FAISS retrieval → answer generation via selected OSS model.
+- Must run in HF build/runtime: single `app.py` entry (or `src/app.py` with `app_file` set), `sdk: gradio`, dependencies via `requirements.txt`/`pyproject`. Persistent storage limited; assume `/data` volume for cached indices.
+
+## 2) Spaces-specific mechanics to carry over
+- Use `@spaces.GPU()` (as in gpt-oss demo) to request ZeroGPU for heavy calls; pair with `device_map="auto"` and `torch_dtype="auto"/bfloat16` to fit managed GPUs.
+- For IBM Granite Docling/Vision: models require `use_auth_token=True` and `trust_remote_code` in some cases; respect `gr.NO_RELOAD` guard to avoid reloading on hot-reload.
+- Streaming: use `TextIteratorStreamer` threaded generation (seen in gpt-oss, gemma) to keep UI responsive.
+- System/developer prompts: gpt-oss demo uses Harmony encoding/preprompt parsing; we can simplify to plain chat unless Harmony is desired. If kept, include `openai_harmony` dependency and message rendering.
+- Media handling: gemma demo enforces image/video limits and uses `<image>` tag counting; docling/granite demos load sample assets and draw boxes—good references for multimodal support.
+- README front matter (`title`, `sdk: gradio`, `sdk_version`, `app_file`) required for Spaces config; ZeroGPU Spaces accept the same.
+
+## 3) Architecture on HF
+- Single `app.py` (or `src/app.py`) hosting:
+  - Model registry: map human-facing names → loader functions/config (pipeline/AutoModel, tokenizer/processor, dtype, device_map, chat template).
+  - Embedding and indexing: MedCPT article encoder (requires GPU). Build FAISS index into `/data/faiss_index.bin` with metadata `/data/doc_metadata.json`; reuse across sessions if present.
+  - Retrieval: embed query with the same encoder, FAISS search (IP for MedCPT), optional rerank (if cross-encoder feasible on available GPU; otherwise skip).
+  - Generation: streaming handler wrapping selected model; minimal prompt template: “Use only retrieved context; answer concisely.”
+  - Gradio UI: upload PDFs, process PDFs (chunk+index), dropdown for embedding model (MedCPT only here) and generator model (OSS list), sliders for k/max_tokens/temp/etc., chat box showing answer + context.
+- Persistence: point caches/indices to `/data` (Space persistent storage). Handle cold start by checking disk before loading.
+
+## 4) Model integration plan (technical nuances)
+- `openai/gpt-oss-20b`:
+  - Load via `pipeline("text-generation", trust_remote_code=True, device_map="auto", torch_dtype="auto")`.
+  - Optionally keep Harmony encoding & `@spaces.GPU()` wrapper for generation; streaming with `TextIteratorStreamer`.
+- `google/gemma-3-12b-it`:
+  - `AutoProcessor.from_pretrained(..., padding_side="left")`, `Gemma3ForConditionalGeneration.from_pretrained(..., device_map="auto", torch_dtype=torch.bfloat16, attn_implementation="eager")`.
+  - Supports images/videos with `<image>` tags; enforce file count limits if multimodal chat is desired.
+- `deepseek-vl2-small` (assumed similar VLM):
+  - Inspect loader; likely `AutoProcessor` + vision-language model. Use `device_map="auto"`, `torch_dtype=torch.bfloat16`, and streaming generation.
+- `ibm-granite/granite-vision*`:
+  - Vision-language; may need `trust_remote_code`; keep sample image handling (resize/pad) patterns from demo.
+- `ibm-granite/granite-docling-258M`:
+  - Uses `AutoProcessor` + `Idefics3ForConditionalGeneration`; requires `use_auth_token=True`, `torch_dtype=torch.bfloat16`, `device_map=device`.
+  - Handles DocTags markup and bounding boxes; we can limit to text answers (no drawing) to reduce complexity initially.
+- For all: wrap generation in a common interface that accepts (query, retrieved_chunks, params) → streamed text. Unify token stopping; cap max_new_tokens modestly to stay within ZeroGPU limits.
+
+## 5) Retrieval/Chunking specifics
+- Chunking: reuse existing 500/50 word overlap; keep MedCPT encoder (article + query) loaded once behind `@spaces.GPU()` guard. Ensure tokenizer truncation to MAX lengths (as in current config).
+- FAISS: build `IndexFlatIP` for MedCPT (768-d). Convert embeddings to float32 before add. Store index/metadata under `/data`.
+- Metadata structure: `[{"doc_id": int, "filename": str, "chunk_id": int, "text": str}]`.
+- Query flow: embed query → FAISS search k → concatenate top-k texts into prompt. Skip rerank if cross-encoder too heavy for ZeroGPU; optional flag if GPU allows.
+
+## 6) Hugging Face build/runtime
+- Files:
+  - `app.py` (entry), `requirements.txt` (gradio, transformers, faiss-cpu, torch==2.3+cpu?; on ZeroGPU torch with CUDA is preinstalled), optional `runtime.txt` (e.g., `python-3.11`), `README.md` with HF front matter, small sample PDFs.
+  - If MedCPT or Granite needs auth, set `HF_TOKEN` secret; use `use_auth_token=True` where required.
+- Caching: set `HF_HOME=/data/.cache` to persist model weights between restarts.
+- No Docker needed unless custom system deps; prefer pure pip for Spaces.
+
+## 7) UI/UX sketch
+- Left rail: API/token textbox (for private models if needed), PDF upload & process button, model dropdown (generator), k slider, max_tokens slider, temperature/top_p, rerank checkbox (if available).
+- Right: chat box with answer; collapsible context display; status messages (index loaded/building).
+- Streaming answers with stop button; show retrieval scores optionally.
+
+## 8) Testing & rollout
+- Dry-run locally on CPU with smallest model (e.g., granite-docling-258M) if possible; otherwise rely on HF logs.
+- Deploy to a test Space (ZeroGPU), confirm model loads, index builds on uploaded PDFs, and chat returns grounded answers.
+- Measure cold-start load times per model; consider pre-pinning a default lightweight model to speed startup.
+
+## 9) Next steps to implement
+1) Scaffold `MedDiscover-HF/app.py` with common loaders, retrieval pipeline, Gradio UI, and `/data` persistence.
+2) Add `requirements.txt` mirroring demos (`gradio>=5.x`, `transformers`, `faiss-cpu`, `torch`, `spaces`, model-specific libs like `openai-harmony`, `opencv-python` if doing video).
+3) Wire model registry for gpt-oss-20b, gemma-3-12b-it, deepseek-vl2-small, granite-vision, granite-docling; test generation stubs with mock context.
+4) Integrate MedCPT embedding + FAISS; add index build/load actions and guard GPU usage with `@spaces.GPU()`.
+5) Push to HF Space; validate ZeroGPU compatibility and memory footprints; tune max token defaults per model.

README.md

@@ -0,0 +1,21 @@
+---
+title: MedDiscover-HF
+emoji: 🩺
+colorFrom: purple
+colorTo: blue
+sdk: gradio
+sdk_version: 5.44.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+short_description: MedDiscover RAG with OSS HF models on ZeroGPU
+---
+
+## MedDiscover-HF
+A Hugging Face Spaces-ready Gradio app that runs MedDiscover-style RAG (MedCPT embeddings + FAISS) with open-weight HF generator models (gpt-oss-20b, gemma-3-12b-it, deepseek-vl2-small, granite vision/docling variants) on ZeroGPU.
+
+### Notes
+- Uses `/data` for cached FAISS index/metadata and HF model cache (`HF_HOME=/data/.cache`).
+- No external API keys required; optionally use `HF_TOKEN` secret if models need auth.
+- Chunking: 500 words with 50-word overlap; MedCPT encoder for embeddings; FAISS IP index.
+- Generation is context-grounded; drop-down selects generator model.

app.py

@@ -0,0 +1,372 @@
+"""
+MedDiscover-HF: Hugging Face Spaces-ready Gradio app.
+ - ZeroGPU-compatible (uses @spaces.GPU for heavy ops)
+ - MedCPT embeddings + FAISS retrieval over uploaded PDFs
+ - OSS generator model dropdown (gpt-oss-20b, gemma-3-12b-it, deepseek-vl2-small, granite vision/docling)
+"""
+
+import os
+import json
+from pathlib import Path
+from threading import Thread
+from typing import List, Dict, Tuple
+
+import faiss
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+from PyPDF2 import PdfReader
+from transformers import (
+    AutoModel,
+    AutoTokenizer,
+    AutoProcessor,
+    TextIteratorStreamer,
+    pipeline,
+)
+
+# ----------------------------
+# Paths and env configuration
+# ----------------------------
+BASE_DIR = Path(__file__).parent
+DATA_DIR = Path(os.getenv("DATA_DIR", "/data"))
+DATA_DIR.mkdir(parents=True, exist_ok=True)
+INDEX_PATH = DATA_DIR / "faiss_index.bin"
+META_PATH = DATA_DIR / "doc_metadata.json"
+HF_TOKEN = os.getenv("HF_TOKEN")  # set in Space secrets if needed
+HF_HOME = os.getenv("HF_HOME", str(DATA_DIR / ".cache"))
+os.environ["HF_HOME"] = HF_HOME
+
+# ----------------------------
+# Chunking / PDF utils
+# ----------------------------
+CHUNK_SIZE = 500
+OVERLAP = 50
+
+
+def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = OVERLAP) -> List[str]:
+    words = text.split()
+    chunks = []
+    start = 0
+    while start < len(words):
+        end = start + chunk_size
+        chunk = words[start:end]
+        if not chunk:
+            break
+        chunks.append(" ".join(chunk))
+        start += (chunk_size - overlap)
+    return chunks
+
+
+def extract_text_from_pdf(path: str) -> str:
+    buff = []
+    try:
+        reader = PdfReader(path)
+        for page in reader.pages:
+            text = page.extract_text()
+            if text:
+                buff.append(text)
+    except Exception as e:  # pragma: no cover
+        return f"Error reading {path}: {e}"
+    return "\n".join(buff)
+
+
+# ----------------------------
+# Embeddings: MedCPT encoders
+# ----------------------------
+MEDCPT_ARTICLE = "ncbi/MedCPT-Article-Encoder"
+MEDCPT_QUERY = "ncbi/MedCPT-Query-Encoder"
+MAX_ART_LEN = 512
+MAX_QUERY_LEN = 64
+EMBED_DIM = 768
+
+_article_tok = None
+_article_model = None
+_query_tok = None
+_query_model = None
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def load_medcpt():
+    global _article_tok, _article_model, _query_tok, _query_model
+    if _article_model and _query_model:
+        return
+    _article_tok = AutoTokenizer.from_pretrained(MEDCPT_ARTICLE, use_auth_token=HF_TOKEN)
+    _article_model = AutoModel.from_pretrained(MEDCPT_ARTICLE, use_auth_token=HF_TOKEN)
+    _article_model.to(DEVICE).eval()
+    _query_tok = AutoTokenizer.from_pretrained(MEDCPT_QUERY, use_auth_token=HF_TOKEN)
+    _query_model = AutoModel.from_pretrained(MEDCPT_QUERY, use_auth_token=HF_TOKEN)
+    _query_model.to(DEVICE).eval()
+
+
+@spaces.GPU()
+def embed_chunks(chunks: List[str]) -> np.ndarray:
+    load_medcpt()
+    all_vecs = []
+    with torch.no_grad():
+        for i in range(0, len(chunks), 8):
+            batch = chunks[i : i + 8]
+            enc = _article_tok(
+                batch,
+                truncation=True,
+                padding=True,
+                return_tensors="pt",
+                max_length=MAX_ART_LEN,
+            ).to(DEVICE)
+            out = _article_model(**enc)
+            vec = out.last_hidden_state[:, 0, :].cpu().numpy()
+            all_vecs.append(vec)
+    if not all_vecs:
+        return np.array([])
+    return np.vstack(all_vecs)
+
+
+@spaces.GPU()
+def embed_query(query: str) -> np.ndarray:
+    load_medcpt()
+    with torch.no_grad():
+        enc = _query_tok(
+            query,
+            truncation=True,
+            padding=True,
+            return_tensors="pt",
+            max_length=MAX_QUERY_LEN,
+        ).to(DEVICE)
+        out = _query_model(**enc)
+        vec = out.last_hidden_state[:, 0, :].cpu().numpy()
+    return vec
+
+
+# ----------------------------
+# FAISS index helpers
+# ----------------------------
+def build_index(embeddings: np.ndarray) -> faiss.IndexFlatIP:
+    if embeddings.dtype != np.float32:
+        embeddings = embeddings.astype(np.float32)
+    index = faiss.IndexFlatIP(embeddings.shape[1])
+    index.add(embeddings)
+    return index
+
+
+def save_index(index: faiss.IndexFlatIP, meta: List[Dict]):
+    faiss.write_index(index, str(INDEX_PATH))
+    META_PATH.write_text(json.dumps(meta, indent=2), encoding="utf-8")
+
+
+def load_index() -> Tuple[faiss.IndexFlatIP, List[Dict]]:
+    if not INDEX_PATH.exists() or not META_PATH.exists():
+        return None, None
+    idx = faiss.read_index(str(INDEX_PATH))
+    meta = json.loads(META_PATH.read_text(encoding="utf-8"))
+    return idx, meta
+
+
+def search(index: faiss.IndexFlatIP, meta: List[Dict], query_vec: np.ndarray, k: int) -> List[Dict]:
+    if query_vec.dtype != np.float32:
+        query_vec = query_vec.astype(np.float32)
+    scores, inds = index.search(query_vec, k)
+    candidates = []
+    for score, ind in zip(scores[0], inds[0]):
+        if ind < 0 or ind >= len(meta):
+            continue
+        item = dict(meta[ind])
+        item["retrieval_score"] = float(score)
+        candidates.append(item)
+    return candidates
+
+
+# ----------------------------
+# Model registry for generators
+# ----------------------------
+class GeneratorWrapper:
+    def __init__(self, name: str, load_fn):
+        self.name = name
+        self._load_fn = load_fn
+        self._pipe = None
+
+    def ensure(self):
+        if self._pipe is None:
+            self._pipe = self._load_fn()
+        return self._pipe
+
+    def generate_stream(self, prompt: str, max_new_tokens: int, temperature: float, top_p: float):
+        pipe = self.ensure()
+        streamer = TextIteratorStreamer(pipe.tokenizer, skip_special_tokens=True, skip_prompt=True)
+        kwargs = {
+            "max_new_tokens": max_new_tokens,
+            "do_sample": True,
+            "temperature": temperature,
+            "top_p": top_p,
+            "streamer": streamer,
+            "return_full_text": False,
+        }
+        thread = Thread(target=pipe, args=(prompt,), kwargs=kwargs)
+        thread.start()
+        for token in streamer:
+            yield token
+
+
+def load_gpt_oss():
+    return pipeline(
+        "text-generation",
+        model="openai/gpt-oss-20b",
+        trust_remote_code=True,
+        device_map="auto",
+        torch_dtype="auto",
+    )
+
+
+def load_gemma():
+    return pipeline(
+        "text-generation",
+        model="google/gemma-3-12b-it",
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+    )
+
+
+def load_deepseek_vl2():
+    return pipeline(
+        "text-generation",
+        model="deepseek-ai/DeepSeek-VL2-small",
+        device_map="auto",
+        torch_dtype="auto",
+        trust_remote_code=True,
+    )
+
+
+def load_granite_vision():
+    return pipeline(
+        "text-generation",
+        model="ibm-granite/granite-3.1-8b-instruct",
+        device_map="auto",
+        torch_dtype="auto",
+        trust_remote_code=True,
+    )
+
+
+def load_granite_docling():
+    # Smaller text-centric model; requires auth for some users.
+    return pipeline(
+        "text-generation",
+        model="ibm-granite/granite-docling-258M",
+        device_map="auto",
+        torch_dtype="auto",
+        use_auth_token=HF_TOKEN,
+    )
+
+
+GENERATORS = {
+    "gpt-oss-20b": GeneratorWrapper("gpt-oss-20b", load_gpt_oss),
+    "gemma-3-12b-it": GeneratorWrapper("gemma-3-12b-it", load_gemma),
+    "deepseek-vl2-small": GeneratorWrapper("deepseek-vl2-small", load_deepseek_vl2),
+    "granite-vision": GeneratorWrapper("granite-vision", load_granite_vision),
+    "granite-docling-258M": GeneratorWrapper("granite-docling-258M", load_granite_docling),
+}
+
+
+# ----------------------------
+# Prompt formatting
+# ----------------------------
+def format_prompt(query: str, contexts: List[Dict]) -> str:
+    context_blocks = []
+    for i, c in enumerate(contexts):
+        context_blocks.append(
+            f"--- Context {i+1} (file={c.get('filename','N/A')} chunk={c.get('chunk_id','?')}) ---\n{c.get('text','')}"
+        )
+    joined = "\n\n".join(context_blocks) if context_blocks else "None."
+    prompt = (
+        "You are MedDiscover, a biomedical assistant. Use ONLY the provided context to answer concisely.\n"
+        "If the context does not contain the answer, reply: 'Not found in provided documents.'\n\n"
+        f"{joined}\n\nQuestion: {query}\nAnswer:"
+    )
+    return prompt
+
+
+# ----------------------------
+# Gradio callbacks
+# ----------------------------
+@spaces.GPU()
+def process_pdfs(files, progress=gr.Progress()):
+    if not files:
+        return "Upload PDFs first."
+    texts = []
+    meta = []
+    doc_id = 0
+    for idx, f in enumerate(files):
+        progress(((idx) / max(len(files), 1)), desc=f"Reading {Path(f.name).name}")
+        text = extract_text_from_pdf(f.name)
+        if not text or text.startswith("Error reading"):
+            continue
+        chunks = chunk_text(text)
+        for cid, chunk in enumerate(chunks):
+            meta.append({"doc_id": doc_id, "filename": Path(f.name).name, "chunk_id": cid, "text": chunk})
+            texts.append(chunk)
+        doc_id += 1
+    if not texts:
+        return "No text extracted."
+    progress(0.7, desc=f"Embedding {len(texts)} chunks")
+    embeds = embed_chunks(texts)
+    if embeds.size == 0:
+        return "Embedding failed."
+    progress(0.85, desc="Building index")
+    idx = build_index(embeds)
+    save_index(idx, meta)
+    progress(1.0, desc="Ready")
+    return f"Processed {doc_id} PDFs. Index size={idx.ntotal}, dim={idx.d}. Saved to /data."
+
+
+def handle_query(query, model_key, k, max_new_tokens, temperature, top_p):
+    if not query or query.strip() == "":
+        return "Enter a query", "No context"
+    idx, meta = load_index()
+    if idx is None or meta is None:
+        return "Index not ready. Process PDFs first.", "No context"
+    qvec = embed_query(query)
+    cands = search(idx, meta, qvec, int(k))
+    prompt = format_prompt(query, cands)
+    wrapper = GENERATORS[model_key]
+    stream = wrapper.generate_stream(prompt, int(max_new_tokens), float(temperature), float(top_p))
+    answer_accum = ""
+    for chunk in stream:
+        answer_accum += chunk
+        yield answer_accum, prompt
+
+
+# ----------------------------
+# Gradio UI
+# ----------------------------
+with gr.Blocks(title="MedDiscover-HF") as demo:
+    gr.Markdown("# 🩺 MedDiscover-HF\nRAG over your PDFs with OSS HF models on ZeroGPU.")
+    with gr.Row():
+        with gr.Column(scale=1):
+            api_info = gr.Markdown("Optional: set `HF_TOKEN` secret in Space settings if needed.")
+            pdfs = gr.File(label="Upload PDFs", file_types=[".pdf"], file_count="multiple")
+            process_btn = gr.Button("Process PDFs (chunk/embed/index)", variant="primary")
+            status = gr.Textbox(label="Status", interactive=False)
+            model_dd = gr.Dropdown(
+                label="Generator Model",
+                choices=list(GENERATORS.keys()),
+                value="gpt-oss-20b",
+                interactive=True,
+            )
+            k_slider = gr.Slider(1, 10, value=3, step=1, label="Top-k chunks")
+            max_tokens = gr.Slider(20, 512, value=128, step=8, label="Max new tokens")
+            temp = gr.Slider(0.1, 1.5, value=0.4, step=0.1, label="Temperature")
+            top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+        with gr.Column(scale=2):
+            query = gr.Textbox(label="Query", lines=3, placeholder="Ask about your documents...")
+            answer = gr.Textbox(label="Answer (streaming)", lines=6)
+            context_box = gr.Textbox(label="Context used in prompt", lines=14)
+            go_btn = gr.Button("Ask", variant="primary")
+
+    process_btn.click(fn=process_pdfs, inputs=pdfs, outputs=status, show_progress="full")
+    go_btn.click(
+        fn=handle_query,
+        inputs=[query, model_dd, k_slider, max_tokens, temp, top_p],
+        outputs=[answer, context_box],
+        concurrency_limit=1,
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch()

requirements.txt

@@ -0,0 +1,9 @@
+gradio>=3.44,<=5.50
+transformers>=4.40.0
+faiss-cpu
+PyPDF2
+spaces
+openai-harmony
+sentencepiece
+accelerate
+pillow