Replace default Streamlit template with medical RAG chat UI

Wires src/streamlit_app.py to RagPipeline, adds the rag package under
src/, and updates the Dockerfile/requirements for the pipeline's deps.
artifacts/ is a placeholder for now — index build is a follow-up.

Dockerfile

@@ -10,6 +10,7 @@ RUN apt-get update && apt-get install -y \
 
 COPY requirements.txt ./
 COPY src/ ./src/
+COPY artifacts/ ./artifacts/
 
 RUN pip3 install -r requirements.txt
 

README.md

@@ -1,19 +1,20 @@
 ---
-title: Medical Assistant Ui
-emoji: 🚀
-colorFrom: red
-colorTo: red
+title: Medical Assistant UI
+emoji: 🩺
+colorFrom: blue
+colorTo: purple
 sdk: docker
 app_port: 8501
 tags:
 - streamlit
 pinned: false
-short_description: Streamlit template space
+short_description: RAG-grounded medical Q&A chat
 ---
 
-# Welcome to Streamlit!
+# Medical Assistant
 
-Edit `/src/streamlit_app.py` to customize this app to your heart's desire. :heart:
+A chat UI over a retrieval-augmented question-answering pipeline built on the
+MedQuAD dataset. Ask a medical question and get an answer grounded in
+retrieved source passages, with sources cited below each response.
 
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+Not a substitute for professional medical advice.

requirements.txt

@@ -1,3 +1,7 @@
-altair
+streamlit
 pandas
-streamlit
+pyarrow
+sentence-transformers
+faiss-cpu
+openai
+python-dotenv

src/rag/__init__.py

@@ -0,0 +1,7 @@
+# Load .env once, at package import, so scripts/notebooks/API all behave alike.
+try:
+    from dotenv import load_dotenv
+
+    load_dotenv()
+except ImportError:  # optional dependency; env vars still work if exported
+    pass

src/rag/chunking.py

@@ -0,0 +1,46 @@
+"""Word-window chunking with overlap (notebook cells 12, 15)."""
+
+import pandas as pd
+
+from rag.config import CONFIG
+
+
+def chunk_text(text: str, chunk_size: int = CONFIG.chunk_size,
+               overlap: int = CONFIG.overlap) -> list[str]:
+    words = text.split()
+    chunks: list[str] = []
+    start = 0
+
+    while start < len(words):
+        end = start + chunk_size
+        chunks.append(" ".join(words[start:end]))
+        if end >= len(words):
+            break
+        start = end - overlap
+
+    return chunks
+
+
+def chunk_corpus(corpus_df: pd.DataFrame) -> pd.DataFrame:
+    """Explode the corpus into one row per chunk, carrying metadata along.
+
+    The metadata is what makes an answer citable later, so it rides with
+    every chunk rather than being re-joined at query time.
+    """
+    records = []
+
+    for row in corpus_df.itertuples(index=False):
+        for i, chunk in enumerate(chunk_text(row.text)):
+            records.append({
+                "chunk_id": f"{row.record_id}_{i}",
+                "record_id": row.record_id,
+                "chunk_number": i,
+                "chunk_text": chunk,
+                "question": row.question,
+                "source": row.source,
+                "focus_area": row.focus_area,
+            })
+
+    chunks_df = pd.DataFrame(records)
+    chunks_df["chunk_word_count"] = chunks_df["chunk_text"].str.split().str.len()
+    return chunks_df

src/rag/config.py

@@ -0,0 +1,52 @@
+"""Single place for every knob in the pipeline.
+
+Everything downstream imports from here, so a change to chunk size or the
+embedding model never means hunting through modules.
+"""
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+ARTIFACTS_DIR = PROJECT_ROOT / "artifacts"
+DATASET_DIR = PROJECT_ROOT / "dataset"
+
+
+@dataclass(frozen=True)
+class Config:
+    # data
+    kaggle_dataset: str = "pythonafroz/medquad-medical-question-answer-for-ai-research"
+    csv_name: str = "medquad.csv"
+    local_csv: Path = DATASET_DIR / "rag1" / "medquad.csv"
+
+    # chunking
+    chunk_size: int = 500
+    overlap: int = 50
+
+    # embeddings
+    embedding_model: str = "all-MiniLM-L6-v2"
+    batch_size: int = 32
+
+    # retrieval
+    top_k: int = 5
+
+    # generation
+    llm_model: str = "gpt-5.6-luna"
+
+    # persisted artifacts
+    index_path: Path = ARTIFACTS_DIR / "faiss.index"
+    chunks_path: Path = ARTIFACTS_DIR / "chunks.parquet"
+
+    @property
+    def openai_api_key(self) -> str:
+        key = os.environ.get("OPENAI_API_KEY")
+        if not key:
+            raise RuntimeError(
+                "OPENAI_API_KEY is not set. Put it in .env or export it; "
+                "never hardcode it in a notebook cell."
+            )
+        return key
+
+
+CONFIG = Config()

