Initial deploy to HF Spaces (clean history, LFS for all binaries)

.gitattributes

@@ -0,0 +1,2 @@
+*.index filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,19 @@
+__pycache__/
+*.pyc
+.venv/
+data/raw/
+data/processed/
+data/eda/
+data/codet5p-ft/
+# data/index/ is excluded globally but force-included below for HF Spaces deploy
+data/index/
+!data/index/
+!data/index/code.index
+!data/index/corpus.parquet
+!data/index/embed_model.txt
+*.parquet
+!data/index/*.parquet
+*.index
+!data/index/*.index
+.ipynb_checkpoints/
+.claude/settings.local.json

CLAUDE.md

@@ -0,0 +1,112 @@
+# Code Generation Assistant — Claude Context
+
+RAG-based Python code generation assistant using CodeSearchNet.
+Compares baseline, RAG, fine-tuned, and agentic approaches.
+
+## Environment
+
+- macOS, no NVIDIA GPU. All local runs must stay small and CPU-friendly.
+- Python 3.9; virtual environment at `.venv/` (never touch system Python).
+- Always activate with `.venv/bin/python` (or `.venv/bin/<tool>`); don't use bare `python`.
+
+## Pipeline run order
+
+```bash
+# 1. Install (one-time)
+python3 -m venv .venv
+.venv/bin/pip install -r requirements.txt
+
+# 2. Smoke test (synthetic data, fast)
+#    Set use_sample: true in config.yaml first.
+.venv/bin/python scripts/01_prepare_data.py
+.venv/bin/python scripts/02_run_eda.py
+
+# 3. Real subset (set use_sample: false, keep max_rows: 5000 in config.yaml)
+.venv/bin/python scripts/01_prepare_data.py   # downloads CodeSearchNet ~457k rows, caps at max_rows
+.venv/bin/python scripts/03_build_index.py    # downloads all-MiniLM-L6-v2, embeds corpus, writes FAISS index
+
+# 4. Launch UI (downloads Qwen2.5-Coder-1.5B-Instruct ~3 GB on first run)
+.venv/bin/python app/gradio_app.py            # serves at http://127.0.0.1:7860
+```
+
+## config.yaml key settings
+
+| Key | Default | Notes |
+|-----|---------|-------|
+| `data.use_sample` | `false` | Set `true` for offline/CI smoke tests |
+| `data.sample_size` | 200 | Rows generated when `use_sample: true` |
+| `data.max_rows` | 5000 | Caps real HF data for local runs (0 = no cap) |
+| `models.embed_model` | `sentence-transformers/all-MiniLM-L6-v2` | Retrieval embedder |
+| `models.gen_model` | `Qwen/Qwen2.5-Coder-1.5B-Instruct` | Code LLM |
+| `models.top_k` | 3 | Retrieved examples per query |
+
+## What each script does
+
+- `scripts/01_prepare_data.py` — load raw dataset (HF or synthetic) → clean → train/val/test split → `data/processed/`
+- `scripts/02_run_eda.py` — compute stats + plots from training split → `data/eda/`
+- `scripts/03_build_index.py` — embed training corpus with MiniLM → FAISS index → `data/index/`
+- `scripts/04_run_eval.py` — retrieval metrics (recall@k, MRR) + pass@1 baseline vs RAG
+- `scripts/05_finetune.py` — fine-tune CodeT5+ on docstring→code (Colab only; too slow locally)
+
+## Source layout
+
+```
+src/config.py          loads config.yaml into a SimpleNamespace
+src/data/load.py       HF dataset fetch + max_rows cap
+src/data/clean.py      filtering funnel (word count, tokens, dedup, etc.)
+src/data/make_sample.py synthetic 200-row sample for smoke tests
+src/eda/analyze.py     stats + matplotlib/seaborn plots
+src/rag/embedder.py    CodeIndex: SentenceTransformer + FAISS (build/save/load/retrieve)
+src/rag/generator.py   CodeAssistant: Qwen LLM wrapper, baseline + RAG prompt builders
+src/eval/             functional_eval.py, retrieval_eval.py, sandbox.py
+src/agent/repair_loop.py  generate → run → self-repair loop
+src/finetune/train_codet5.py  (Colab only)
+app/gradio_app.py      Gradio chat UI (main local + HF Spaces deploy target)
+app/api.py             FastAPI REST service (uvicorn)
+app/streamlit_app.py   Streamlit UI
+```
+
+## HuggingFace downloads (one-time, cached in ~/.cache/huggingface/)
+
+| Asset | Size | When |
+|-------|------|------|
+| `code_search_net` dataset | ~2 GB | `01_prepare_data.py` with `use_sample: false` |
+| `sentence-transformers/all-MiniLM-L6-v2` | ~90 MB | `03_build_index.py` (first run) |
+| `Qwen/Qwen2.5-Coder-1.5B-Instruct` | ~3 GB | `app/gradio_app.py` (first run) |
+
+## Data directories (excluded from git)
+
+```
+data/raw/          raw parquet from HF
+data/processed/    train/val/test.parquet + cleaning_funnel.csv
+data/eda/          PNG plots + eda_stats.json
+data/index/        code.index (FAISS) + corpus.parquet + embed_model.txt
+```
+
+## Deployment options
+
+```bash
+# Gradio (local or push to HF Spaces as app.py)
+.venv/bin/python app/gradio_app.py
+
+# FastAPI
+.venv/bin/uvicorn app.api:app --host 0.0.0.0 --port 8000
+
+# Streamlit
+.venv/bin/streamlit run app/streamlit_app.py
+
+# Docker
+docker build -t cga . && docker run -p 8000:8000 cga
+```
+
+## Full-dataset training / eval
+
+Do NOT run locally — use Colab:
+- `scripts/04_run_eval.py` on full CodeSearchNet is slow; fine for small subsets.
+- `scripts/05_finetune.py` (CodeT5+) requires a GPU.
+- The notebook (`notebooks/`) is for Colab EDA, training, and reporting eval numbers.
+
+## Known warnings (non-fatal)
+
+- `urllib3 NotOpenSSLWarning` — macOS LibreSSL vs OpenSSL; safe to ignore.
+- `Some parameters are on the meta device` — CPU offload of Qwen weights; expected on macOS without GPU.

Dockerfile

@@ -0,0 +1,19 @@
+# Container for the FastAPI service. Build: docker build -t code-gen-assistant .
+# Run:   docker run -p 8000:8000 code-gen-assistant
+# For generated-code execution safety, run with --network none if you don't need
+# the model to download at runtime (bake the index/model into the image instead).
+FROM python:3.11-slim
+
+WORKDIR /app
+
+# System deps for tokenizers / faiss wheels.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential git && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 8000
+CMD ["uvicorn", "app.api:app", "--host", "0.0.0.0", "--port", "8000"]

README.md

@@ -0,0 +1,90 @@
+---
+title: Code Generation Assistant
+sdk: gradio
+app_file: app/gradio_app.py
+pinned: false
+---
+
+# Code Generation Assistant
+
+Generate Python code from natural-language descriptions, grounded in
+**CodeSearchNet** via retrieval (RAG), with functional evaluation and a
+deployable chat interface.
+
+## Approaches compared
+1. **Baseline** - frozen code LLM, zero/few-shot
+2. **RAG** - retrieve similar CodeSearchNet examples, condition the LLM
+3. **Fine-tuned** - CodeT5+ trained on `docstring -> code`
+4. **Agentic** - generate -> run -> read error -> repair loop
+
+## Evaluation
+- **CodeBLEU** (similarity to reference) - in the notebook
+- **pass@k** on HumanEval / MBPP (functional correctness) - `src/eval/functional_eval.py`
+- **recall@k / MRR** (retrieval quality) - `src/eval/retrieval_eval.py`
+
+> CodeSearchNet ships no unit tests, so pass@k is measured on HumanEval/MBPP
+> (which do), while CodeSearchNet powers retrieval + the similarity metric.
+
+## Pipeline (run in order)
+```bash
+pip install -r requirements.txt
+python scripts/01_prepare_data.py   # load -> clean -> split
+python scripts/02_run_eda.py        # stats + plots
+python scripts/03_build_index.py    # embed corpus -> FAISS index (persisted)
+python scripts/04_run_eval.py       # retrieval + functional pass@1 (baseline vs RAG)
+python scripts/05_finetune.py       # (optional) fine-tune CodeT5+
+```
+
+## Project layout
+```
+config.yaml              # single source of truth (models, paths, thresholds)
+src/
+  data/        load.py clean.py make_sample.py     # Phase 1
+  eda/         analyze.py                          # Phase 1
+  rag/         embedder.py  generator.py           # Phase 3 + 5 (CodeAssistant)
+  eval/        functional_eval.py retrieval_eval.py sandbox.py   # Phase 2
+  agent/       repair_loop.py                       # Phase 6
+  finetune/    train_codet5.py                       # Phase 4
+app/
+  api.py            # FastAPI REST service
+  gradio_app.py     # Gradio chat UI (Hugging Face Spaces)
+  streamlit_app.py  # Streamlit chat UI
+scripts/            # numbered phase entrypoints
+notebooks/          # experimentation notebook
+Dockerfile          # container for the API
+```
+
+## Is the notebook the right vehicle?
+The notebook is for **experimentation, EDA, and reporting eval numbers** - keep
+it for your capstone appendix. For anything you *deploy*, the logic belongs in
+the `src/` package (importable, testable, version-controlled). The apps in
+`app/` all import the same `CodeAssistant`, so there is one implementation, not
+three copies. Workflow: prototype in the notebook -> harden into `src/` ->
+push to GitHub -> deploy an app.
+
+## Deployment interfaces
+
+**1. Gradio on Hugging Face Spaces (easiest).** Copy `app/gradio_app.py` to a new
+Gradio Space as `app.py`, add `requirements.txt`, pick a GPU tier. Public chat UI
+in minutes, no servers to manage.
+
+**2. FastAPI (production / integration).**
+```bash
+uvicorn app.api:app --host 0.0.0.0 --port 8000   # docs at /docs
+curl -X POST localhost:8000/generate -H 'Content-Type: application/json' \
+     -d '{"intent": "function to check if a number is prime", "mode": "rag"}'
+```
+Use this when another system (an IDE plugin, a CI bot, a web frontend) needs to
+call the assistant programmatically. Endpoints: `GET /health`, `POST /generate`
+(supports `mode`, `repair`, `return_sources`).
+
+**3. Streamlit.** `streamlit run app/streamlit_app.py` - deploys free to
+Streamlit Community Cloud from a GitHub repo.
+
+**4. Docker (any cloud).** `docker build -t cga . && docker run -p 8000:8000 cga`.
+
+## Security note
+Generated code is executed (for pass@k and the repair loop) via a subprocess +
+timeout in `src/eval/sandbox.py`. That guards against hangs but is **not** a
+security sandbox. For public deployment, run execution inside a container with
+`--network none` and dropped privileges, or disable the repair feature.

app/api.py

@@ -0,0 +1,78 @@
+"""FastAPI service - the production interface.
+
+Exposes the CodeAssistant over REST so any client (web app, CI bot, IDE plugin)
+can call it. Run locally:
+    uvicorn app.api:app --host 0.0.0.0 --port 8000
+Then POST to /generate. Interactive docs at /docs.
+
+The model loads once at startup (lifespan) and is reused across requests.
+"""
+from __future__ import annotations
+
+import os
+import sys
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+from fastapi import FastAPI
+from pydantic import BaseModel, Field
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+
+# Allow tests / CI to skip the heavy model load.
+_STATE: dict = {"assistant": None}
+
+
+def get_assistant():
+    return _STATE["assistant"]
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    if os.environ.get("CGA_SKIP_MODEL") != "1":
+        from src.rag.generator import CodeAssistant
+        _STATE["assistant"] = CodeAssistant.from_config(with_index=True)
+    yield
+    _STATE["assistant"] = None
+
+
+app = FastAPI(title="Code Generation Assistant", version="1.0", lifespan=lifespan)
+
+
+class GenerateRequest(BaseModel):
+    intent: str = Field(..., description="Natural-language description of the code")
+    mode: str = Field("rag", description="'rag' or 'baseline'")
+    repair: bool = Field(False, description="Run the agentic repair loop")
+    return_sources: bool = Field(True)
+
+
+class GenerateResponse(BaseModel):
+    code: str
+    mode: str
+    sources: list = []
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok", "model_loaded": _STATE["assistant"] is not None}
+
+
+@app.post("/generate", response_model=GenerateResponse)
+def generate(req: GenerateRequest):
+    assistant = get_assistant()
+    if assistant is None:
+        return GenerateResponse(code="# model not loaded", mode=req.mode)
+
+    if req.repair:
+        from src.agent.repair_loop import make_repair_generator, repair_loop
+        gen_fn = make_repair_generator(assistant)
+        trace = repair_loop(req.intent, gen_fn,
+                            check_program_fn=lambda c: c, max_iters=3)
+        return GenerateResponse(code=trace.final_code, mode=f"{req.mode}+repair")
+
+    result = assistant.generate(req.intent, mode=req.mode,
+                                return_sources=req.return_sources)
+    if isinstance(result, tuple):
+        code, sources = result
+        return GenerateResponse(code=code, mode=req.mode, sources=sources)
+    return GenerateResponse(code=result, mode=req.mode)

app/gradio_app.py

@@ -0,0 +1,52 @@
+"""Gradio chat UI - the easiest deploy target (Hugging Face Spaces).
+
+To deploy on HF Spaces:
+  1. Create a new Space (SDK: Gradio).
+  2. Push this file as `app.py` at the repo root, plus requirements.txt.
+  3. Pick a GPU tier if you want the 7B model; the 1.5B runs on CPU (slowly).
+
+Locally:  python app/gradio_app.py
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import gradio as gr
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+from src.rag.generator import CodeAssistant  # noqa: E402
+
+assistant = CodeAssistant.from_config(with_index=True)
+
+
+def respond(message, history, mode, use_repair):
+    if use_repair:
+        from src.agent.repair_loop import make_repair_generator, repair_loop
+        gen_fn = make_repair_generator(assistant)
+        trace = repair_loop(message, gen_fn, check_program_fn=lambda c: c, max_iters=3)
+        code, note = trace.final_code, f"(repair: {'passed' if trace.success else 'best effort'} in {trace.iterations} iters)"
+        return f"```python\n{code}\n```\n\n_{note}_"
+
+    result = assistant.generate(message, mode=mode, return_sources=True)
+    if isinstance(result, tuple):
+        code, sources = result
+        src_md = "\n".join(f"- ({s['score']:.2f}) {s['docstring']}" for s in sources)
+        return f"```python\n{code}\n```\n\n**Retrieved examples:**\n{src_md}"
+    return f"```python\n{result}\n```"
+
+
+with gr.Blocks(title="Code Generation Assistant") as demo:
+    gr.Markdown("# Code Generation Assistant\nDescribe what you want; it writes Python, grounded in CodeSearchNet.")
+    with gr.Row():
+        mode = gr.Radio(["rag", "baseline"], value="rag", label="Mode")
+        use_repair = gr.Checkbox(label="Agentic repair (run + self-fix)", value=False)
+    gr.ChatInterface(
+        fn=respond,
+        additional_inputs=[mode, use_repair],
+        examples=[["Write a function to check if a number is prime."],
+                  ["Parse a CSV file and return a list of dicts."]],
+    )
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0")

app/streamlit_app.py

@@ -0,0 +1,63 @@
+"""Streamlit chat UI - alternative front-end (Streamlit Community Cloud).
+
+Run:  streamlit run app/streamlit_app.py
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import streamlit as st
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+
+
+@st.cache_resource
+def get_assistant():
+    from src.rag.generator import CodeAssistant
+    return CodeAssistant.from_config(with_index=True)
+
+
+st.set_page_config(page_title="Code Generation Assistant", page_icon="</>")
+st.title("Code Generation Assistant")
+
+with st.sidebar:
+    mode = st.radio("Mode", ["rag", "baseline"], index=0)
+    use_repair = st.checkbox("Agentic repair (run + self-fix)", value=False)
+
+assistant = get_assistant()
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+for m in st.session_state.messages:
+    with st.chat_message(m["role"]):
+        st.markdown(m["content"])
+
+if prompt := st.chat_input("Describe the function you want..."):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    with st.chat_message("assistant"), st.spinner("Generating..."):
+        if use_repair:
+            from src.agent.repair_loop import make_repair_generator, repair_loop
+            gen_fn = make_repair_generator(assistant)
+            trace = repair_loop(prompt, gen_fn, check_program_fn=lambda c: c, max_iters=3)
+            code = trace.final_code
+            st.code(code, language="python")
+            st.caption(f"Repair: {'passed' if trace.success else 'best effort'} in {trace.iterations} iters")
+            answer = code
+        else:
+            result = assistant.generate(prompt, mode=mode, return_sources=True)
+            if isinstance(result, tuple):
+                code, sources = result
+                st.code(code, language="python")
+                with st.expander("Retrieved examples"):
+                    for s in sources:
+                        st.write(f"({s['score']:.2f}) {s['docstring']}")
+            else:
+                code = result
+                st.code(code, language="python")
+            answer = code
+    st.session_state.messages.append({"role": "assistant", "content": f"```python\n{answer}\n```"})

config.yaml

@@ -0,0 +1,53 @@
+# ============================================================
+# Code Generation Assistant - Project Configuration
+# ============================================================
+# Central config so every phase reads the same settings.
+# Change `use_sample: true` to test the pipeline on synthetic
+# data; set false to pull the real CodeSearchNet from HuggingFace.
+
+data:
+  # HuggingFace dataset id. The canonical id is "code_search_net".
+  # If that fails to load, try the community mirror:
+  #   "code-search-net/code_search_net"
+  hf_dataset_id: "code_search_net"
+  # Languages to include. Start with python only; add "java" etc. later.
+  languages:
+    - python
+  # When true, the pipeline uses a small synthetic sample instead of
+  # downloading the real dataset (useful for offline testing / CI).
+  use_sample: false
+  sample_size: 200          # rows generated when use_sample is true
+  max_rows: 5000            # cap on real HF data (0 = no cap); keeps local runs fast
+
+cleaning:
+  min_doc_words: 3          # drop docstrings shorter than this (words)
+  max_doc_words: 120        # drop overly long docstrings (likely noise)
+  min_code_chars: 20        # drop trivially short functions
+  max_code_tokens: 512      # drop functions longer than this (token budget)
+  drop_exact_duplicates: true
+  drop_non_ascii_docs: true # drop docstrings that are mostly non-ASCII
+  # Substrings that flag low-quality / autogenerated docstrings.
+  doc_blocklist:
+    - "todo"
+    - "fixme"
+    - "auto-generated"
+    - "autogenerated"
+    - "do not edit"
+
+split:
+  train: 0.8
+  val: 0.1
+  test: 0.1
+  seed: 42
+
+models:
+  embed_model: "sentence-transformers/all-MiniLM-L6-v2"
+  gen_model: "Qwen/Qwen2.5-Coder-1.5B-Instruct"
+  top_k: 3
+
+paths:
+  data_dir: "data"
+  raw_dir: "data/raw"
+  processed_dir: "data/processed"
+  eda_dir: "data/eda"
+  index_dir: "data/index"

data/index/code.index

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44f5653ddb9f84a940c75aa9c483df8a8193252695529e011d49a1821d73fa7a
+size 5343789

data/index/corpus.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfbd0e7044327f45bdb99a3a4d0ac0aa60b635f88cb709eaa5b8f9ba5d39f3e2
+size 1783627

data/index/embed_model.txt

@@ -0,0 +1 @@
+sentence-transformers/all-MiniLM-L6-v2

notebooks/Code_Generation_Assistant.ipynb

@@ -0,0 +1,652 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "43fed051",
+   "metadata": {},
+   "source": [
+    "# Code Generation Assistant \n",
+    "\n",
+    "**Generate Python code from natural-language descriptions, grounded in CodeSearchNet.**\n",
+    "\n",
+    "This notebook runs the core vertical slice of the capstone top-to-bottom:\n",
+    "\n",
+    "1. **Phase 1** - load + clean CodeSearchNet, EDA\n",
+    "2. **Phase 3** - embed the corpus + build a FAISS retrieval index\n",
+    "3. **Phase 5** - RAG: retrieve similar examples and condition a code LLM\n",
+    "4. **Eval** - baseline (no retrieval) vs RAG, scored with CodeBLEU\n",
+    "5. **Interactive** - ask it to write code\n",
+    "6. **Phase 4 (optional)** - fine-tune CodeT5+ on docstring->code\n",
+    "\n",
+    "> CodeSearchNet was built for code *search*, so it ships `(docstring, code)` pairs\n",
+    "> and **no unit tests**. We treat the docstring summary as the intent and the\n",
+    "> function body as the target. Because it is natively a retrieval corpus, RAG is\n",
+    "> the most natural architecture here.\n",
+    "\n",
+    "**Runtime:** set `Runtime -> Change runtime type -> T4 GPU`. No API key required -\n",
+    "generation uses a small local model (`Qwen2.5-Coder-1.5B-Instruct`)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "06dd65ca",
+   "metadata": {},
+   "source": [
+    "## 0. Setup\n",
+    "\n",
+    "Installs everything. `codebleu` is optional (it builds tree-sitter parsers); if it\n",
+    "fails the eval falls back to a token-overlap metric so the notebook still runs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2733fd6d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip -q install datasets transformers accelerate sentence-transformers faiss-cpu pandas matplotlib seaborn\n",
+    "# codebleu needs a tree-sitter parser to actually run; install it too (optional - has a fallback)\n",
+    "!pip -q install codebleu tree-sitter tree-sitter-python || echo \"codebleu/parser install failed - will use fallback metric\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1125e98a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "print(\"CUDA available:\", torch.cuda.is_available())\n",
+    "print(\"Device:\", torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"CPU (generation will be slow)\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "88a40082",
+   "metadata": {},
+   "source": [
+    "## 1. Config\n",
+    "\n",
+    "One place for every knob. `MAX_ROWS` keeps the Colab run fast - raise it (or set to\n",
+    "`None`) for a fuller run. `python` only to start; depth over breadth."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "894bfb5f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from dataclasses import dataclass, field\n",
+    "from typing import Tuple\n",
+    "\n",
+    "@dataclass\n",
+    "class Config:\n",
+    "    # data\n",
+    "    candidate_dataset_ids: Tuple[str, ...] = (\n",
+    "        \"code-search-net/code_search_net\",  # parquet mirror (most reliable)\n",
+    "        \"code_search_net\",                   # canonical (may need older datasets)\n",
+    "    )\n",
+    "    language: str = \"python\"\n",
+    "    max_rows: int = 8000        # subset for speed; set None for full split\n",
+    "    # cleaning\n",
+    "    min_doc_words: int = 3\n",
+    "    max_doc_words: int = 120\n",
+    "    min_code_chars: int = 20\n",
+    "    max_code_tokens: int = 400\n",
+    "    doc_blocklist: Tuple[str, ...] = (\"todo\", \"fixme\", \"auto-generated\",\n",
+    "                                      \"autogenerated\", \"do not edit\")\n",
+    "    # split\n",
+    "    seed: int = 42\n",
+    "    train: float = 0.8\n",
+    "    val: float = 0.1\n",
+    "    # models\n",
+    "    embed_model: str = \"sentence-transformers/all-MiniLM-L6-v2\"\n",
+    "    gen_model: str = \"Qwen/Qwen2.5-Coder-1.5B-Instruct\"\n",
+    "    top_k: int = 3\n",
+    "\n",
+    "CFG = Config()\n",
+    "CFG"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4382632d",
+   "metadata": {},
+   "source": [
+    "## 2. Phase 1a - Load CodeSearchNet\n",
+    "\n",
+    "Tries the parquet mirror first, then the canonical id. If both fail on your\n",
+    "`datasets` version, run `!pip install \"datasets<3\"` and re-run, or download the\n",
+    "raw release from the CodeSearchNet GitHub repo."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "158ab53b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "import pandas as pd\n",
+    "\n",
+    "USE_COLS = {\n",
+    "    \"func_documentation_string\": \"docstring\",\n",
+    "    \"func_code_string\": \"code\",\n",
+    "    \"language\": \"language\",\n",
+    "    \"repository_name\": \"repo\",\n",
+    "    \"func_code_url\": \"url\",\n",
+    "}\n",
+    "\n",
+    "def load_codesearchnet(cfg):\n",
+    "    last_err = None\n",
+    "    for ds_id in cfg.candidate_dataset_ids:\n",
+    "        try:\n",
+    "            print(f\"[load] trying '{ds_id}' ({cfg.language}) ...\")\n",
+    "            ds = load_dataset(ds_id, cfg.language, split=\"train\", trust_remote_code=True)\n",
+    "            if cfg.max_rows:\n",
+    "                ds = ds.select(range(min(cfg.max_rows, len(ds))))\n",
+    "            df = ds.to_pandas()\n",
+    "            keep = [c for c in USE_COLS if c in df.columns]\n",
+    "            df = df[keep].rename(columns=USE_COLS)\n",
+    "            for col in USE_COLS.values():\n",
+    "                if col not in df.columns:\n",
+    "                    df[col] = \"\"\n",
+    "            print(f\"[load] OK - {len(df):,} rows from '{ds_id}'\")\n",
+    "            return df\n",
+    "        except Exception as e:\n",
+    "            print(f\"[load] failed: {e}\")\n",
+    "            last_err = e\n",
+    "    raise RuntimeError(f\"All dataset ids failed. Last error: {last_err}\")\n",
+    "\n",
+    "raw = load_codesearchnet(CFG)\n",
+    "raw.head(2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "32dba780",
+   "metadata": {},
+   "source": [
+    "## 3. Phase 1b - Clean & filter\n",
+    "\n",
+    "CodeSearchNet is noisy. We keep only the **summary first line** of each docstring\n",
+    "as the intent (the rest is usually `:param:`/`:return:` boilerplate), then apply\n",
+    "quality filters and dedup. The **funnel** logs how many rows each filter removes -\n",
+    "keep it for your write-up."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2aa0af5e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import re\n",
+    "\n",
+    "WORD_RE = re.compile(r\"\\b\\w+\\b\")\n",
+    "\n",
+    "def first_line(t):\n",
+    "    return t.strip().split(\"\\n\")[0].strip() if isinstance(t, str) else \"\"\n",
+    "\n",
+    "def word_count(t):\n",
+    "    return len(WORD_RE.findall(t)) if isinstance(t, str) else 0\n",
+    "\n",
+    "def ascii_ratio(t):\n",
+    "    if not t:\n",
+    "        return 1.0\n",
+    "    return sum(1 for ch in t if ord(ch) < 128) / len(t)\n",
+    "\n",
+    "def approx_tokens(c):\n",
+    "    return len(re.findall(r\"\\w+|[^\\s\\w]\", c)) if isinstance(c, str) else 0\n",
+    "\n",
+    "def clean(df, cfg):\n",
+    "    funnel = [(\"raw\", len(df))]\n",
+    "    df = df.copy()\n",
+    "    df[\"docstring\"] = df[\"docstring\"].map(first_line)\n",
+    "    df[\"code\"] = df[\"code\"].fillna(\"\").astype(str)\n",
+    "\n",
+    "    df = df[(df[\"docstring\"].str.len() > 0) & (df[\"code\"].str.len() > 0)]\n",
+    "    funnel.append((\"non_empty\", len(df)))\n",
+    "\n",
+    "    wc = df[\"docstring\"].map(word_count)\n",
+    "    df = df[(wc >= cfg.min_doc_words) & (wc <= cfg.max_doc_words)]\n",
+    "    funnel.append((\"doc_word_window\", len(df)))\n",
+    "\n",
+    "    df = df[df[\"code\"].str.len() >= cfg.min_code_chars]\n",
+    "    funnel.append((\"min_code_chars\", len(df)))\n",
+    "\n",
+    "    df = df[df[\"code\"].map(approx_tokens) <= cfg.max_code_tokens]\n",
+    "    funnel.append((\"max_code_tokens\", len(df)))\n",
+    "\n",
+    "    pat = \"|\".join(re.escape(t) for t in cfg.doc_blocklist)\n",
+    "    df = df[~df[\"docstring\"].str.lower().str.contains(pat, regex=True)]\n",
+    "    funnel.append((\"doc_blocklist\", len(df)))\n",
+    "\n",
+    "    df = df[df[\"docstring\"].map(ascii_ratio) >= 0.9]\n",
+    "    funnel.append((\"ascii_docs\", len(df)))\n",
+    "\n",
+    "    df = df.drop_duplicates(subset=[\"code\"]).drop_duplicates(subset=[\"docstring\"])\n",
+    "    funnel.append((\"dedup\", len(df)))\n",
+    "\n",
+    "    funnel_df = pd.DataFrame(funnel, columns=[\"step\", \"rows_remaining\"])\n",
+    "    return df.reset_index(drop=True), funnel_df\n",
+    "\n",
+    "clean_df, funnel = clean(raw, CFG)\n",
+    "print(funnel.to_string(index=False))\n",
+    "print(\"\\nClean rows:\", len(clean_df))\n",
+    "clean_df.head(2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "20747f0a",
+   "metadata": {},
+   "source": [
+    "## 4. Phase 1c - EDA\n",
+    "\n",
+    "Quick look at the cleaned corpus: docstring length, code length, and the cleaning\n",
+    "funnel. Save these for the report appendix."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f684c430",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import seaborn as sns\n",
+    "sns.set_theme(style=\"whitegrid\")\n",
+    "\n",
+    "doc_words = clean_df[\"docstring\"].map(word_count)\n",
+    "code_lines = clean_df[\"code\"].str.count(\"\\n\") + 1\n",
+    "\n",
+    "fig, axes = plt.subplots(1, 3, figsize=(16, 4))\n",
+    "sns.histplot(doc_words, bins=40, ax=axes[0]); axes[0].set(title=\"Docstring length (words)\", xlabel=\"words\")\n",
+    "sns.histplot(code_lines.clip(upper=80), bins=40, ax=axes[1]); axes[1].set(title=\"Code length (lines, clipped 80)\", xlabel=\"lines\")\n",
+    "axes[2].barh(funnel[\"step\"], funnel[\"rows_remaining\"]); axes[2].invert_yaxis(); axes[2].set(title=\"Cleaning funnel\")\n",
+    "plt.tight_layout(); plt.show()\n",
+    "\n",
+    "print({\n",
+    "    \"rows\": len(clean_df),\n",
+    "    \"doc_words_median\": int(doc_words.median()),\n",
+    "    \"code_lines_median\": int(code_lines.median()),\n",
+    "})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "db098ca9",
+   "metadata": {},
+   "source": [
+    "## 5. Train / val / test split\n",
+    "\n",
+    "The **train** pool doubles as the retrieval corpus for RAG. We evaluate on **test**\n",
+    "so retrieved examples never leak the answer."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0c18c3e2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def split(df, cfg):\n",
+    "    df = df.sample(frac=1.0, random_state=cfg.seed).reset_index(drop=True)\n",
+    "    n = len(df); n_tr = int(n * cfg.train); n_va = int(n * cfg.val)\n",
+    "    return (df.iloc[:n_tr].reset_index(drop=True),\n",
+    "            df.iloc[n_tr:n_tr+n_va].reset_index(drop=True),\n",
+    "            df.iloc[n_tr+n_va:].reset_index(drop=True))\n",
+    "\n",
+    "train_df, val_df, test_df = split(clean_df, CFG)\n",
+    "print(f\"train={len(train_df)}  val={len(val_df)}  test={len(test_df)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b2d3b684",
+   "metadata": {},
+   "source": [
+    "## 6. Phase 3 - Embeddings + FAISS index\n",
+    "\n",
+    "Embed each docstring in the train pool and build a cosine-similarity index\n",
+    "(`IndexFlatIP` on L2-normalised vectors). The default embedder is small and fast;\n",
+    "for a stronger code-aware corpus, swap `embed_model` to\n",
+    "`Salesforce/codet5p-110m-embedding` (ties into the CodeT5 family)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "68145a4c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sentence_transformers import SentenceTransformer\n",
+    "import faiss\n",
+    "import numpy as np\n",
+    "\n",
+    "embedder = SentenceTransformer(CFG.embed_model)\n",
+    "corpus = train_df.reset_index(drop=True)\n",
+    "\n",
+    "corpus_emb = embedder.encode(\n",
+    "    corpus[\"docstring\"].tolist(),\n",
+    "    batch_size=64, show_progress_bar=True,\n",
+    "    convert_to_numpy=True, normalize_embeddings=True,\n",
+    ").astype(\"float32\")\n",
+    "\n",
+    "index = faiss.IndexFlatIP(corpus_emb.shape[1])\n",
+    "index.add(corpus_emb)\n",
+    "print(\"Indexed vectors:\", index.ntotal, \"| dim:\", corpus_emb.shape[1])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "386988c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def retrieve(query, k=None):\n",
+    "    k = k or CFG.top_k\n",
+    "    q = embedder.encode([query], convert_to_numpy=True,\n",
+    "                        normalize_embeddings=True).astype(\"float32\")\n",
+    "    scores, idx = index.search(q, k)\n",
+    "    out = corpus.iloc[idx[0]].copy()\n",
+    "    out[\"score\"] = scores[0]\n",
+    "    return out\n",
+    "\n",
+    "# sanity check\n",
+    "retrieve(\"read a json file from disk and return a dict\")[[\"docstring\", \"score\"]]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "45bad6b2",
+   "metadata": {},
+   "source": [
+    "## 7. Phase 5a - Load the code LLM\n",
+    "\n",
+    "`Qwen2.5-Coder-1.5B-Instruct` fits on a free T4. For higher quality (and a Colab Pro\n",
+    "GPU) bump `gen_model` to `Qwen/Qwen2.5-Coder-7B-Instruct`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c42beea4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+    "\n",
+    "tok = AutoTokenizer.from_pretrained(CFG.gen_model)\n",
+    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "    CFG.gen_model, torch_dtype=\"auto\", device_map=\"auto\"\n",
+    ")\n",
+    "\n",
+    "def chat_generate(messages, max_new_tokens=320):\n",
+    "    text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
+    "    inputs = tok(text, return_tensors=\"pt\").to(model.device)\n",
+    "    out = model.generate(**inputs, max_new_tokens=max_new_tokens,\n",
+    "                         do_sample=False, pad_token_id=tok.eos_token_id)\n",
+    "    new = out[0][inputs.input_ids.shape[1]:]\n",
+    "    return tok.decode(new, skip_special_tokens=True)\n",
+    "\n",
+    "def extract_code(text):\n",
+    "    \"\"\"Strip markdown fences if the model wrapped the code.\"\"\"\n",
+    "    m = re.search(r\"```(?:python)?\\n(.*?)```\", text, re.DOTALL)\n",
+    "    return m.group(1).strip() if m else text.strip()\n",
+    "\n",
+    "print(\"Model loaded:\", CFG.gen_model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2622d4fd",
+   "metadata": {},
+   "source": [
+    "## 8. Phase 5b - Baseline vs RAG prompts\n",
+    "\n",
+    "Same model, two prompting strategies. The RAG prompt injects the top-k retrieved\n",
+    "`(docstring, code)` pairs as dynamic few-shot context."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "54db7627",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SYS = (\"You are an expert Python coding assistant. Write a single, correct, \"\n",
+    "       \"self-contained Python function for the request. Output only code.\")\n",
+    "\n",
+    "def baseline_messages(intent):\n",
+    "    return [{\"role\": \"system\", \"content\": SYS},\n",
+    "            {\"role\": \"user\", \"content\": f\"# Task: {intent}\"}]\n",
+    "\n",
+    "def rag_messages(intent, k=None):\n",
+    "    ex = retrieve(intent, k)\n",
+    "    blocks = [f\"# Task: {r.docstring}\\n{r.code}\" for _, r in ex.iterrows()]\n",
+    "    context = \"\\n\\n\".join(blocks)\n",
+    "    user = (f\"Here are similar reference examples:\\n\\n{context}\\n\\n\"\n",
+    "            f\"# Now write a function for this task:\\n# Task: {intent}\")\n",
+    "    return [{\"role\": \"system\", \"content\": SYS},\n",
+    "            {\"role\": \"user\", \"content\": user}]\n",
+    "\n",
+    "demo = \"Write a function that returns the n-th Fibonacci number.\"\n",
+    "print(\"=== BASELINE ===\")\n",
+    "print(extract_code(chat_generate(baseline_messages(demo))))\n",
+    "print(\"\\n=== RAG ===\")\n",
+    "print(extract_code(chat_generate(rag_messages(demo))))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f2f6fda3",
+   "metadata": {},
+   "source": [
+    "## 9. Eval - CodeBLEU, baseline vs RAG\n",
+    "\n",
+    "We score generated code against the reference on held-out **test** rows. CodeBLEU\n",
+    "weights AST + data-flow match, not just text overlap. If `codebleu` did not install,\n",
+    "we fall back to a token-overlap F1 so the cell still runs.\n",
+    "\n",
+    "> Caveat: CodeSearchNet has no unit tests, so this measures *similarity to the\n",
+    "> reference*, not functional correctness. For pass@k, add a HumanEval/MBPP harness\n",
+    "> (Phase 2) - flagged in the next-steps cell."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8530179f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Try CodeBLEU; fall back to token-F1 if the metric OR its parser is unavailable.\n",
+    "score, METRIC = None, None\n",
+    "try:\n",
+    "    from codebleu import calc_codebleu\n",
+    "    # actually CALL it once - this is what needs the tree-sitter parser\n",
+    "    _ = calc_codebleu([\"def f(): return 1\"], [\"def f(): return 1\"], lang=\"python\")\n",
+    "    def score(ref, hyp):\n",
+    "        return calc_codebleu([ref], [hyp], lang=\"python\")[\"codebleu\"]\n",
+    "    METRIC = \"CodeBLEU\"\n",
+    "except Exception as e:\n",
+    "    print(\"CodeBLEU unavailable, using token-F1 fallback:\", e)\n",
+    "    def _toks(s):\n",
+    "        return set(re.findall(r\"\\w+\", s))\n",
+    "    def score(ref, hyp):\n",
+    "        a, b = _toks(ref), _toks(hyp)\n",
+    "        if not a or not b:\n",
+    "            return 0.0\n",
+    "        inter = len(a & b)\n",
+    "        p, rec = inter / len(b), inter / len(a)\n",
+    "        return 0.0 if p + rec == 0 else 2 * p * rec / (p + rec)\n",
+    "    METRIC = \"token-F1 (fallback)\"\n",
+    "print(\"Using metric:\", METRIC)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b1f22892",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "N_EVAL = 15  # keep small on free Colab; raise for the real run\n",
+    "sample = test_df.sample(min(N_EVAL, len(test_df)), random_state=CFG.seed)\n",
+    "\n",
+    "rows = []\n",
+    "for _, r in sample.iterrows():\n",
+    "    base = extract_code(chat_generate(baseline_messages(r.docstring)))\n",
+    "    rag = extract_code(chat_generate(rag_messages(r.docstring)))\n",
+    "    rows.append({\"baseline\": score(r.code, base), \"rag\": score(r.code, rag)})\n",
+    "\n",
+    "res = pd.DataFrame(rows)\n",
+    "print(f\"Mean {METRIC} over {len(res)} test tasks:\")\n",
+    "print(res.mean().round(4).to_string())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "86b042c5",
+   "metadata": {},
+   "source": [
+    "## 10. Interactive - ask it to write code\n",
+    "\n",
+    "Edit the string and run. This uses the RAG pipeline and shows the retrieved\n",
+    "examples so the grounding is visible."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13cf8cc1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def ask(intent, show_sources=True):\n",
+    "    if show_sources:\n",
+    "        print(\"Retrieved examples:\")\n",
+    "        for _, r in retrieve(intent).iterrows():\n",
+    "            print(f\"  - ({r.score:.2f}) {r.docstring}\")\n",
+    "        print(\"-\" * 50)\n",
+    "    print(extract_code(chat_generate(rag_messages(intent))))\n",
+    "\n",
+    "ask(\"Write a function to check whether a string is a valid IPv4 address.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "61adc276",
+   "metadata": {},
+   "source": [
+    "## 11. (Optional) Phase 4 - Fine-tune CodeT5+\n",
+    "\n",
+    "A compact demonstration of the fine-tuning arm: train `codet5p-220m` on a small\n",
+    "`docstring -> code` subset for a few steps so you can see the loop work, then\n",
+    "generate. For the real capstone result, raise `subset`/`epochs` and run on a Pro\n",
+    "GPU. **This section is slow - skip on a first pass.**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c2f5217e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Set to True to run fine-tuning.\n",
+    "RUN_FINETUNE = False\n",
+    "\n",
+    "if RUN_FINETUNE:\n",
+    "    from transformers import (AutoTokenizer, AutoModelForSeq2SeqLM,\n",
+    "                              Seq2SeqTrainer, Seq2SeqTrainingArguments,\n",
+    "                              DataCollatorForSeq2Seq)\n",
+    "    from datasets import Dataset\n",
+    "\n",
+    "    ck = \"Salesforce/codet5p-220m\"\n",
+    "    t5_tok = AutoTokenizer.from_pretrained(ck)\n",
+    "    t5 = AutoModelForSeq2SeqLM.from_pretrained(ck)\n",
+    "\n",
+    "    subset = train_df.head(2000)\n",
+    "    def to_features(batch):\n",
+    "        x = t5_tok(batch[\"docstring\"], max_length=64, truncation=True, padding=\"max_length\")\n",
+    "        y = t5_tok(text_target=batch[\"code\"], max_length=256, truncation=True, padding=\"max_length\")\n",
+    "        x[\"labels\"] = y[\"input_ids\"]\n",
+    "        return x\n",
+    "\n",
+    "    hf = Dataset.from_pandas(subset[[\"docstring\", \"code\"]]).map(\n",
+    "        to_features, batched=True, remove_columns=[\"docstring\", \"code\"])\n",
+    "\n",
+    "    args = Seq2SeqTrainingArguments(\n",
+    "        output_dir=\"codet5p-ft\", per_device_train_batch_size=8,\n",
+    "        num_train_epochs=1, learning_rate=5e-5, logging_steps=20,\n",
+    "        fp16=torch.cuda.is_available(), report_to=\"none\", save_strategy=\"no\")\n",
+    "\n",
+    "    trainer = Seq2SeqTrainer(\n",
+    "        model=t5, args=args, train_dataset=hf,\n",
+    "        data_collator=DataCollatorForSeq2Seq(t5_tok, model=t5))\n",
+    "    trainer.train()\n",
+    "\n",
+    "    def t5_generate(intent):\n",
+    "        ids = t5_tok(intent, return_tensors=\"pt\").input_ids.to(t5.device)\n",
+    "        out = t5.generate(ids, max_length=256)\n",
+    "        return t5_tok.decode(out[0], skip_special_tokens=True)\n",
+    "\n",
+    "    print(t5_generate(\"Return the factorial of a non-negative integer n.\"))\n",
+    "else:\n",
+    "    print(\"Fine-tuning skipped. Set RUN_FINETUNE = True to run it.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ed8b964b",
+   "metadata": {},
+   "source": [
+    "## 12. Next steps + deploying to VS Code\n",
+    "\n",
+    "**What's still to add for the full capstone:**\n",
+    "- **Phase 2 functional eval:** wire up HumanEval / MBPP for real `pass@k` (they ship\n",
+    "  unit tests, unlike CodeSearchNet). This is the metric graders trust most.\n",
+    "- **Phase 6 agentic loop:** generate -> run in a sandbox -> read traceback -> repair.\n",
+    "- **Retrieval quality:** measure recall@k / MRR on the search task to justify the embedder.\n",
+    "\n",
+    "**Lifting this into VS Code for deployment:**\n",
+    "1. The functions here map cleanly onto the repo modules: `clean()` -> `src/data/clean.py`,\n",
+    "   `retrieve()` + index build -> `src/rag/retriever.py`, `chat_generate()`/prompts ->\n",
+    "   `src/rag/generator.py`.\n",
+    "2. Persist the FAISS index (`faiss.write_index(index, \"index.faiss\")`) and the corpus\n",
+    "   so you don't rebuild on every start.\n",
+    "3. Wrap `ask()` in a **Streamlit** app (`app.py`) for the Phase 7 chat UI:\n",
+    "   `streamlit run app.py`.\n",
+    "4. Keep `config.yaml` as the single source of truth across notebook and app."
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

requirements.txt

@@ -0,0 +1,30 @@
+# --- Phase 1: data + EDA ---
+datasets>=2.18
+pandas>=2.0
+pyarrow>=14.0
+matplotlib>=3.7
+seaborn>=0.13
+pyyaml>=6.0
+
+# --- Phase 2: evaluation ---
+numpy>=1.24
+codebleu>=0.7
+tree-sitter>=0.21
+tree-sitter-python>=0.21
+
+# --- Phase 3: retrieval ---
+sentence-transformers>=2.7
+faiss-cpu>=1.8
+
+# --- Phase 4-6: generation, fine-tuning, agent ---
+transformers>=4.40
+torch>=2.2
+accelerate>=0.30
+peft>=0.10
+
+# --- Deployment ---
+fastapi>=0.110
+uvicorn>=0.29
+pydantic>=2.0
+gradio>=4.0
+streamlit>=1.33

scripts/01_prepare_data.py

@@ -0,0 +1,42 @@
+"""Phase 1 entrypoint: load -> clean -> split -> save processed parquet files.
+
+Usage:
+    python scripts/01_prepare_data.py
+Outputs train/val/test parquet files + a cleaning_funnel.csv into data/processed/.
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+from src.config import load_config
+from src.data.clean import clean, split
+from src.data.load import load_raw
+
+
+def main():
+    cfg = load_config()
+    print("=" * 60)
+    print("PHASE 1: DATA PREPARATION")
+    print("=" * 60)
+
+    raw = load_raw(cfg)
+    cleaned, funnel = clean(raw, cfg)
+
+    print("\nCleaning funnel:")
+    print(funnel.to_string(index=False))
+
+    splits = split(cleaned, cfg)
+    out = Path(cfg.paths.processed_dir)
+    for name, part in splits.items():
+        path = out / f"{name}.parquet"
+        part.to_parquet(path, index=False)
+        print(f"  saved {name}: {len(part):,} rows -> {path}")
+
+    funnel.to_csv(out / "cleaning_funnel.csv", index=False)
+    print(f"\nDone. Processed data in {out}")
+
+
+if __name__ == "__main__":
+    main()

scripts/02_run_eda.py

@@ -0,0 +1,40 @@
+"""Phase 1 EDA entrypoint: analyse the processed training split.
+
+Usage:
+    python scripts/02_run_eda.py
+Reads data/processed/train.parquet, writes plots + eda_stats.json to data/eda/.
+"""
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+import pandas as pd
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+from src.config import load_config
+from src.eda.analyze import run_eda
+
+
+def main():
+    cfg = load_config()
+    print("=" * 60)
+    print("PHASE 1: EXPLORATORY DATA ANALYSIS")
+    print("=" * 60)
+
+    train_path = Path(cfg.paths.processed_dir) / "train.parquet"
+    funnel_path = Path(cfg.paths.processed_dir) / "cleaning_funnel.csv"
+    if not train_path.exists():
+        sys.exit("train.parquet not found. Run scripts/01_prepare_data.py first.")
+
+    df = pd.read_parquet(train_path)
+    funnel = pd.read_csv(funnel_path) if funnel_path.exists() else None
+
+    stats = run_eda(df, cfg, funnel)
+    print(json.dumps({k: v for k, v in stats.items() if k != "plots"}, indent=2))
+    print(f"\nPlots saved to {cfg.paths.eda_dir}")
+
+
+if __name__ == "__main__":
+    main()

scripts/03_build_index.py

@@ -0,0 +1,14 @@
+"""Phase 3 entrypoint: build + persist the FAISS retrieval index.
+
+Usage:  python scripts/03_build_index.py
+Reads data/processed/train.parquet, writes the index to data/index/.
+"""
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+from src.rag.embedder import build_index_from_processed
+
+if __name__ == "__main__":
+    print("=" * 60, "\nPHASE 3: BUILD RETRIEVAL INDEX\n", "=" * 60, sep="")
+    build_index_from_processed()

scripts/04_run_eval.py

@@ -0,0 +1,49 @@
+"""Phase 2 entrypoint: run the eval suite and print a comparison table.
+
+Usage:  python scripts/04_run_eval.py
+Loads the index + assistant, then reports:
+  - retrieval recall@k / MRR
+  - functional pass@1 on HumanEval (baseline vs RAG)
+"""
+import sys
+from pathlib import Path
+
+import pandas as pd
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+from src.config import load_config
+from src.eval.functional_eval import evaluate
+from src.eval.retrieval_eval import evaluate_cross_modal
+from src.rag.generator import CodeAssistant
+
+
+def main():
+    cfg = load_config()
+    print("=" * 60, "\nPHASE 2: EVALUATION\n", "=" * 60, sep="")
+
+    assistant = CodeAssistant.from_config(cfg, with_index=True)
+
+    # 1. Cross-modal retrieval quality on held-out test pairs.
+    test = pd.read_parquet(Path(cfg.paths.processed_dir) / "test.parquet")
+    pairs = test[["docstring", "code"]].dropna().sample(
+        n=min(500, len(test)), random_state=42
+    ).reset_index(drop=True)
+    print("\n[retrieval]", evaluate_cross_modal(assistant.index.embedder, pairs))
+
+    # 2. Functional pass@1: baseline vs RAG on HumanEval.
+    LIMIT = 12  # raise for the real run
+    base = evaluate(lambda i: assistant.generate(i, mode="baseline"),
+                    "humaneval", limit=LIMIT)
+    rag = evaluate(lambda i: assistant.generate(i, mode="rag"),
+                   "humaneval", limit=LIMIT)
+
+    table = pd.DataFrame([
+        {"system": "baseline", "pass@1": base.pass_at_1},
+        {"system": "rag", "pass@1": rag.pass_at_1},
+    ])
+    print(f"\n[functional] HumanEval pass@1 over {LIMIT} problems:")
+    print(table.to_string(index=False))
+
+
+if __name__ == "__main__":
+    main()

scripts/05_finetune.py

@@ -0,0 +1,14 @@
+"""Phase 4 entrypoint: fine-tune CodeT5+.
+
+Usage:  python scripts/05_finetune.py
+Trains on data/processed/train.parquet, saves to data/codet5p-ft/.
+"""
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+from src.finetune.train_codet5 import finetune
+
+if __name__ == "__main__":
+    print("=" * 60, "\nPHASE 4: FINE-TUNE CODET5+\n", "=" * 60, sep="")
+    finetune(subset_size=5000, epochs=1)

scripts/retrieval_only_eval.py

@@ -0,0 +1,43 @@
+"""Cross-modal retrieval eval — no LLM generation, finishes in seconds.
+
+Runs twice:
+  1. Raw code (inflated — docstring is embedded inside func_code_string).
+  2. Docstrings stripped from candidate code (cleaner semantic signal).
+
+Usage:  python scripts/retrieval_only_eval.py
+"""
+import sys
+from pathlib import Path
+
+import pandas as pd
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+from src.config import load_config
+from src.eval.retrieval_eval import evaluate_cross_modal
+from src.rag.embedder import CodeIndex
+
+cfg = load_config()
+
+print("[load] reading test split ...")
+test = pd.read_parquet(Path(cfg.paths.processed_dir) / "test.parquet")
+pairs = (
+    test[["docstring", "code"]]
+    .dropna()
+    .sample(n=min(500, len(test)), random_state=42)
+    .reset_index(drop=True)
+)
+print(f"       {len(pairs)} pairs")
+
+print("[load] loading embedder ...")
+idx = CodeIndex.load(cfg.paths.index_dir)
+
+print()
+for strip in (False, True):
+    label = "stripped (leakage-free)" if strip else "raw code  (⚠ lexical leakage)"
+    r = evaluate_cross_modal(idx.embedder, pairs, k_values=(1, 5, 10),
+                             strip_code_docstrings=strip)
+    print(f"\n=== {label} ===")
+    print(f"  N     : {r['n_pairs']}")
+    print(f"  MRR   : {r['mrr']:.4f}")
+    for k in (1, 5, 10):
+        print(f"  R@{k:2d}  : {r[f'recall@{k}']:.4f}")

scripts/test_repair_loop.py

@@ -0,0 +1,79 @@
+"""Throwaway script: exercise the agentic repair loop on a tricky problem.
+
+The intent is chosen to be genuinely hard for a small model: balanced-bracket
+checking with all three delimiter types. The check_fn appends a strict test
+harness so the sandbox can report a traceback on failure.
+
+Run:
+    python scripts/test_repair_loop.py
+Expected output: iteration trace showing whether the model self-corrects.
+"""
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+from src.agent.repair_loop import make_repair_generator, repair_loop
+from src.rag.generator import CodeAssistant
+
+# ── test harness appended to every generated function ─────────────────────
+_TEST_HARNESS = '''
+
+# --- auto-injected test harness ---
+assert is_balanced("") is True,       "empty string should be balanced"
+assert is_balanced("()") is True,     "simple parens"
+assert is_balanced("([{}])") is True, "nested mix"
+assert is_balanced("([)]") is False,  "wrong order"
+assert is_balanced("(") is False,     "unclosed"
+assert is_balanced("{[}]") is False,  "interleaved wrong"
+assert is_balanced("((()))") is True, "deep nesting"
+print("all tests passed")
+'''
+
+INTENT = (
+    "Write a Python function called `is_balanced` that returns True if the input "
+    "string has balanced parentheses (), square brackets [], and curly braces {}, "
+    "and False otherwise. Use a stack."
+)
+
+
+def check_fn(code: str) -> str:
+    """Append the test harness to the generated function."""
+    return code + _TEST_HARNESS
+
+
+def main():
+    print("=" * 60)
+    print("REPAIR LOOP TEST")
+    print("Intent:", INTENT[:80] + "...")
+    print("=" * 60)
+
+    print("\n[setup] Loading CodeAssistant (model + index)...")
+    assistant = CodeAssistant.from_config(with_index=True)
+    generate_fn = make_repair_generator(assistant)
+
+    print("[repair] Starting loop (max_iters=3)...\n")
+    trace = repair_loop(INTENT, generate_fn, check_fn, max_iters=3, timeout=10.0)
+
+    print("\n" + "=" * 60)
+    print("ITERATION TRACE")
+    print("=" * 60)
+    for i, (code, error) in enumerate(trace.history, 1):
+        print(f"\n--- Iteration {i} ---")
+        print("Generated code:")
+        print(code)
+        if error:
+            print(f"\nError:\n  {error}")
+        else:
+            print("\nResult: PASSED")
+
+    print("\n" + "=" * 60)
+    print(f"Outcome: {'SUCCESS' if trace.success else 'BEST EFFORT (did not pass)'}")
+    print(f"Iterations used: {trace.iterations}")
+    print("=" * 60)
+    if trace.success:
+        print("\nFinal passing code:")
+        print(trace.final_code)
+
+
+if __name__ == "__main__":
+    main()

src/agent/repair_loop.py

@@ -0,0 +1,58 @@
+"""Phase 6: agentic generate -> execute -> repair loop.
+
+Model-agnostic: takes a `generate_fn(intent, feedback) -> code` callable, so it
+works with CodeAssistant or any other generator (and is unit-testable with a mock).
+"""
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable
+
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+from src.eval.sandbox import run_code  # noqa: E402
+
+
+@dataclass
+class RepairTrace:
+    final_code: str
+    success: bool
+    iterations: int
+    history: list = field(default_factory=list)  # list of (code, error)
+
+
+def repair_loop(
+    intent: str,
+    generate_fn: Callable[[str, str | None], str],
+    check_program_fn: Callable[[str], str],
+    max_iters: int = 3,
+    timeout: float = 8.0,
+) -> RepairTrace:
+    """Iteratively generate and self-correct.
+
+    generate_fn(intent, feedback)  -> candidate code
+        feedback is None on the first call, else the previous error string.
+    check_program_fn(code)         -> a runnable program string (code + a smoke
+        test or the harness test) used to decide pass/fail.
+    """
+    feedback = None
+    history = []
+    for i in range(1, max_iters + 1):
+        code = generate_fn(intent, feedback)
+        result = run_code(check_program_fn(code), timeout=timeout)
+        history.append((code, result.error))
+        if result.ok:
+            return RepairTrace(code, True, i, history)
+        feedback = result.error  # feed the traceback back in
+    return RepairTrace(history[-1][0], False, max_iters, history)
+
+
+def make_repair_generator(assistant):
+    """Adapt a CodeAssistant into a generate_fn for the repair loop."""
+    def generate_fn(intent: str, feedback: str | None) -> str:
+        if feedback:
+            intent = (f"{intent}\n\n# Your previous attempt failed with this error:\n"
+                      f"# {feedback}\n# Fix it and return the corrected function.")
+        return assistant.generate(intent, mode="rag")
+    return generate_fn

src/config.py

@@ -0,0 +1,39 @@
+"""Load the project YAML config into a simple attribute-style object."""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import yaml
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def _to_namespace(obj):
+    if isinstance(obj, dict):
+        return SimpleNamespace(**{k: _to_namespace(v) for k, v in obj.items()})
+    if isinstance(obj, list):
+        return [_to_namespace(v) for v in obj]
+    return obj
+
+
+def load_config(path: str | os.PathLike | None = None) -> SimpleNamespace:
+    """Read config.yaml and return a nested namespace (cfg.data.languages, ...)."""
+    path = Path(path) if path else PROJECT_ROOT / "config.yaml"
+    with open(path, "r", encoding="utf-8") as f:
+        raw = yaml.safe_load(f)
+    cfg = _to_namespace(raw)
+    # Resolve paths relative to project root and ensure they exist.
+    for attr in ("data_dir", "raw_dir", "processed_dir", "eda_dir", "index_dir"):
+        abspath = PROJECT_ROOT / getattr(cfg.paths, attr)
+        setattr(cfg.paths, attr, str(abspath))
+        abspath.mkdir(parents=True, exist_ok=True)
+    return cfg
+
+
+if __name__ == "__main__":
+    c = load_config()
+    print("Languages:", c.data.languages)
+    print("Use sample:", c.data.use_sample)
+    print("Processed dir:", c.paths.processed_dir)

src/data/clean.py

@@ -0,0 +1,116 @@
+"""Phase 1b: clean and filter the raw (docstring, code) pairs.
+
+CodeSearchNet is noisy. A clean, filtered subset trains and retrieves better
+than the raw dump. Every filter records how many rows it removed so you can
+report the funnel in your EDA / write-up.
+"""
+from __future__ import annotations
+
+import re
+import sys
+from pathlib import Path
+
+import pandas as pd
+
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+from src.config import load_config  # noqa: E402
+
+_WORD_RE = re.compile(r"\b\w+\b")
+
+
+def _first_line(text: str) -> str:
+    """CodeSearchNet docstrings often have a summary first line + details.
+    For NL->code we keep the summary line (the actual 'intent')."""
+    return text.strip().split("\n")[0].strip() if isinstance(text, str) else ""
+
+
+def _word_count(text: str) -> int:
+    return len(_WORD_RE.findall(text)) if isinstance(text, str) else 0
+
+
+def _ascii_ratio(text: str) -> float:
+    if not text:
+        return 1.0
+    ascii_chars = sum(1 for ch in text if ord(ch) < 128)
+    return ascii_chars / len(text)
+
+
+def _approx_tokens(code: str) -> int:
+    """Cheap proxy for token count (whitespace + punctuation split)."""
+    return len(re.findall(r"\w+|[^\s\w]", code)) if isinstance(code, str) else 0
+
+
+def clean(df: pd.DataFrame, cfg=None) -> tuple[pd.DataFrame, pd.DataFrame]:
+    """Return (cleaned_df, funnel_df). funnel_df logs rows removed per step."""
+    cfg = cfg or load_config()
+    cc = cfg.cleaning
+    funnel = [("raw", len(df))]
+    df = df.copy()
+
+    # Use only the summary line of each docstring as the NL intent.
+    df["docstring"] = df["docstring"].map(_first_line)
+    df["code"] = df["code"].fillna("").astype(str)
+
+    # 1. Drop empty docstring or code.
+    df = df[(df["docstring"].str.len() > 0) & (df["code"].str.len() > 0)]
+    funnel.append(("non_empty", len(df)))
+
+    # 2. Docstring word-count window.
+    wc = df["docstring"].map(_word_count)
+    df = df[(wc >= cc.min_doc_words) & (wc <= cc.max_doc_words)]
+    funnel.append(("doc_word_window", len(df)))
+
+    # 3. Minimum code length.
+    df = df[df["code"].str.len() >= cc.min_code_chars]
+    funnel.append(("min_code_chars", len(df)))
+
+    # 4. Maximum code tokens (budget for the generator's context).
+    df = df[df["code"].map(_approx_tokens) <= cc.max_code_tokens]
+    funnel.append(("max_code_tokens", len(df)))
+
+    # 5. Blocklisted / autogenerated docstrings.
+    pattern = "|".join(re.escape(t) for t in cc.doc_blocklist)
+    if pattern:
+        df = df[~df["docstring"].str.lower().str.contains(pattern, regex=True)]
+    funnel.append(("doc_blocklist", len(df)))
+
+    # 6. Drop mostly-non-ASCII docstrings (non-English noise).
+    if cc.drop_non_ascii_docs:
+        df = df[df["docstring"].map(_ascii_ratio) >= 0.9]
+    funnel.append(("ascii_docs", len(df)))
+
+    # 7. Exact duplicate removal (same code or same docstring).
+    if cc.drop_exact_duplicates:
+        df = df.drop_duplicates(subset=["code"]).drop_duplicates(subset=["docstring"])
+    funnel.append(("dedup", len(df)))
+
+    df = df.reset_index(drop=True)
+    funnel_df = pd.DataFrame(funnel, columns=["step", "rows_remaining"])
+    funnel_df["removed"] = funnel_df["rows_remaining"].shift(1).fillna(
+        funnel_df["rows_remaining"].iloc[0]
+    ).astype(int) - funnel_df["rows_remaining"]
+    return df, funnel_df
+
+
+def split(df: pd.DataFrame, cfg=None) -> dict[str, pd.DataFrame]:
+    """Random train/val/test split per the config ratios."""
+    cfg = cfg or load_config()
+    df = df.sample(frac=1.0, random_state=cfg.split.seed).reset_index(drop=True)
+    n = len(df)
+    n_train = int(n * cfg.split.train)
+    n_val = int(n * cfg.split.val)
+    return {
+        "train": df.iloc[:n_train].reset_index(drop=True),
+        "val": df.iloc[n_train:n_train + n_val].reset_index(drop=True),
+        "test": df.iloc[n_train + n_val:].reset_index(drop=True),
+    }
+
+
+if __name__ == "__main__":
+    from src.data.load import load_raw
+
+    cfg = load_config()
+    raw = load_raw(cfg)
+    cleaned, funnel = clean(raw, cfg)
+    print(funnel.to_string(index=False))
+    print("cleaned rows:", len(cleaned))

src/data/load.py

@@ -0,0 +1,84 @@
+"""Phase 1a: load raw CodeSearchNet (or the synthetic sample).
+
+We normalise everything to a pandas DataFrame with two key columns:
+  - docstring : natural-language description (model INPUT)
+  - code      : function body (model TARGET)
+plus useful metadata (language, repo, url) for traceability.
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pandas as pd
+
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+from src.config import load_config  # noqa: E402
+from src.data.make_sample import make_sample  # noqa: E402
+
+# Map CodeSearchNet's verbose column names to our canonical names.
+_COLUMN_MAP = {
+    "func_documentation_string": "docstring",
+    "func_code_string": "code",
+    "language": "language",
+    "repository_name": "repo",
+    "func_code_url": "url",
+}
+
+
+def _from_huggingface(cfg) -> pd.DataFrame:
+    """Stream CodeSearchNet per-language from HuggingFace and concatenate."""
+    from datasets import load_dataset
+
+    max_rows = getattr(cfg.data, "max_rows", 0)
+    frames = []
+    for lang in cfg.data.languages:
+        print(f"[load] downloading CodeSearchNet '{lang}' ...")
+        # CodeSearchNet ships train/validation/test; we pull all and re-split later.
+        ds = load_dataset(cfg.data.hf_dataset_id, lang)
+        for split in ds.keys():
+            df = ds[split].to_pandas()
+            keep = [c for c in _COLUMN_MAP if c in df.columns]
+            df = df[keep].rename(columns=_COLUMN_MAP)
+            frames.append(df)
+    out = pd.concat(frames, ignore_index=True)
+    print(f"[load] total raw rows: {len(out):,}")
+    if max_rows and max_rows > 0 and len(out) > max_rows:
+        out = out.sample(n=max_rows, random_state=42).reset_index(drop=True)
+        print(f"[load] capped to {max_rows:,} rows (max_rows setting)")
+    return out
+
+
+def _from_sample(cfg) -> pd.DataFrame:
+    print(f"[load] using synthetic sample (n={cfg.data.sample_size})")
+    df = make_sample(cfg.data.sample_size, cfg.split.seed)
+    keep = [c for c in _COLUMN_MAP if c in df.columns]
+    return df[keep].rename(columns=_COLUMN_MAP)
+
+
+def load_raw(cfg=None) -> pd.DataFrame:
+    cfg = cfg or load_config()
+    if cfg.data.use_sample:
+        df = _from_sample(cfg)
+    else:
+        try:
+            df = _from_huggingface(cfg)
+        except Exception as e:  # noqa: BLE001
+            print(
+                f"[load] HuggingFace load failed ({e}).\n"
+                f"        Tip: try hf_dataset_id: 'code-search-net/code_search_net' "
+                f"in config.yaml, or set use_sample: true.",
+                file=sys.stderr,
+            )
+            raise
+    # Guarantee the columns downstream code expects, even if metadata missing.
+    for col in ("docstring", "code", "language", "repo", "url"):
+        if col not in df.columns:
+            df[col] = ""
+    return df
+
+
+if __name__ == "__main__":
+    df = load_raw()
+    print(df.shape)
+    print(df.head(3).to_string())

src/data/make_sample.py

@@ -0,0 +1,81 @@
+"""Generate a small synthetic dataset in CodeSearchNet's schema.
+
+This lets you test the full pipeline (clean -> EDA -> later phases) without
+downloading the real ~2M-row dataset. The schema matches the HuggingFace
+`code_search_net` columns we actually use downstream.
+"""
+from __future__ import annotations
+
+import random
+
+import pandas as pd
+
+# A handful of realistic (docstring, code) templates plus some deliberately
+# "dirty" rows so the cleaning step has something to remove.
+_GOOD = [
+    (
+        "Return the factorial of a non-negative integer n.",
+        "def factorial(n):\n    if n < 0:\n        raise ValueError('n must be >= 0')\n    result = 1\n    for i in range(2, n + 1):\n        result *= i\n    return result",
+    ),
+    (
+        "Compute the nth Fibonacci number using iteration.",
+        "def fib(n):\n    a, b = 0, 1\n    for _ in range(n):\n        a, b = b, a + b\n    return a",
+    ),
+    (
+        "Check whether a given string is a palindrome, ignoring case.",
+        "def is_palindrome(s):\n    s = s.lower()\n    return s == s[::-1]",
+    ),
+    (
+        "Read a JSON file from disk and return the parsed dictionary.",
+        "import json\n\ndef read_json(path):\n    with open(path) as f:\n        return json.load(f)",
+    ),
+    (
+        "Return a list of unique elements preserving original order.",
+        "def dedupe(items):\n    seen = set()\n    out = []\n    for x in items:\n        if x not in seen:\n            seen.add(x)\n            out.append(x)\n    return out",
+    ),
+    (
+        "Flatten a nested list of arbitrary depth into a single list.",
+        "def flatten(lst):\n    out = []\n    for x in lst:\n        if isinstance(x, list):\n            out.extend(flatten(x))\n        else:\n            out.append(x)\n    return out",
+    ),
+]
+
+# Rows that should be filtered out by the cleaning rules.
+_DIRTY = [
+    ("", "def noop():\n    pass"),                          # empty docstring
+    ("ok", "def f():\n    return 1"),                        # too-short docstring
+    ("TODO: write this later", "def g():\n    pass"),        # blocklisted
+    ("auto-generated do not edit", "def h():\n    pass"),    # blocklisted
+    ("Returns x.", "x"),                                     # too-short code
+    ("说明：返回输入值的两倍。", "def dbl(x):\n    return x * 2"),  # non-ascii doc
+]
+
+
+def make_sample(n: int = 200, seed: int = 42) -> pd.DataFrame:
+    """Build a DataFrame with the CodeSearchNet columns we rely on."""
+    rng = random.Random(seed)
+    rows = []
+    for i in range(n):
+        # ~15% dirty rows so cleaning has work to do.
+        if rng.random() < 0.15:
+            doc, code = rng.choice(_DIRTY)
+        else:
+            doc, code = rng.choice(_GOOD)
+        rows.append(
+            {
+                "repository_name": f"acme/repo{i % 10}",
+                "func_path_in_repository": f"src/module_{i}.py",
+                "func_name": (code.split("(")[0].replace("def ", "").strip()
+                              if code.startswith("def ") else f"sym_{i}"),
+                "language": "python",
+                "func_code_string": code,
+                "func_documentation_string": doc,
+                "func_code_url": f"https://github.com/acme/repo{i % 10}/blob/main/src/module_{i}.py",
+            }
+        )
+    return pd.DataFrame(rows)
+
+
+if __name__ == "__main__":
+    df = make_sample()
+    print(df.shape)
+    print(df.head(3).to_string())

src/eda/analyze.py

@@ -0,0 +1,111 @@
+"""Phase 1c: exploratory data analysis.
+
+Produces (a) a stats dict you can dump to JSON for the report, and
+(b) PNG plots saved to the eda dir. Keep these in your capstone appendix.
+"""
+from __future__ import annotations
+
+import json
+import re
+import sys
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")  # headless backend for servers/CI
+import matplotlib.pyplot as plt  # noqa: E402
+import pandas as pd  # noqa: E402
+import seaborn as sns  # noqa: E402
+
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+from src.config import load_config  # noqa: E402
+
+sns.set_theme(style="whitegrid")
+_WORD_RE = re.compile(r"\b\w+\b")
+
+
+def _doc_words(s: str) -> int:
+    return len(_WORD_RE.findall(s))
+
+
+def compute_stats(df: pd.DataFrame) -> dict:
+    doc_words = df["docstring"].map(_doc_words)
+    code_lines = df["code"].str.count("\n") + 1
+    code_chars = df["code"].str.len()
+    return {
+        "n_rows": int(len(df)),
+        "languages": df["language"].value_counts().to_dict(),
+        "docstring_words": {
+            "mean": round(float(doc_words.mean()), 2),
+            "median": int(doc_words.median()),
+            "p95": int(doc_words.quantile(0.95)),
+            "max": int(doc_words.max()),
+        },
+        "code_lines": {
+            "mean": round(float(code_lines.mean()), 2),
+            "median": int(code_lines.median()),
+            "p95": int(code_lines.quantile(0.95)),
+            "max": int(code_lines.max()),
+        },
+        "code_chars": {
+            "mean": round(float(code_chars.mean()), 2),
+            "median": int(code_chars.median()),
+        },
+    }
+
+
+def make_plots(df: pd.DataFrame, out_dir: str, funnel: pd.DataFrame | None = None):
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    saved = []
+
+    # Docstring length distribution.
+    fig, ax = plt.subplots(figsize=(7, 4))
+    sns.histplot(df["docstring"].map(_doc_words), bins=40, ax=ax)
+    ax.set(title="Docstring length (words)", xlabel="words", ylabel="count")
+    p = out / "docstring_length.png"
+    fig.tight_layout(); fig.savefig(p, dpi=120); plt.close(fig); saved.append(str(p))
+
+    # Code length distribution (lines).
+    fig, ax = plt.subplots(figsize=(7, 4))
+    sns.histplot((df["code"].str.count("\n") + 1).clip(upper=80), bins=40, ax=ax)
+    ax.set(title="Code length (lines, clipped at 80)", xlabel="lines", ylabel="count")
+    p = out / "code_length.png"
+    fig.tight_layout(); fig.savefig(p, dpi=120); plt.close(fig); saved.append(str(p))
+
+    # Language distribution.
+    fig, ax = plt.subplots(figsize=(7, 4))
+    df["language"].value_counts().plot(kind="bar", ax=ax)
+    ax.set(title="Rows per language", xlabel="language", ylabel="count")
+    p = out / "language_distribution.png"
+    fig.tight_layout(); fig.savefig(p, dpi=120); plt.close(fig); saved.append(str(p))
+
+    # Cleaning funnel (if provided).
+    if funnel is not None:
+        fig, ax = plt.subplots(figsize=(7, 4))
+        ax.barh(funnel["step"], funnel["rows_remaining"])
+        ax.invert_yaxis()
+        ax.set(title="Cleaning funnel (rows remaining)", xlabel="rows")
+        p = out / "cleaning_funnel.png"
+        fig.tight_layout(); fig.savefig(p, dpi=120); plt.close(fig); saved.append(str(p))
+
+    return saved
+
+
+def run_eda(df: pd.DataFrame, cfg=None, funnel: pd.DataFrame | None = None) -> dict:
+    cfg = cfg or load_config()
+    stats = compute_stats(df)
+    plots = make_plots(df, cfg.paths.eda_dir, funnel)
+    stats["plots"] = plots
+    with open(Path(cfg.paths.eda_dir) / "eda_stats.json", "w") as f:
+        json.dump(stats, f, indent=2)
+    return stats
+
+
+if __name__ == "__main__":
+    from src.data.clean import clean
+    from src.data.load import load_raw
+
+    cfg = load_config()
+    cleaned, funnel = clean(load_raw(cfg), cfg)
+    print(json.dumps(run_eda(cleaned, cfg, funnel), indent=2))

src/eval/functional_eval.py

@@ -0,0 +1,109 @@
+"""Phase 2: functional correctness eval (pass@k) on HumanEval / MBPP.
+
+Unlike CodeBLEU (similarity), this measures whether generated code actually RUNS
+and PASSES unit tests - the claim that carries a capstone defense.
+
+generate_fn(intent) -> code is injected, so this is decoupled from the model and
+unit-testable with a mock.
+"""
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+
+import numpy as np
+
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+from src.eval.sandbox import run_code  # noqa: E402
+
+
+def pass_at_k(n: int, c: int, k: int) -> float:
+    """Unbiased pass@k estimator (Codex paper). n samples, c correct."""
+    if n - c < k:
+        return 1.0
+    return 1.0 - float(np.prod(1.0 - k / np.arange(n - c + 1, n + 1)))
+
+
+# ---- per-benchmark program builders ------------------------------------
+def humaneval_program(problem, candidate_code: str) -> str:
+    # candidate_code already defines the full function (entry_point).
+    return f"{candidate_code}\n\n{problem['test']}\n\ncheck({problem['entry_point']})\n"
+
+
+def mbpp_program(problem, candidate_code: str) -> str:
+    setup = problem.get("test_setup_code", "") or ""
+    tests = "\n".join(problem.get("test_list", []))
+    return f"{candidate_code}\n\n{setup}\n{tests}\n"
+
+
+# hf_id is a tuple of candidate IDs tried in order; newer datasets versions
+# require the namespaced form, older ones or mirrors may still use the bare id.
+_BENCH = {
+    "humaneval": {
+        "hf_id": ("openai/openai_humaneval", "openai_humaneval"), "split": "test",
+        "intent_col": "prompt", "program": humaneval_program,
+    },
+    "mbpp": {
+        "hf_id": ("google-research-datasets/mbpp", "mbpp"), "split": "test",
+        "intent_col": "text", "program": mbpp_program,
+    },
+}
+
+
+@dataclass
+class EvalResult:
+    benchmark: str
+    n_problems: int
+    n_samples: int
+    pass_at_1: float
+    pass_at_k: float
+    k: int
+
+
+def evaluate(
+    generate_fn: Callable[[str], str],
+    benchmark: str = "humaneval",
+    limit: int | None = 20,
+    n_samples: int = 1,
+    k: int = 1,
+    timeout: float = 8.0,
+) -> EvalResult:
+    """Run pass@k. Use n_samples>1 + a sampling generate_fn for k>1."""
+    from datasets import load_dataset
+
+    spec = _BENCH[benchmark]
+    hf_ids = spec["hf_id"]
+    ds = None
+    for hf_id in hf_ids:
+        try:
+            ds = load_dataset(hf_id, split=spec["split"], trust_remote_code=True)
+            break
+        except Exception:  # noqa: BLE001
+            continue
+    if ds is None:
+        raise RuntimeError(
+            f"Could not load benchmark '{benchmark}' from any of {hf_ids}. "
+            "Check your HuggingFace token and dataset availability."
+        )
+    if limit:
+        ds = ds.select(range(min(limit, len(ds))))
+
+    p1_scores, pk_scores = [], []
+    for problem in ds:
+        intent = problem[spec["intent_col"]]
+        correct = 0
+        for _ in range(n_samples):
+            code = generate_fn(intent)
+            program = spec["program"](problem, code)
+            if run_code(program, timeout=timeout).ok:
+                correct += 1
+        p1_scores.append(pass_at_k(n_samples, correct, 1))
+        pk_scores.append(pass_at_k(n_samples, correct, k))
+
+    return EvalResult(
+        benchmark=benchmark, n_problems=len(ds), n_samples=n_samples,
+        pass_at_1=round(float(np.mean(p1_scores)), 4),
+        pass_at_k=round(float(np.mean(pk_scores)), 4), k=k,
+    )

src/eval/retrieval_eval.py

@@ -0,0 +1,109 @@
+"""Phase 3 eval: cross-modal retrieval quality (recall@k, MRR).
+
+Design: N held-out (docstring, code) pairs form a closed candidate pool.
+Each query docstring is ranked against all N code candidates; the paired
+code is the positive and the other N-1 are distractors.  This tests the
+embedder's ability to bridge natural-language intent → code, without the
+confound of looking up exact code that is already in the FAISS index.
+
+⚠️  Leakage caveat: CodeSearchNet's func_code_string includes the Python
+docstring verbatim inside the function body (the triple-quoted string right
+after `def`).  Embedding the raw code therefore lets the embedder trivially
+find the match via lexical overlap — recall@1 ≈ 0.96 is an artefact, NOT a
+measure of true code understanding.
+
+Call with strip_code_docstrings=True to remove triple-quoted strings and
+# comments from candidate code before embedding.  That number (~0.3-0.5
+recall@1) reflects the embedder's actual semantic matching ability.
+
+Usage (standalone):
+    python scripts/retrieval_only_eval.py
+"""
+from __future__ import annotations
+
+import re
+import sys
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+
+# Matches the first triple-quoted string in a Python function body.
+_TRIPLE_QUOTE_RE = re.compile(r'"""[\s\S]*?"""|\'\'\'[\s\S]*?\'\'\'')
+_COMMENT_RE = re.compile(r'#[^\n]*')
+
+
+def _strip_code_docstring(code: str) -> str:
+    """Remove the first triple-quoted docstring and all # comments from Python code."""
+    code = _TRIPLE_QUOTE_RE.sub('', code, count=1)
+    code = _COMMENT_RE.sub('', code)
+    return code
+
+
+def evaluate_cross_modal(
+    embedder,
+    pairs: pd.DataFrame,
+    k_values: tuple[int, ...] = (1, 5, 10),
+    batch_size: int = 64,
+    strip_code_docstrings: bool = False,
+) -> dict:
+    """Cross-modal retrieval eval: docstring queries → code candidates.
+
+    Args:
+        embedder:               SentenceTransformer (or anything with .encode()).
+        pairs:                  DataFrame with 'docstring' and 'code' columns (N rows).
+        k_values:               Recall cut-offs to report.
+        batch_size:             Encoding batch size.
+        strip_code_docstrings:  If True, remove triple-quoted docstrings and #
+                                comments from candidate code before embedding.
+                                Use this for a leakage-free signal; see module
+                                docstring for why the raw number is inflated.
+
+    Returns dict with keys  mrr, recall@k (for each k), n_pairs, stripped.
+    """
+    n = len(pairs)
+
+    candidates = pairs["code"].tolist()
+    if strip_code_docstrings:
+        candidates = [_strip_code_docstring(c) for c in candidates]
+
+    print(f"[eval] encoding {n} docstrings as queries ...")
+    q_emb = embedder.encode(
+        pairs["docstring"].tolist(),
+        batch_size=batch_size, show_progress_bar=True,
+        convert_to_numpy=True, normalize_embeddings=True,
+    ).astype("float32")
+
+    print(f"[eval] encoding {n} code candidates"
+          f"{' (docstrings stripped)' if strip_code_docstrings else ''} ...")
+    c_emb = embedder.encode(
+        candidates,
+        batch_size=batch_size, show_progress_bar=True,
+        convert_to_numpy=True, normalize_embeddings=True,
+    ).astype("float32")
+
+    # Cosine similarity matrix (N × N); both sides are already L2-normalised,
+    # so inner product == cosine similarity.
+    sim = q_emb @ c_emb.T  # shape (N, N)
+
+    reciprocal_ranks: list[float] = []
+    hits: dict[int, int] = {k: 0 for k in k_values}
+
+    for i in range(n):
+        order = sim[i].argsort()[::-1]
+        rank = int(np.where(order == i)[0][0]) + 1  # 1-indexed
+        reciprocal_ranks.append(1.0 / rank)
+        for k in k_values:
+            if rank <= k:
+                hits[k] += 1
+
+    result: dict = {
+        "mrr": round(float(np.mean(reciprocal_ranks)), 4),
+        "n_pairs": n,
+        "stripped": strip_code_docstrings,
+    }
+    for k in k_values:
+        result[f"recall@{k}"] = round(hits[k] / n, 4)
+    return result

src/eval/sandbox.py

@@ -0,0 +1,46 @@
+"""Run untrusted generated code in a subprocess with a timeout.
+
+NOTE: this is isolation-by-subprocess + timeout, NOT a real security sandbox.
+It protects against hangs and lets us capture tracebacks for the repair loop and
+pass@k scoring. For production / public deployment, run inside a container with
+no network and dropped privileges (see Dockerfile + README security note).
+"""
+from __future__ import annotations
+
+import subprocess
+import sys
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class ExecResult:
+    ok: bool
+    stdout: str
+    stderr: str
+    timed_out: bool = False
+
+    @property
+    def error(self) -> str:
+        """Short error summary for feeding back into a repair prompt."""
+        if self.timed_out:
+            return "Execution timed out."
+        return self.stderr.strip().splitlines()[-1] if self.stderr.strip() else ""
+
+
+def run_code(program: str, timeout: float = 8.0) -> ExecResult:
+    """Execute a full Python program string; return pass/fail + captured output."""
+    with tempfile.TemporaryDirectory() as d:
+        path = Path(d) / "candidate.py"
+        path.write_text(program)
+        try:
+            proc = subprocess.run(
+                [sys.executable, str(path)],
+                capture_output=True, text=True, timeout=timeout,
+            )
+            return ExecResult(ok=proc.returncode == 0, stdout=proc.stdout,
+                              stderr=proc.stderr)
+        except subprocess.TimeoutExpired as e:
+            return ExecResult(ok=False, stdout=e.stdout or "", stderr="Timeout",
+                              timed_out=True)

src/finetune/train_codet5.py

@@ -0,0 +1,73 @@
+"""Phase 4: fine-tune CodeT5+ on docstring -> code.
+
+This is the second experimental arm (a tuned small model) to compare against
+frozen-LLM + RAG. Runs on a single mid-range GPU; raise subset/epochs for the
+real result.
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pandas as pd
+
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+from src.config import load_config  # noqa: E402
+
+CHECKPOINT = "Salesforce/codet5p-220m"
+
+
+def finetune(subset_size: int = 5000, epochs: int = 1, out_dir: str = "data/codet5p-ft",
+             cfg=None):
+    from datasets import Dataset
+    from transformers import (AutoModelForSeq2SeqLM, AutoTokenizer,
+                              DataCollatorForSeq2Seq, Seq2SeqTrainer,
+                              Seq2SeqTrainingArguments)
+    import torch
+
+    cfg = cfg or load_config()
+    train_path = Path(cfg.paths.processed_dir) / "train.parquet"
+    df = pd.read_parquet(train_path).head(subset_size)
+
+    tok = AutoTokenizer.from_pretrained(CHECKPOINT)
+    model = AutoModelForSeq2SeqLM.from_pretrained(CHECKPOINT)
+
+    def to_features(batch):
+        x = tok(batch["docstring"], max_length=64, truncation=True, padding="max_length")
+        y = tok(text_target=batch["code"], max_length=256, truncation=True,
+                padding="max_length")
+        x["labels"] = y["input_ids"]
+        return x
+
+    ds = Dataset.from_pandas(df[["docstring", "code"]]).map(
+        to_features, batched=True, remove_columns=["docstring", "code"])
+
+    args = Seq2SeqTrainingArguments(
+        output_dir=out_dir, per_device_train_batch_size=8, num_train_epochs=epochs,
+        learning_rate=5e-5, logging_steps=50, save_strategy="epoch",
+        fp16=torch.cuda.is_available(), report_to="none")
+
+    trainer = Seq2SeqTrainer(
+        model=model, args=args, train_dataset=ds,
+        data_collator=DataCollatorForSeq2Seq(tok, model=model))
+    trainer.train()
+    trainer.save_model(out_dir)
+    tok.save_pretrained(out_dir)
+    print(f"[finetune] saved to {out_dir}")
+    return out_dir
+
+
+def make_t5_generate_fn(model_dir: str):
+    """Return generate_fn(intent)->code for plugging a tuned CodeT5+ into eval."""
+    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+    model.eval()
+
+    def generate_fn(intent: str) -> str:
+        ids = tok(intent, return_tensors="pt", truncation=True, max_length=64).input_ids
+        out = model.generate(ids.to(model.device), max_length=256)
+        return tok.decode(out[0], skip_special_tokens=True)
+
+    return generate_fn

src/rag/embedder.py

@@ -0,0 +1,85 @@
+"""Phase 3: build / save / load the FAISS retrieval index.
+
+The index plus the corpus DataFrame are persisted so deployment doesn't rebuild
+embeddings on every start (rebuilding is the slow part).
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+from src.config import load_config  # noqa: E402
+
+
+class CodeIndex:
+    """Wraps a sentence-transformer embedder + a FAISS cosine index."""
+
+    def __init__(self, embed_model: str):
+        from sentence_transformers import SentenceTransformer
+
+        self.embed_model = embed_model
+        self.embedder = SentenceTransformer(embed_model)
+        self.index = None
+        self.corpus: pd.DataFrame | None = None
+
+    def build(self, corpus: pd.DataFrame, text_col: str = "docstring", batch_size: int = 64):
+        import faiss
+
+        self.corpus = corpus.reset_index(drop=True)
+        emb = self.embedder.encode(
+            self.corpus[text_col].tolist(),
+            batch_size=batch_size, show_progress_bar=True,
+            convert_to_numpy=True, normalize_embeddings=True,
+        ).astype("float32")
+        self.index = faiss.IndexFlatIP(emb.shape[1])
+        self.index.add(emb)
+        return self
+
+    def retrieve(self, query: str, k: int = 3) -> pd.DataFrame:
+        if self.index is None or self.corpus is None:
+            raise RuntimeError("Index not built/loaded. Call build() or load().")
+        q = self.embedder.encode(
+            [query], convert_to_numpy=True, normalize_embeddings=True
+        ).astype("float32")
+        scores, idx = self.index.search(q, k)
+        out = self.corpus.iloc[idx[0]].copy()
+        out["score"] = scores[0]
+        return out
+
+    def save(self, out_dir: str):
+        import faiss
+
+        out = Path(out_dir)
+        out.mkdir(parents=True, exist_ok=True)
+        faiss.write_index(self.index, str(out / "code.index"))
+        self.corpus.to_parquet(out / "corpus.parquet", index=False)
+        (out / "embed_model.txt").write_text(self.embed_model)
+        print(f"[index] saved to {out}")
+
+    @classmethod
+    def load(cls, in_dir: str) -> "CodeIndex":
+        import faiss
+
+        in_dir = Path(in_dir)
+        embed_model = (in_dir / "embed_model.txt").read_text().strip()
+        obj = cls(embed_model)
+        obj.index = faiss.read_index(str(in_dir / "code.index"))
+        obj.corpus = pd.read_parquet(in_dir / "corpus.parquet")
+        print(f"[index] loaded {obj.index.ntotal} vectors from {in_dir}")
+        return obj
+
+
+def build_index_from_processed(cfg=None) -> CodeIndex:
+    """Build the index from data/processed/train.parquet."""
+    cfg = cfg or load_config()
+    train_path = Path(cfg.paths.processed_dir) / "train.parquet"
+    if not train_path.exists():
+        sys.exit("train.parquet missing. Run scripts/01_prepare_data.py first.")
+    corpus = pd.read_parquet(train_path)
+    idx = CodeIndex(cfg.models.embed_model).build(corpus)
+    idx.save(cfg.paths.index_dir)
+    return idx

src/rag/generator.py

@@ -0,0 +1,96 @@
+"""Phase 5: the CodeAssistant service - the heart of the deployable app.
+
+Wraps a code LLM + an optional retrieval index and exposes:
+  - generate(intent, mode="baseline"|"rag")
+  - the prompt builders, so eval/agent code can reuse them.
+
+Designed to be imported by the FastAPI / Gradio / Streamlit front-ends so all
+surfaces share one implementation.
+"""
+from __future__ import annotations
+
+import re
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+from src.config import load_config  # noqa: E402
+from src.rag.embedder import CodeIndex  # noqa: E402
+
+SYSTEM_PROMPT = (
+    "You are an expert Python coding assistant. Write a single, correct, "
+    "self-contained Python function for the request. Output only code."
+)
+
+_FENCE_RE = re.compile(r"```(?:python)?\n(.*?)```", re.DOTALL)
+
+
+def extract_code(text: str) -> str:
+    """Strip markdown fences if the model wrapped its answer."""
+    m = _FENCE_RE.search(text)
+    return m.group(1).strip() if m else text.strip()
+
+
+class CodeAssistant:
+    def __init__(self, gen_model: str, index: CodeIndex | None = None,
+                 top_k: int = 3, device_map: str = "auto"):
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        self.gen_model = gen_model
+        self.index = index
+        self.top_k = top_k
+        self.tok = AutoTokenizer.from_pretrained(gen_model)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            gen_model, dtype="auto", device_map=device_map
+        )
+
+    # ---- prompt builders ------------------------------------------------
+    def baseline_messages(self, intent: str):
+        return [{"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": f"# Task: {intent}"}]
+
+    def rag_messages(self, intent: str, k: int | None = None):
+        if self.index is None:
+            return self.baseline_messages(intent)
+        ex = self.index.retrieve(intent, k or self.top_k)
+        blocks = [f"# Task: {r.docstring}\n{r.code}" for _, r in ex.iterrows()]
+        context = "\n\n".join(blocks)
+        user = (f"Here are similar reference examples:\n\n{context}\n\n"
+                f"# Now write a function for this task:\n# Task: {intent}")
+        return [{"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": user}]
+
+    # ---- generation -----------------------------------------------------
+    def _generate(self, messages, max_new_tokens=320, temperature=0.0):
+        text = self.tok.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True)
+        inputs = self.tok(text, return_tensors="pt").to(self.model.device)
+        do_sample = temperature and temperature > 0
+        kwargs = dict(max_new_tokens=max_new_tokens, do_sample=do_sample,
+                      pad_token_id=self.tok.eos_token_id)
+        if do_sample:
+            kwargs["temperature"] = temperature
+        out = self.model.generate(**inputs, **kwargs)
+        new = out[0][inputs.input_ids.shape[1]:]
+        return self.tok.decode(new, skip_special_tokens=True)
+
+    def generate(self, intent: str, mode: str = "rag", max_new_tokens=320,
+                 temperature=0.0, return_sources=False):
+        msgs = self.rag_messages(intent) if mode == "rag" else self.baseline_messages(intent)
+        code = extract_code(self._generate(msgs, max_new_tokens, temperature))
+        if return_sources and mode == "rag" and self.index is not None:
+            srcs = self.index.retrieve(intent, self.top_k)[["docstring", "score"]]
+            return code, srcs.to_dict("records")
+        return code
+
+    @classmethod
+    def from_config(cls, cfg=None, with_index: bool = True) -> "CodeAssistant":
+        cfg = cfg or load_config()
+        index = None
+        if with_index:
+            idx_dir = Path(cfg.paths.index_dir)
+            if (idx_dir / "code.index").exists():
+                index = CodeIndex.load(str(idx_dir))
+            else:
+                print("[assistant] no saved index found; running baseline-only.")
+        return cls(cfg.models.gen_model, index=index, top_k=cfg.models.top_k)