src/rag/data/clean.py

@@ -0,0 +1,22 @@
+"""Cleaning + corpus construction (notebook cells 6-7)."""
+
+import pandas as pd
+
+
+def clean(df: pd.DataFrame) -> pd.DataFrame:
+    """Drop exact duplicate rows and records with no answer."""
+    return df.drop_duplicates().dropna(subset=["answer"])
+
+
+def build_corpus(df: pd.DataFrame) -> pd.DataFrame:
+    """Fuse question + answer into the single `text` field we embed."""
+    corpus = df.copy()
+    corpus["text"] = (
+        "Question: " + corpus["question"].astype(str)
+        + "\n\n"
+        + "Answer: " + corpus["answer"].astype(str)
+    )
+    corpus = corpus.reset_index(drop=True)
+    corpus["record_id"] = corpus.index
+    corpus["word_count"] = corpus["text"].str.split().str.len()
+    return corpus

src/rag/data/loader.py

@@ -0,0 +1,18 @@
+"""Get the raw MedQuAD dataframe (notebook cell 2)."""
+
+import os
+
+import pandas as pd
+
+from rag.config import CONFIG
+
+
+def load_raw() -> pd.DataFrame:
+    """Prefer the local copy in dataset/; fall back to a Kaggle download."""
+    if CONFIG.local_csv.exists():
+        return pd.read_csv(CONFIG.local_csv)
+
+    import kagglehub
+
+    path = kagglehub.dataset_download(CONFIG.kaggle_dataset)
+    return pd.read_csv(os.path.join(path, CONFIG.csv_name))

src/rag/embeddings.py

@@ -0,0 +1,22 @@
+"""Sentence-transformer embeddings (notebook cells 21-27)."""
+
+from functools import lru_cache
+
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+from rag.config import CONFIG
+
+
+@lru_cache(maxsize=1)
+def get_model() -> SentenceTransformer:
+    """Cached so importing this module twice doesn't reload 90MB of weights."""
+    return SentenceTransformer(CONFIG.embedding_model)
+
+
+def embed(texts: list[str], show_progress: bool = False) -> np.ndarray:
+    return get_model().encode(
+        texts,
+        batch_size=CONFIG.batch_size,
+        show_progress_bar=show_progress,
+    ).astype("float32")

src/rag/generator.py

@@ -0,0 +1,17 @@
+"""LLM call (notebook cells 39-44)."""
+
+from functools import lru_cache
+
+from openai import OpenAI
+
+from rag.config import CONFIG
+
+
+@lru_cache(maxsize=1)
+def get_client() -> OpenAI:
+    return OpenAI(api_key=CONFIG.openai_api_key)
+
+
+def generate(prompt: str, model: str = CONFIG.llm_model) -> str:
+    response = get_client().responses.create(model=model, input=prompt)
+    return response.output_text

src/rag/index.py

@@ -0,0 +1,31 @@
+"""FAISS index build / save / load (notebook cells 30-31)."""
+
+import faiss
+import numpy as np
+import pandas as pd
+
+from rag.config import CONFIG
+
+
+def build_index(embeddings: np.ndarray) -> faiss.Index:
+    """Inner product over L2-normalized vectors == cosine similarity."""
+    matrix = embeddings.astype("float32")
+    faiss.normalize_L2(matrix)
+
+    index = faiss.IndexFlatIP(matrix.shape[1])
+    index.add(matrix)
+    return index
+
+
+def save(index: faiss.Index, chunks_df: pd.DataFrame) -> None:
+    CONFIG.index_path.parent.mkdir(parents=True, exist_ok=True)
+    faiss.write_index(index, str(CONFIG.index_path))
+    chunks_df.to_parquet(CONFIG.chunks_path, index=False)
+
+
+def load() -> tuple[faiss.Index, pd.DataFrame]:
+    if not CONFIG.index_path.exists():
+        raise FileNotFoundError(
+            f"No index at {CONFIG.index_path}. Run: python scripts/build_index.py"
+        )
+    return faiss.read_index(str(CONFIG.index_path)), pd.read_parquet(CONFIG.chunks_path)

src/rag/pipeline.py

@@ -0,0 +1,44 @@
+"""The two halves of the diagram in the notebook, as functions.
+
+`build()` is the offline path (dataset -> index). `RagPipeline.answer()` is
+the online path (question -> grounded answer + sources).
+"""
+
+from dataclasses import dataclass
+
+from rag import index as index_store
+from rag.chunking import chunk_corpus
+from rag.config import CONFIG
+from rag.data.clean import build_corpus, clean
+from rag.data.loader import load_raw
+from rag.embeddings import embed
+from rag.generator import generate
+from rag.prompts import build_prompt
+from rag.retriever import Retrieved, Retriever
+
+
+def build() -> None:
+    """Offline: download -> clean -> corpus -> chunk -> embed -> index -> disk."""
+    corpus = build_corpus(clean(load_raw()))
+    chunks_df = chunk_corpus(corpus)
+    embeddings = embed(chunks_df["chunk_text"].tolist(), show_progress=True)
+    index_store.save(index_store.build_index(embeddings), chunks_df)
+    print(f"Indexed {len(chunks_df)} chunks from {len(corpus)} records.")
+
+
+@dataclass
+class Answer:
+    query: str
+    text: str
+    sources: list[Retrieved]
+
+
+class RagPipeline:
+    def __init__(self) -> None:
+        idx, chunks_df = index_store.load()
+        self.retriever = Retriever(idx, chunks_df)
+
+    def answer(self, query: str, top_k: int = CONFIG.top_k) -> Answer:
+        results = self.retriever.retrieve(query, top_k=top_k)
+        text = generate(build_prompt(query, results))
+        return Answer(query=query, text=text, sources=results)

src/rag/prompts.py

@@ -0,0 +1,44 @@
+"""Context assembly and the grounding prompt (notebook cells 42-43).
+
+Kept in its own module because the wording here is the main lever against
+hallucination — it should be diffable and reviewable on its own.
+"""
+
+from rag.retriever import Retrieved
+
+SYSTEM_RULES = """You are a medical information assistant.
+
+Answer the user's question using ONLY the medical context provided below.
+
+Rules:
+- Do not add medical facts that are not supported by the context.
+- Do not make a diagnosis.
+- Do not prescribe treatment or medication.
+- If the context is insufficient to answer the question,
+  clearly say that the retrieved information is insufficient.
+- Keep the answer clear and concise.
+- After the answer, list the sources used."""
+
+
+def build_context(results: list[Retrieved]) -> str:
+    parts = [
+        f"""Source: {r.source}
+Focus area: {r.focus_area}
+Original question: {r.question}
+
+Medical context:
+{r.chunk_text}"""
+        for r in results
+    ]
+    return "\n\n---\n\n".join(parts)
+
+
+def build_prompt(query: str, results: list[Retrieved]) -> str:
+    return f"""{SYSTEM_RULES}
+
+Medical context:
+{build_context(results)}
+
+User question:
+{query}
+"""

src/rag/retriever.py

@@ -0,0 +1,48 @@
+"""Query -> top-k chunks with metadata (notebook cell 32)."""
+
+from dataclasses import dataclass
+
+import faiss
+import pandas as pd
+
+from rag.config import CONFIG
+from rag.embeddings import embed
+
+
+@dataclass
+class Retrieved:
+    rank: int
+    score: float
+    chunk_id: str
+    record_id: int
+    question: str
+    source: str
+    focus_area: str
+    chunk_text: str
+
+
+class Retriever:
+    def __init__(self, index: faiss.Index, chunks_df: pd.DataFrame):
+        self.index = index
+        self.chunks_df = chunks_df
+
+    def retrieve(self, query: str, top_k: int = CONFIG.top_k) -> list[Retrieved]:
+        query_embedding = embed([query])
+        faiss.normalize_L2(query_embedding)
+
+        scores, indices = self.index.search(query_embedding, top_k)
+
+        results = []
+        for rank, (score, idx) in enumerate(zip(scores[0], indices[0]), start=1):
+            row = self.chunks_df.iloc[idx]
+            results.append(Retrieved(
+                rank=rank,
+                score=float(score),
+                chunk_id=row["chunk_id"],
+                record_id=int(row["record_id"]),
+                question=row["question"],
+                source=row["source"],
+                focus_area=row["focus_area"],
+                chunk_text=row["chunk_text"],
+            ))
+        return results

src/streamlit_app.py

@@ -1,40 +1,45 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
 import streamlit as st
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+from rag.pipeline import RagPipeline
+
+st.set_page_config(page_title="Medical Assistant", page_icon="🩺")
+
+
+@st.cache_resource
+def load_pipeline() -> RagPipeline:
+    return RagPipeline()
+
+
+pipeline = load_pipeline()
+
+st.title("Medical Assistant")
+st.caption(
+    "Ask a medical question and get an answer grounded in the MedQuAD dataset, "
+    "with sources cited below each response. Not a substitute for professional "
+    "medical advice."
+)
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+if question := st.chat_input("What are the symptoms of glaucoma?"):
+    st.session_state.messages.append({"role": "user", "content": question})
+    with st.chat_message("user"):
+        st.markdown(question)
+
+    with st.chat_message("assistant"):
+        result = pipeline.answer(question)
+        sources = "\n".join(f"- {s.source} — {s.focus_area}" for s in result.sources)
+        answer = f"{result.text}\n\n**Sources**\n{sources}"
+        st.markdown(answer)
+
+    st.session_state.messages.append({"role": "assistant", "content": answer})