Feat/last sprint (#12)

* hf docker build

* gradio sdk deploy for zerogpu mode

* gradio structure wip

* server gpu and tasks

* fix sdk mode

* wip studio fix gradio

* fix studio

* index fix ui

* ui coach

* ui voice

* index ui

* cleaning stuff

* clean

* clean stuff experiment

* ui css app

* fix css

.cursor/plans/gradio_sdk_deploy_58daaf6e.plan.md

@@ -0,0 +1,268 @@
+---
+name: Gradio SDK Deploy
+overview: "Add Gradio SDK deployment files on main alongside the existing Dockerfile, switch README to `sdk: gradio` for ZeroGPU Spaces, and add `@spaces.GPU` wrappers on LLM entry points so the full Studio + Classic app runs on HF without Docker."
+todos:
+  - id: root-gradio-files
+    content: Add root app.py, requirements.txt, packages.txt with editable workspace installs and Debian deps
+    status: completed
+  - id: readme-gradio-sdk
+    content: "Fix README YAML frontmatter and switch to sdk: gradio (sdk_version 6.16.0, app_file: app.py)"
+    status: completed
+  - id: spaces-runtime
+    content: Add gradio_space/spaces_runtime.py with gpu_task decorator and is_hf_gradio_runtime()
+    status: completed
+  - id: zerogpu-decorators
+    content: Apply @gpu_task to LLM entry points in model_loading, research_helpers, and tab handlers; skip preload on HF Gradio runtime in server.py
+    status: completed
+  - id: usage-docs
+    content: Update USAGE.md with Gradio SDK + ZeroGPU deploy steps; demote Docker section to later phase
+    status: completed
+  - id: local-smoke
+    content: Validate pip install + python app.py locally before pushing to HF Space
+    status: completed
+  - id: hf-space-create
+    content: Create Gradio Space under build-small-hackathon with ZeroGPU hardware and env vars; verify Studio + Classic smoke tests
+    status: cancelled
+isProject: false
+---
+
+# Gradio SDK + ZeroGPU deployment (same branch as Docker)
+
+## Goal
+
+Ship the **full app** (Studio at `/`, Classic at `/classic`, all tabs) via **Gradio SDK** on Hugging Face with **ZeroGPU**, while keeping [`Dockerfile`](Dockerfile) on `main` untouched for a later Docker Space phase.
+
+## Same-branch constraint (important)
+
+HF reads **one** `sdk:` value from root [`README.md`](README.md). Both deploy paths can live on the same branch as files, but **only one SDK is active per branch at a time**:
+
+| Files on `main` | Active when README says |
+|-----------------|-------------------------|
+| `app.py`, `requirements.txt`, `packages.txt` | `sdk: gradio` |
+| [`Dockerfile`](Dockerfile) | `sdk: docker` + `app_port: 7860` |
+
+**Phase 1 (now):** set `sdk: gradio` — Gradio Space builds from `app.py`.  
+**Phase 2 (later):** flip README to `sdk: docker` for Docker Space, or use a **second HF Space on a second branch** if you need both live at once.
+
+```mermaid
+flowchart TB
+  subgraph repo [main branch]
+    AppPy[app.py]
+    ReqTxt[requirements.txt]
+    DockerFile[Dockerfile]
+    Shared[apps/gradio-space + libs + skills]
+  end
+  subgraph phase1 [Phase 1 active]
+    ReadmeG[sdk: gradio in README]
+    HFGradio[HF Gradio SDK build]
+    ZeroGPU[ZeroGPU hardware]
+  end
+  subgraph phase2 [Phase 2 later]
+    ReadmeD[sdk: docker in README]
+    HFDocker[HF Docker build]
+    GPUBasic[GPU Basic hardware]
+  end
+  Shared --> AppPy
+  Shared --> DockerFile
+  ReadmeG --> HFGradio --> ZeroGPU
+  ReadmeD --> HFDocker --> GPUBasic
+```
+
+---
+
+## Phase 1 — Add Gradio SDK root files
+
+### 1. Root [`app.py`](app.py)
+
+Thin entry point that reuses the existing server (no UI rewrite):
+
+```python
+from gradio_space.server import main
+
+if __name__ == "__main__":
+    main()
+```
+
+HF Gradio SDK executes `app.py`; [`server.py`](apps/gradio-space/src/gradio_space/server.py) already calls `server.launch()` on port 7860 with Studio + `/classic`.
+
+### 2. Root [`requirements.txt`](requirements.txt)
+
+Pip-install workspace packages via editable paths (HF clones the full repo):
+
+```text
+-e ./libs/inference
+-e ./libs/researchmind
+-e ./libs/agent
+-e ./libs/echocoach[piper,whisper]
+-e ./apps/gradio-space
+# plus transitive deps from libs/*/pyproject.toml (torch, transformers, sentence-transformers, python-pptx, etc.)
+```
+
+Rules (per [HF Spaces dependencies](https://huggingface.co/docs/hub/spaces-dependencies)):
+
+- **Do not pin** `gradio`, `spaces`, or `huggingface_hub` — HF preinstalls them.
+- Pin heavy libs that matter for reproducibility: `torch`, `transformers`, `accelerate`, `sentence-transformers`, etc.
+- Keep `llama-cpp-python` for preset parity (HF image has `cmake`; build may be slow).
+
+Optional: add [`scripts/sync-requirements.sh`](scripts/sync-requirements.sh) later to regenerate from `pyproject.toml` files — not required for v1.
+
+### 3. Root [`packages.txt`](packages.txt)
+
+Debian deps beyond HF defaults (mirror [`Dockerfile`](Dockerfile) apt lines):
+
+```text
+ffmpeg
+libsndfile1
+```
+
+### 4. Fix + switch README frontmatter
+
+Current [`README.md`](README.md) has a blank line after `---` and still declares Docker. Update to:
+
+```yaml
+---
+title: Lesson Agent
+emoji: 📚
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: "6.16.0"
+app_file: app.py
+python_version: "3.12"
+pinned: false
+license: apache-2.0
+---
+```
+
+Remove `app_port` (Docker-only). Keep [`Dockerfile`](Dockerfile) in repo for phase 2.
+
+---
+
+## Phase 2 — ZeroGPU runtime hooks
+
+ZeroGPU requires all CUDA work inside `@spaces.GPU`. The decorator is a **no-op** locally and on dedicated GPU Spaces, so it is safe to apply everywhere.
+
+### New module: [`apps/gradio-space/src/gradio_space/spaces_runtime.py`](apps/gradio-space/src/gradio_space/spaces_runtime.py)
+
+```python
+def gpu_task(*, duration: int = 180, size: str = "large"):
+    """Apply @spaces.GPU when the HF spaces runtime is present."""
+    ...
+
+def is_hf_gradio_runtime() -> bool:
+    """True on HF Gradio SDK Spaces (skip startup model preload)."""
+    ...
+```
+
+Use `duration=180`–`300` for agent/slide flows; `duration=60` for simple chat.
+
+### Skip startup preload on HF Gradio runtime
+
+[`server.py`](apps/gradio-space/src/gradio_space/server.py) currently calls `preload_active_model()` before launch — this fails on ZeroGPU (no GPU at process start):
+
+```69:69:apps/gradio-space/src/gradio_space/server.py
+    preload_active_model()
+```
+
+Change to:
+
+```python
+if not is_hf_gradio_runtime():
+    preload_active_model()
+```
+
+First user request lazy-loads inside a `@spaces.GPU`-decorated handler.
+
+### Decorate LLM entry points (not every `backend.chat` call)
+
+Wrap **top-level handlers** so multi-step agent loops run inside one GPU allocation:
+
+| Module | Functions to decorate |
+|--------|----------------------|
+| [`model_loading.py`](apps/gradio-space/src/gradio_space/model_loading.py) | `chat`, `reload_model` |
+| [`research_helpers.py`](apps/gradio-space/src/gradio_space/research_helpers.py) | `run_research_question`, `rag_aware_chat` |
+| [`tabs/education_pptx.py`](apps/gradio-space/src/gradio_space/tabs/education_pptx.py) | `generate_lesson_slides`, `discover_lesson_sources` |
+| [`tabs/research_mind.py`](apps/gradio-space/src/gradio_space/tabs/research_mind.py) | `discover_sources`, `ask_question`, `auto_search_ingest` |
+| [`tabs/echo_coach.py`](apps/gradio-space/src/gradio_space/tabs/echo_coach.py) | `analyze_pitch` |
+| [`tabs/teacher_voice.py`](apps/gradio-space/src/gradio_space/tabs/teacher_voice.py) | text/audio turn handlers |
+
+Studio APIs in [`api/studio.py`](apps/gradio-space/src/gradio_space/api/studio.py) call these helpers — decorating the tab/helper layer avoids duplicating decorators on ~20 API wrappers.
+
+**Generator caveat:** `generate_lesson_slides` uses `yield` for progress. If ZeroGPU rejects generator handlers, extract GPU work into a plain `@gpu_task` function and keep the outer generator for UI progress only (test on Space Logs after first deploy).
+
+**Embeddings (ResearchMind ingest):** sentence-transformers can stay on CPU for v1; only LLM paths need `@spaces.GPU` initially.
+
+---
+
+## Phase 3 — Space configuration
+
+Create Space under [build-small-hackathon](https://huggingface.co/build-small-hackathon):
+
+| Setting | Value |
+|---------|-------|
+| SDK | **Gradio** (Blank template) |
+| Hardware | **ZeroGPU** (creator needs PRO/Team) |
+| Repo | GitHub `main` (or push to Space git) |
+
+**Environment variables** (Settings → Variables):
+
+| Variable | Value |
+|----------|-------|
+| `ACTIVE_MODEL` | `minicpm5-1b` |
+| `ALLOW_MODEL_SWITCH` | `false` |
+| `RESEARCHMIND_DATA_DIR` | `/tmp/researchmind` |
+
+Default preset in [`models.yaml`](models.yaml) is already `minicpm5-1b` (transformers) — good fit for ZeroGPU.
+
+---
+
+## Phase 4 — Docs and local smoke test
+
+Update [`USAGE.md`](USAGE.md):
+
+- New **Gradio SDK deployment** section (primary path): `app.py`, `requirements.txt`, ZeroGPU, env vars.
+- Move existing Docker section to **"Docker SDK (later)"** — note README must switch to `sdk: docker` + `app_port: 7860`.
+- Local Gradio SDK smoke test:
+
+```bash
+python -m venv .venv && source .venv/bin/activate
+pip install -r requirements.txt
+ACTIVE_MODEL=minicpm5-1b ALLOW_MODEL_SWITCH=false python app.py
+```
+
+Keep existing `uv run` workflow for day-to-day dev unchanged.
+
+Update [`.cursor/plans/hf_space_publish_e8a57bab.plan.md`](.cursor/plans/hf_space_publish_e8a57bab.plan.md) todos to reflect Gradio-first ordering.
+
+---
+
+## Phase 5 — Verify on Space
+
+1. **Logs** — pip install succeeds; app starts on `0.0.0.0:7860`.
+2. **`/` Studio** — loads static UI.
+3. **`/classic`** — all tabs render.
+4. **Smoke flows** — slides generation, research chat, EchoCoach sample clip, teacher voice text turn.
+5. **ZeroGPU** — first LLM request allocates GPU (may be slow on cold start); watch for "No CUDA GPUs" (means handler is outside `@spaces.GPU`).
+
+---
+
+## Phase 6 — Docker later (no code removal)
+
+When ready for Docker Space:
+
+1. Change README to `sdk: docker`, `app_port: 7860` (remove `sdk_version` / `app_file`).
+2. Create a **second Space** (or reuse after README flip) with **GPU Basic** hardware.
+3. Existing [`Dockerfile`](Dockerfile) + `uv sync` path unchanged; no `@spaces.GPU` needed on dedicated GPU.
+
+Both file sets remain on `main`; only README `sdk:` toggles which build HF runs.
+
+---
+
+## Risk notes
+
+| Risk | Mitigation |
+|------|------------|
+| `pip install llama-cpp-python` slow/fails on HF | Accept slow build; default `minicpm5-1b` avoids GGUF at runtime |
+| EchoCoach deps (piper, whisper) heavy | Full scope requested; pin versions; fix from Space Logs if needed |
+| ZeroGPU + generator slide progress | Refactor GPU block to non-generator helper if build succeeds but inference fails |
+| Two live Spaces same branch | Not supported with different SDKs — use README flip or second branch for concurrent Docker + Gradio |

.cursor/plans/hf_space_publish_e8a57bab.plan.md

@@ -0,0 +1,208 @@
+---
+name: HF Space Publish
+overview: Fix two repo blockers (README Space card YAML and missing `researchmind` in Dockerfile), validate locally with Docker, push to GitHub, then create a Docker Space under build-small-hackathon linked to GitHub with GPU hardware and MiniCPM5-1B env vars.
+todos:
+  - id: fix-readme-yaml
+    content: "Fix root README.md frontmatter: change `## title:` to `title:`"
+    status: completed
+  - id: fix-dockerfile-researchmind
+    content: Add libs/researchmind COPY lines to Dockerfile
+    status: completed
+  - id: local-docker-smoke
+    content: Run docker build + docker run locally on port 7860 with ACTIVE_MODEL=minicpm5-1b
+    status: in_progress
+  - id: push-github
+    content: Push fixed branch to GitHub repo
+    status: pending
+  - id: create-space
+    content: Create Docker Space under build-small-hackathon, link GitHub, set GPU basic
+    status: pending
+  - id: configure-env
+    content: "Set Space secrets: ACTIVE_MODEL=minicpm5-1b, ALLOW_MODEL_SWITCH=false, RESEARCHMIND_DATA_DIR=/tmp/researchmind"
+    status: pending
+  - id: verify-live
+    content: Check Space Logs, test / and /classic, confirm slide generation works
+    status: pending
+isProject: false
+---
+
+# Publish Gradio app to Hugging Face Space
+
+## Current state
+
+Your repo is **mostly ready** for a Docker Space:
+
+- Root [`Dockerfile`](Dockerfile) exposes port **7860** and runs `python -m gradio_space.app`
+- Root [`README.md`](README.md) has Space metadata (`sdk: docker`, `app_port: 7860`)
+- Default model in [`models.yaml`](models.yaml) is **`minicpm5-1b`** (transformers, `openbmb/MiniCPM5-1B`)
+
+Two issues will likely **break the Space build or card** until fixed:
+
+### Blocker 1 — README YAML is malformed
+
+The Space card frontmatter must use `title:`, not a markdown heading:
+
+```yaml
+# Current (wrong)
+## title: Lesson Agent
+
+# Required (correct)
+title: Lesson Agent
+```
+
+HF reads YAML from the **root** [`README.md`](README.md) only. Keep [`apps/gradio-space/README.md`](apps/gradio-space/README.md) as dev docs.
+
+### Blocker 2 — Dockerfile missing `researchmind`
+
+[`libs/agent`](libs/agent/pyproject.toml) depends on `researchmind`, but the Dockerfile only copies `inference`, `agent`, and `echocoach`. `uv sync` inside the image will fail without:
+
+```dockerfile
+COPY libs/researchmind/pyproject.toml libs/researchmind/README.md libs/researchmind/
+COPY libs/researchmind/src libs/researchmind/src
+```
+
+Add these lines alongside the other `libs/*` COPY blocks in [`Dockerfile`](Dockerfile).
+
+---
+
+## Architecture (what gets deployed)
+
+```mermaid
+flowchart LR
+  subgraph hf [HuggingFaceSpace]
+    DockerBuild[DockerBuild]
+    Container[Container_port7860]
+  end
+  GitHub[GitHub_repo] --> DockerBuild
+  DockerBuild --> Container
+  Container --> StudioUI["/ Studio UI"]
+  Container --> ClassicUI["/classic Gradio tabs"]
+  Container --> HubModel["Hub: openbmb/MiniCPM5-1B"]
+  HubModel --> Container
+```
+
+Entrypoint (unchanged):
+
+```44:44:Dockerfile
+CMD ["uv", "run", "--package", "gradio-space", "python", "-m", "gradio_space.app"]
+```
+
+This launches [`gradio_space.server`](apps/gradio-space/src/gradio_space/server.py): Studio at `/`, Classic tabs at `/classic`.
+
+---
+
+## Phase 1 — Fix repo files (before push)
+
+| File | Change |
+|------|--------|
+| [`README.md`](README.md) | Fix frontmatter: `title: Lesson Agent` (remove `##`); keep `sdk: docker`, `app_port: 7860` |
+| [`Dockerfile`](Dockerfile) | Add `libs/researchmind` pyproject + src COPY lines |
+
+Optional but recommended in README frontmatter (already present except title):
+
+```yaml
+---
+title: Lesson Agent
+emoji: 📚
+colorFrom: blue
+colorTo: green
+sdk: docker
+app_port: 7860
+pinned: false
+license: apache-2.0
+---
+```
+
+---
+
+## Phase 2 — Validate locally with Docker
+
+From repo root:
+
+```bash
+docker build -t hackathon-space .
+docker run --rm -p 7860:7860 \
+  -e ACTIVE_MODEL=minicpm5-1b \
+  -e ALLOW_MODEL_SWITCH=false \
+  hackathon-space
+```
+
+Open [http://localhost:7860](http://localhost:7860) (`/` Studio, `/classic` tabs). First model load downloads weights from Hub — expect several minutes on first run.
+
+If build fails, check Logs for `researchmind` or `uv sync` errors (confirms Blocker 2 fix).
+
+---
+
+## Phase 3 — Push to GitHub
+
+1. Create a GitHub repo (if not already linked)
+2. Push `main` with at minimum:
+   - `Dockerfile`, `README.md`, `pyproject.toml`, `uv.lock`
+   - `apps/gradio-space/`, `libs/`, `skills/`, `models.yaml`, `voice_models.yaml`
+
+Do **not** commit `.env`, local `models/*.gguf`, or large artifacts (`.dockerignore` already excludes these).
+
+---
+
+## Phase 4 — Create and link the Space
+
+1. Go to [build-small-hackathon](https://huggingface.co/build-small-hackathon) → **New Space**
+2. Settings:
+   - **Name:** e.g. `lesson-agent` or `small-model-hackathon`
+   - **SDK:** **Docker** (not Gradio SDK — monorepo needs root Dockerfile)
+   - **Hardware:** **GPU basic** (required for transformers `minicpm5-1b`)
+3. Under **Repository** → connect your GitHub repo and branch (`main`)
+4. HF will auto-build from root `Dockerfile` on each push
+
+---
+
+## Phase 5 — Space environment variables
+
+In Space **Settings → Variables and secrets** (Repository secrets, not `.env` in git):
+
+| Variable | Value | Why |
+|----------|-------|-----|
+| `ACTIVE_MODEL` | `minicpm5-1b` | Pins model for visitors |
+| `ALLOW_MODEL_SWITCH` | `false` | Hides dev model dropdown |
+| `AGENT_OUTPUTS_DIR` | `/tmp/agent_outputs` | Already set in Dockerfile; optional override |
+| `RESEARCHMIND_DATA_DIR` | `/tmp/researchmind` | Ephemeral RAG store on Space (recommended) |
+
+No secrets required for the default MiniCPM5 preset unless you switch to a gated model.
+
+---
+
+## Phase 6 — Verify publish
+
+1. Open Space **Logs** — wait for `Running on local URL: 0.0.0.0:7860`
+2. Open the Space URL
+3. Smoke test:
+   - `/` — Studio loads
+   - Generate slides with a simple topic (e.g. "Photosynthesis, grade 8, 5 slides")
+   - `/classic` — tabs render
+4. First inference may be slow while `openbmb/MiniCPM5-1B` downloads
+
+### Optional: faster restarts
+
+If cold starts are painful, add a **Storage Bucket** in Space settings so Hub model cache persists across restarts.
+
+---
+
+## Troubleshooting
+
+| Symptom | Fix |
+|---------|-----|
+| Space card shows wrong title / no Docker | Fix README YAML (`title:` not `## title:`) |
+| Docker build fails at `uv sync` | Add `researchmind` to Dockerfile |
+| Build OK but app crashes on Research tab | Confirm `researchmind` src is copied |
+| First request very slow | Normal — model download; use Storage Bucket |
+| OOM on GPU | Try smaller batch or switch preset to GGUF on CPU |
+
+Full reference: [`USAGE.md`](USAGE.md) sections "Docker smoke test" and "Hugging Face Space deployment".
+
+---
+
+## What you do NOT need
+
+- Plain Gradio SDK (`app.py` + `requirements.txt` at root) — wrong fit for this monorepo
+- Committing GGUF files — models download from Hub at runtime via `ACTIVE_MODEL` / `models.yaml`
+- Changing the CMD — current entrypoint already serves Studio + Classic

.env.example

@@ -66,14 +66,4 @@ ALLOW_MODEL_SWITCH=false
 # For Cohere Transcribe ASR: huggingface-cli login + accept model terms, then:
 # ECHOCOACH_ASR_PRESET=cohere-transcribe
 
-# --- Ensemble research (research/ensemble/) ---
-# Base LLM resolution (first match wins): ENSEMBLE_LLM, LLM_PATH, BASE, MODEL_ID, ACTIVE_MODEL
-# LLM_PATH=./models/finetuned/minicpm5-1b-lora-merged
-# ENSEMBLE_LLM=Qwen/Qwen2.5-0.5B-Instruct
-# ENSEMBLE_PRESET=minicpm5-1b
-# ENSEMBLE_OUT=./models/ensemble/minicpm5-1b-jepa-pretrain
-# ENSEMBLE_QA=./research/data/benchmark-qa.jsonl
-# ENSEMBLE_KB=./research/data/benchmark-kb.jsonl
-# ENSEMBLE_CKPT=./models/ensemble/jepa-lesson-pretrain
-
 BASE=openbmb/MiniCPM5-1B

.gitignore

@@ -1,4 +1,5 @@
 .venv/
+.venv-gradio/
 __pycache__/
 *.py[cod]
 .env

Dockerfile

@@ -20,11 +20,13 @@ COPY apps/gradio-space/pyproject.toml apps/gradio-space/README.md apps/gradio-sp
 COPY libs/inference/pyproject.toml libs/inference/README.md libs/inference/
 COPY libs/agent/pyproject.toml libs/agent/README.md libs/agent/
 COPY libs/echocoach/pyproject.toml libs/echocoach/README.md libs/echocoach/
+COPY libs/researchmind/pyproject.toml libs/researchmind/README.md libs/researchmind/
 COPY apps/gradio-space/src apps/gradio-space/src
 COPY apps/gradio-space/static apps/gradio-space/static
 COPY libs/inference/src libs/inference/src
 COPY libs/agent/src libs/agent/src
 COPY libs/echocoach/src libs/echocoach/src
+COPY libs/researchmind/src libs/researchmind/src
 COPY skills skills
 
 RUN useradd -m -u 1000 user && \

README.md

@@ -1,13 +1,15 @@
 ---
-
-## title: Lesson Agent
+title: Lesson Agent
 emoji: 📚
 colorFrom: blue
 colorTo: green
-sdk: docker
-app_port: 7860
+sdk: gradio
+sdk_version: "6.16.0"
+app_file: app.py
+python_version: "3.12"
 pinned: false
 license: apache-2.0
+---
 
 # Lesson Agent
 
@@ -15,7 +17,7 @@ license: apache-2.0
 
 A local skill-based agent helps a teacher you know turn a **topic + grade level** into a downloadable **PowerPoint** — powered by a small transformers model (`MiniCPM5-1B` by default), no cloud LLM API.
 
-See **[USAGE.md](USAGE.md)** for local run, Docker smoke test, and HF Space deployment.
+See **[USAGE.md](USAGE.md)** for local run, Gradio SDK / ZeroGPU Space deployment, and Docker (later).
 
 ## Prerequisites
 
@@ -59,7 +61,7 @@ libs/agent/          # Skill agent runner, tools, trace recorder
 libs/researchmind/   # Scraper, chunk/embed, MemRAG SQLite store, retrieval
 libs/inference/      # Transformers + llama.cpp backends
 skills/              # SKILL.md + references/ + scripts/ per task
-research/            # Fine-tune, ensemble experiments, agentic evals (optional)
+research/            # Fine-tune and agentic evals (optional)
 ```
 
 ### ResearchMind (offline after ingest)
@@ -87,15 +89,12 @@ See [`.env.example`](.env.example) and [`models.yaml`](models.yaml) for model pr
 
 ## Hugging Face Space deployment
 
-1. Create a Space under [build-small-hackathon](https://huggingface.co/build-small-hackathon) with **Docker** SDK.
-2. Link this repository (root `Dockerfile` + root `README.md` YAML above).
-3. Hardware: **GPU basic** recommended for transformers (`minicpm5-1b`).
-4. Optional secrets: `ACTIVE_MODEL`, `N_GPU_LAYERS` (if using GGUF preset).
+1. Create a Space under [build-small-hackathon](https://huggingface.co/build-small-hackathon) with **Gradio** SDK (Blank template).
+2. Link this repository — HF builds from root `app.py` + `requirements.txt` (README YAML above).
+3. Hardware: **ZeroGPU** for burst GPU inference, or **GPU basic** for always-on GPU.
+4. Set `ACTIVE_MODEL=minicpm5-1b`, `ALLOW_MODEL_SWITCH=false`, `RESEARCHMIND_DATA_DIR=/tmp/researchmind`.
 
-```bash
-docker build -t hackathon-space .
-docker run --rm -p 7860:7860 -e ACTIVE_MODEL=minicpm5-1b hackathon-space
-```
+A root `Dockerfile` is kept for a later **Docker SDK** deploy (flip README to `sdk: docker`). See [USAGE.md](USAGE.md).
 
 ## Hackathon checklist
 

USAGE.md

@@ -1,6 +1,6 @@
 # Usage
 
-How to run the **Lesson Agent** Gradio app locally, test it in Docker, and deploy to a Hugging Face Space for the [Build Small Hackathon](https://huggingface.co/build-small-hackathon).
+How to run the **Lesson Agent** Gradio app locally, deploy to a Hugging Face Space (Gradio SDK + ZeroGPU), and optionally test with Docker later for the [Build Small Hackathon](https://huggingface.co/build-small-hackathon).
 
 The primary UI is the **Lesson slides** tab (topic → local model outline → downloadable `.pptx`). Use **ResearchMind** for corpus Q&A, **TeacherVoice** for spoken back-and-forth tutoring, **EchoCoach** for one-shot pitch analysis, or ground lessons directly from the Lesson tab. The **Chat (debug)** tab tests the underlying model.
 
@@ -223,98 +223,121 @@ INFERENCE_BACKEND=transformers MODEL_ID=Qwen/Qwen2.5-3B-Instruct \
 
 ---
 
-## Docker (local prod-like test)
+## Gradio SDK local smoke test (matches HF Space build)
 
-Run the same container image HF Spaces will build:
+Before pushing to Hugging Face, verify the Gradio SDK entry point:
 
 ```bash
-docker build -t hackathon-space .
-docker run --rm -p 7860:7860 \
-  -e MODEL_REPO=Qwen/Qwen2.5-3B-Instruct-GGUF \
-  -e MODEL_FILE=qwen2.5-3b-instruct-q4_k_m.gguf \
-  -e N_CTX=4096 \
-  -e N_GPU_LAYERS=0 \
-  hackathon-space
+python -m venv .venv-gradio && source .venv-gradio/bin/activate
+pip install -r requirements.txt
+ACTIVE_MODEL=minicpm5-1b ALLOW_MODEL_SWITCH=false python app.py
 ```
 
-Open [http://localhost:7860](http://localhost:7860) — Studio at `/`, Classic tabs at `/classic`. Stop with `Ctrl+C`.
-
-To use a pre-downloaded local model inside Docker, mount it and set `MODEL_PATH`:
+Open [http://localhost:7860](http://localhost:7860) — Studio at `/`, Classic at `/classic`.
 
-```bash
-docker run --rm -p 7860:7860 \
-  -v "$(pwd)/models:/app/models:ro" \
-  -e MODEL_PATH=/app/models/qwen2.5-3b-instruct-q4_k_m.gguf \
-  hackathon-space
-```
+Day-to-day development can still use `uv run` (see above); this path mirrors what HF installs from `requirements.txt`.
 
 ---
 
-## Hugging Face Space deployment
+## Hugging Face Space deployment (Gradio SDK + ZeroGPU)
 
-This repo uses the **Docker SDK**. The Space card metadata lives in the YAML frontmatter at the top of [README.md](README.md).
+The Space card metadata lives in the YAML frontmatter at the top of [README.md](README.md) (`sdk: gradio`, `app_file: app.py`).
 
 ### 1. Push code to GitHub
 
-Make sure `main` (or your deploy branch) contains at minimum:
+Make sure `main` contains at minimum:
 
-- `Dockerfile`
-- `README.md` (with `sdk: docker` and `app_port: 7860`)
-- `pyproject.toml`, `uv.lock`
-- `apps/gradio-space/` and `libs/inference/`
+- `app.py`, `requirements.txt`, `packages.txt`
+- `README.md` (with `sdk: gradio`, `sdk_version`, `app_file: app.py`)
+- `models.yaml`, `skills/`
+- `apps/gradio-space/` and all `libs/*` packages
+
+The root `Dockerfile` stays in the repo for a later Docker SDK deploy (see below).
 
 ### 2. Create the Space
 
 1. Go to [build-small-hackathon](https://huggingface.co/build-small-hackathon)
 2. **New Space**
-3. Name: e.g. `small-model-hackathon`
-4. SDK: **Docker**
-5. Link your GitHub repo, or push directly to the Space repo
+3. Name: e.g. `lesson-agent` or `small-model-hackathon`
+4. SDK: **Gradio** (Blank template)
+5. Hardware: **ZeroGPU** (creator needs PRO/Team) or **GPU basic**
+6. Link your GitHub repo, or push directly to the Space git remote
 
 CLI alternative (if you have `hf` installed and org access):
 
 ```bash
 hf repo create build-small-hackathon/<your-space-name> \
   --repo-type space \
-  --space_sdk docker
+  --space_sdk gradio
 ```
 
-### 3. Configure hardware
+### 3. Set Space environment variables
+
+In the Space **Settings → Variables and secrets**:
+
+| Variable | Value |
+| -------- | ----- |
+| `ACTIVE_MODEL` | `minicpm5-1b` |
+| `ALLOW_MODEL_SWITCH` | `false` |
+| `RESEARCHMIND_DATA_DIR` | `/tmp/researchmind` |
 
+Default preset in [`models.yaml`](models.yaml) is `minicpm5-1b` (transformers) — suitable for ZeroGPU.
 
-| Setting  | Recommendation                                               |
-| -------- | ------------------------------------------------------------ |
-| Hardware | **CPU basic** to start (llama.cpp with `N_GPU_LAYERS=0`)     |
-| Upgrade  | GPU Space if you set `N_GPU_LAYERS > 0` for faster inference |
+### 4. Build and verify
 
+HF installs from `requirements.txt` and runs root `app.py`. Check the **Logs** tab for:
 
-### 4. Set Space environment variables
+- Successful pip install (first build may take several minutes — `llama-cpp-python` compiles)
+- `Running on local URL: 0.0.0.0:7860`
 
-In the Space **Settings → Variables and secrets**:
+Smoke test on the live Space:
 
+1. **`/`** — Studio UI loads
+2. **`/classic`** — all tabs render
+3. Generate slides with a simple topic (e.g. "Photosynthesis, grade 8, 5 slides")
+4. First LLM request may be slow (model download + ZeroGPU queue)
 
-| Variable            | Value                             |
-| ------------------- | --------------------------------- |
-| `INFERENCE_BACKEND` | `llama_cpp`                       |
-| `MODEL_REPO`        | `Qwen/Qwen2.5-3B-Instruct-GGUF`   |
-| `MODEL_FILE`        | `qwen2.5-3b-instruct-q4_k_m.gguf` |
-| `N_CTX`             | `4096`                            |
-| `N_GPU_LAYERS`      | `0` (or higher on GPU hardware)   |
+### 5. ZeroGPU notes
 
+LLM handlers use `@spaces.GPU` via [`gradio_space/spaces_runtime.py`](apps/gradio-space/src/gradio_space/spaces_runtime.py). If you see **No CUDA GPUs are available**, an inference path is running outside a decorated handler.
 
-### 5. Build and verify
+Startup model preload is skipped on HF Gradio runtime; the first user request loads the model inside a GPU task.
+
+### 6. Optional: persistent model cache
+
+Attach a **Storage Bucket** in Space settings so Hub model weights survive restarts.
+
+---
 
-HF builds from the root `Dockerfile` and runs:
+## Docker SDK deployment (later)
+
+Both deploy paths live on the same branch. HF reads **one** `sdk:` from README — switch to Docker when you are ready for a dedicated-GPU Space.
+
+1. Change [README.md](README.md) frontmatter to `sdk: docker`, `app_port: 7860` (remove `sdk_version` / `app_file`)
+2. Create or reconfigure a Space with **Docker** SDK and **GPU basic** hardware
+3. Set the same env vars (`ACTIVE_MODEL=minicpm5-1b`, etc.)
+
+### Local Docker smoke test
 
 ```bash
-uv run --package gradio-space python -m gradio_space.app
+docker build -t hackathon-space .
+docker run --rm -p 7860:7860 \
+  -e ACTIVE_MODEL=minicpm5-1b \
+  -e ALLOW_MODEL_SWITCH=false \
+  -e RESEARCHMIND_DATA_DIR=/tmp/researchmind \
+  hackathon-space
 ```
 
-Check the **Logs** tab while the Space builds. Once running, open the Space URL and send a test chat message. The first message may take several minutes on CPU while the GGUF downloads.
+Open [http://localhost:7860](http://localhost:7860) — Studio at `/`, Classic tabs at `/classic`. Stop with `Ctrl+C`.
 
-### 6. Optional: persistent model cache
+To use a pre-downloaded local GGUF model inside Docker, mount it and set `MODEL_PATH`:
 
-If cold starts are too slow, attach a **Storage Bucket** in Space settings so downloaded GGUF files survive restarts.
+```bash
+docker run --rm -p 7860:7860 \
+  -v "$(pwd)/models:/app/models:ro" \
+  -e MODEL_PATH=/app/models/qwen2.5-3b-instruct-q4_k_m.gguf \
+  hackathon-space
+```
 
 ---
 
@@ -323,29 +346,24 @@ If cold starts are too slow, attach a **Storage Bucket** in Space settings so do
 
 | Symptom                                  | Likely cause                      | Fix                                                                  |
 | ---------------------------------------- | --------------------------------- | -------------------------------------------------------------------- |
-| First chat hangs / slow                  | GGUF downloading from Hub         | Pre-download locally; on Space, wait or use Storage Bucket           |
-| `Failed to load model` in chat           | Wrong `MODEL_REPO` / `MODEL_FILE` | Check env vars match a valid GGUF on Hub                             |
-| Docker build fails on `llama-cpp-python` | Missing build tools               | Dockerfile already installs `build-essential` and `cmake`            |
-| Space build fails                        | Missing `uv.lock` or README YAML  | Ensure `sdk: docker` is in root `README.md` frontmatter              |
-| `transformers` backend error             | Optional deps not installed       | Run `uv sync --package inference --extra transformers`               |
-| Port already in use locally              | Another process on 7860           | `PORT=7861 uv run --package gradio-space python -m gradio_space.app` |
+| First chat hangs / slow                  | Model downloading from Hub        | Wait on Space; use Storage Bucket for cache                            |
+| `Failed to load model` in chat           | Wrong `ACTIVE_MODEL` preset       | Use `minicpm5-1b` or valid key from `models.yaml`                    |
+| Space build fails on pip install         | `llama-cpp-python` compile        | Check Logs; default preset avoids GGUF at runtime                    |
+| Space build fails                        | Malformed README YAML             | Ensure `sdk: gradio` and `app_file: app.py` in README frontmatter    |
+| No CUDA GPUs on ZeroGPU                  | Handler outside `@spaces.GPU`     | LLM entry points must use `gpu_task` in `spaces_runtime.py`          |
+| Docker build fails on `llama-cpp-python` | Missing build tools               | Dockerfile installs `build-essential` and `cmake`                    |
+| Port already in use locally              | Another process on 7860           | `PORT=7861 python app.py` or `uv run ...`                            |
 
 
 ---
 
 ## Entrypoint summary
 
-All three environments use the same command:
-
-```bash
-uv run --package gradio-space python -m gradio_space.app
-```
-
-
-| Environment | How to run                                                 |
-| ----------- | ---------------------------------------------------------- |
-| Local dev   | `uv run --package gradio-space python -m gradio_space.app` |
-| Docker      | `docker run -p 7860:7860 hackathon-space`                  |
-| HF Space    | Built and started automatically from `Dockerfile` `CMD`    |
+| Environment | How to run |
+| ----------- | ---------- |
+| Local dev (uv) | `uv run --package gradio-space python -m gradio_space.app` |
+| Local Gradio SDK smoke | `pip install -r requirements.txt && python app.py` |
+| HF Gradio Space | HF runs root `app.py` automatically |
+| Docker (later) | `docker run -p 7860:7860 hackathon-space` (after README `sdk: docker`) |
 
 

app.py

@@ -0,0 +1,6 @@
+"""Hugging Face Gradio SDK entry point (ZeroGPU / Gradio Spaces)."""
+
+from gradio_space.server import main
+
+if __name__ == "__main__":
+    main()

apps/gradio-space/src/gradio_space/model_loading.py

@@ -1,3 +1,4 @@
+from gradio_space.spaces_runtime import gpu_task
 from inference.config import get_app_config, get_model_config
 from inference.factory import get_backend, reset_backend
 from inference.response_clean import strip_reasoning_output
@@ -74,6 +75,7 @@ def warmup(model_key: str | None = None) -> str:
     )
 
 
+@gpu_task(duration=120)
 def reload_model(model_key: str) -> str:
     """Clear cached backend and reload weights for settings panel."""
     global _current_model_key
@@ -120,6 +122,7 @@ def _history_to_messages(history: list) -> list[dict[str, str]]:
     return messages
 
 
+@gpu_task(duration=60)
 def chat(message: str, history: list, model_key: str) -> str:
     load_error = ensure_model_loaded(model_key)
     if load_error:

apps/gradio-space/src/gradio_space/research_helpers.py

@@ -8,6 +8,7 @@ import gradio as gr
 from agent.models import ResearchIngestResult
 from agent.runner import AgentRunner
 from gradio_space.model_loading import chat, ensure_model_loaded, get_active_model_key
+from gradio_space.spaces_runtime import gpu_task
 from inference.factory import get_backend
 from researchmind.ingest import IngestPipeline
 
@@ -209,6 +210,7 @@ def rag_scope_hint(session_id: str, doc_ids: list[str] | None) -> str:
     return "RAG scope: **entire** indexed corpus (all sessions)."
 
 
+@gpu_task(duration=180)
 def run_research_question(
     question: str,
     *,

apps/gradio-space/src/gradio_space/server.py

@@ -12,6 +12,7 @@ from gradio import mount_gradio_app
 from gradio_space.api.studio import register_studio_apis
 from gradio_space.app import build_demo
 from gradio_space.model_loading import preload_active_model
+from gradio_space.spaces_runtime import is_hf_gradio_runtime
 from gradio_space.tabs.education_pptx import gradio_allowed_paths
 from gradio_space.tabs.echo_coach import echo_coach_allowed_paths
 from gradio_space.tabs.research_mind import researchmind_allowed_paths
@@ -66,7 +67,8 @@ def create_server() -> gr.Server:
 
 
 def main() -> None:
-    preload_active_model()
+    if not is_hf_gradio_runtime():
+        preload_active_model()
     server = create_server()
     port = int(os.environ.get("PORT", "7860"))
     server_name = os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0")

apps/gradio-space/src/gradio_space/spaces_runtime.py

@@ -0,0 +1,37 @@
+"""Hugging Face Spaces ZeroGPU helpers."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+from typing import ParamSpec, TypeVar
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+def is_hf_gradio_runtime() -> bool:
+    """True on Hugging Face Gradio SDK Spaces (skip startup model preload)."""
+    try:
+        import spaces  # noqa: F401
+    except ImportError:
+        return False
+    return bool(os.environ.get("SPACE_ID"))
+
+
+def gpu_task(
+    *,
+    duration: int = 180,
+    size: str = "large",
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Apply @spaces.GPU when the HF spaces runtime is present (no-op elsewhere)."""
+
+    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
+        try:
+            import spaces
+
+            return spaces.GPU(duration=duration, size=size)(fn)
+        except ImportError:
+            return fn
+
+    return decorator

apps/gradio-space/src/gradio_space/tabs/echo_coach.py

@@ -7,6 +7,7 @@ import gradio as gr
 from echocoach.config import get_echo_coach_config
 from echocoach.pipeline import run_echo_coach
 from gradio_space.model_loading import ensure_model_loaded, get_active_model_key
+from gradio_space.spaces_runtime import gpu_task
 from gradio_space.ui.components import (
     build_advanced_panel,
     build_recording_block,
@@ -64,6 +65,7 @@ def load_sample_pitch() -> tuple[str | None, str]:
     )
 
 
+@gpu_task(duration=180)
 def analyze_pitch(
     audio_path: str | None,
     language: str,

apps/gradio-space/src/gradio_space/tabs/education_pptx.py

@@ -16,6 +16,7 @@ from gradio_space.research_helpers import (
     resolve_session,
     resolve_topic,
 )
+from gradio_space.spaces_runtime import gpu_task
 from gradio_space.ui.components import build_advanced_panel, DOC_CHOICE_LIST_CLASSES, WorkspaceWidgets
 from inference.factory import get_backend
 from researchmind.config import get_config
@@ -158,6 +159,7 @@ def update_source_visibility(source_mode_label: str, search_workflow_label: str)
     )
 
 
+@gpu_task(duration=120)
 def discover_lesson_sources(
     topic: str,
     session_id: str,
@@ -208,6 +210,7 @@ def discover_lesson_sources(
         return msg, gr.update(choices=[], value=[]), refresh_sessions(session_id)
 
 
+@gpu_task(duration=300)
 def generate_lesson_slides(
     topic: str,
     grade: str,

apps/gradio-space/src/gradio_space/tabs/research_mind.py

@@ -23,6 +23,7 @@ from gradio_space.research_helpers import (
     run_research_question,
     trace_summary_markdown,
 )
+from gradio_space.spaces_runtime import gpu_task
 from gradio_space.ui.components import build_advanced_panel, DOC_CHOICE_LIST_CLASSES, WorkspaceWidgets
 from inference.factory import get_backend
 
@@ -35,6 +36,7 @@ def _require_topic(topic: str | None) -> str | None:
     return None
 
 
+@gpu_task(duration=120)
 def discover_sources(
     topic: str,
     session_id: str,
@@ -118,6 +120,7 @@ def discover_sources(
         )
 
 
+@gpu_task(duration=180)
 def auto_search_ingest(
     topic: str,
     session_id: str,
@@ -279,6 +282,7 @@ def ingest_selected(
         )
 
 
+@gpu_task(duration=180)
 def ask_question(
     question: str,
     session_id: str,

apps/gradio-space/src/gradio_space/tabs/teacher_voice.py

@@ -18,6 +18,7 @@ from gradio_space.research_helpers import (
     resolve_topic,
     trace_as_dict,
 )
+from gradio_space.spaces_runtime import gpu_task
 from gradio_space.tabs.research_mind import (
     auto_search_ingest,
     discover_sources,
@@ -87,6 +88,7 @@ def _turn_error(history: list | None, message: str) -> tuple:
     )
 
 
+@gpu_task(duration=180)
 def send_turn(
     audio_path: str | None,
     history: list,
@@ -142,6 +144,7 @@ def send_turn(
     return _turn_result(result)
 
 
+@gpu_task(duration=180)
 def send_text_turn(
     message: str,
     history: list,

apps/gradio-space/static/studio/index.html

@@ -291,83 +291,107 @@
     </section>
 
     <section class="col col-studio">
-      <h2 class="section-label">Step 3 · Studio Controls</h2>
-      <div class="card">
-        <p class="card-title">RAG Scope</p>
-        <label class="toggle-row">
-          <span>Cross-Reference Sources</span>
-          <input id="use-rag" type="checkbox" checked />
-        </label>
-        <p class="status-text">Session and documents use workspace defaults above unless overridden per tool.</p>
-      </div>
-      <div class="card">
-        <p class="card-title">Teacher Voice Mode</p>
-        <div class="mode-cards" id="voice-modes">
-          <button type="button" class="mode-card" data-mode="explain">Explain</button>
-          <button type="button" class="mode-card active" data-mode="lesson">Coach</button>
-          <button type="button" class="mode-card" data-mode="pitch">Practice</button>
-        </div>
-        <label class="field voice-topic-wrap" id="voice-topic-wrap">
-          <span>Focus topic</span>
-          <input id="voice-topic" type="text" class="input" placeholder="Uses workspace topic when empty" />
-        </label>
-        <details class="voice-rag-sources" id="voice-rag-sources">
-          <summary>ResearchMind sources (optional)</summary>
-          <p class="status-text">Set focus topic, then discover or ingest sources. Enable RAG above to ground answers in your library.</p>
-          <div class="ingest-action-row">
-            <button type="button" id="btn-voice-discover" class="btn btn-secondary">Discover on web</button>
-            <button type="button" id="btn-voice-auto-ingest" class="btn btn-secondary">Auto-ingest from web</button>
+      <div class="voice-layout view-voice-only">
+        <aside class="voice-rail">
+          <div class="card voice-rag-card">
+            <p class="card-title">RAG Scope</p>
+            <label class="toggle-row">
+              <span>Cross-Reference Sources</span>
+              <input id="use-rag" type="checkbox" checked />
+            </label>
+            <p class="status-text">Uses workspace session and documents unless overridden below.</p>
           </div>
-          <div id="voice-url-choices-panel" class="url-choices-panel hidden">
-            <div id="voice-url-choices-list" class="url-choices-list"></div>
+          <div class="card voice-rail-controls">
+            <p class="card-title">Mode</p>
+            <div class="mode-cards voice-mode-cards" id="voice-modes">
+              <button type="button" class="mode-card" data-mode="explain">Explain</button>
+              <button type="button" class="mode-card active" data-mode="lesson">Coach</button>
+              <button type="button" class="mode-card" data-mode="pitch">Practice</button>
+            </div>
+            <label class="field voice-topic-wrap" id="voice-topic-wrap">
+              <span>Focus topic</span>
+              <input id="voice-topic" type="text" class="input" placeholder="Uses workspace topic when empty" />
+            </label>
+            <details class="voice-rag-sources" id="voice-rag-sources">
+              <summary>Add sources (optional)</summary>
+              <p class="status-text">Discover or ingest sources to ground answers in your library.</p>
+              <div class="ingest-action-row">
+                <button type="button" id="btn-voice-discover" class="btn btn-secondary">Discover on web</button>
+                <button type="button" id="btn-voice-auto-ingest" class="btn btn-secondary">Auto-ingest</button>
+              </div>
+              <div id="voice-url-choices-panel" class="url-choices-panel hidden">
+                <div id="voice-url-choices-list" class="url-choices-list"></div>
+              </div>
+              <label class="field">
+                <span>Paste URLs (one per line)</span>
+                <textarea id="voice-urls-text" class="input" rows="2" placeholder="https://…"></textarea>
+              </label>
+              <label class="upload-zone upload-zone-compact">
+                <input id="voice-ingest-file" type="file" accept=".pdf,.docx" multiple hidden />
+                <span class="material-symbols-outlined">upload_file</span>
+                <span>Upload PDF or Doc</span>
+              </label>
+              <button type="button" id="btn-voice-ingest" class="btn btn-secondary btn-block">Ingest sources</button>
+              <p id="voice-ingest-status" class="status-text"></p>
+            </details>
           </div>
-          <label class="field">
-            <span>Paste URLs (one per line)</span>
-            <textarea id="voice-urls-text" class="input" rows="2" placeholder="https://…"></textarea>
-          </label>
-          <label class="upload-zone upload-zone-compact">
-            <input id="voice-ingest-file" type="file" accept=".pdf,.docx" multiple hidden />
-            <span class="material-symbols-outlined">upload_file</span>
-            <span>Upload PDF or Doc</span>
-          </label>
-          <button type="button" id="btn-voice-ingest" class="btn btn-secondary btn-block">Ingest sources</button>
-          <p id="voice-ingest-status" class="status-text"></p>
-        </details>
-        <div id="voice-chat-messages" class="research-chat-messages voice-chat-messages">
-          <p class="research-chat-empty">Type a message or record audio, then send.</p>
-        </div>
-        <label class="field voice-panel" id="voice-panel">
-          <span>Ask the teacher</span>
-          <textarea id="voice-message" class="input" rows="3" placeholder="What is the difference between pretraining and finetuning a small model?"></textarea>
-          <div class="recording-row">
-            <button type="button" id="btn-voice-record-start" class="btn btn-secondary">Start mic</button>
-            <button type="button" id="btn-voice-record-stop" class="btn btn-secondary" disabled>Stop mic</button>
-            <input id="voice-audio-upload" type="file" accept="audio/*" class="input input-compact" />
+        </aside>
+        <div class="voice-main">
+          <div class="card voice-main-card">
+            <div class="voice-card-head">
+              <h2 class="section-label">Teacher Voice</h2>
+              <p class="voice-card-desc">Talk with the teacher using text or voice — grounded in your sources when RAG is on.</p>
+            </div>
+            <div id="voice-chat-messages" class="research-chat-messages voice-chat-messages">
+              <p class="research-chat-empty">Type a message or record audio, then send.</p>
+            </div>
+            <div class="voice-compose" id="voice-panel">
+              <label class="field">
+                <span>Ask the teacher</span>
+                <textarea id="voice-message" class="input" rows="2" placeholder="What is the difference between pretraining and finetuning a small model?"></textarea>
+              </label>
+              <div class="voice-input-toolbar">
+                <div class="recording-row voice-recording-row">
+                  <button type="button" id="btn-voice-record-start" class="btn btn-secondary">Start mic</button>
+                  <button type="button" id="btn-voice-record-stop" class="btn btn-secondary" disabled>Stop mic</button>
+                  <input id="voice-audio-upload" type="file" accept="audio/*" class="input input-compact" />
+                </div>
+                <p id="voice-record-status" class="status-text voice-record-status"></p>
+              </div>
+              <div class="voice-send-row">
+                <button type="button" id="btn-voice-send" class="btn btn-secondary">Send text</button>
+                <button type="button" id="btn-voice-audio-send" class="btn btn-primary">Send voice turn</button>
+              </div>
+              <p id="voice-turn-status" class="status-text"></p>
+              <div class="voice-replay-row">
+                <button type="button" id="btn-voice-speak-full" class="btn btn-secondary">Speak full reply</button>
+                <button type="button" id="btn-voice-speak-quick" class="btn btn-secondary">Speak first sentence</button>
+                <button type="button" id="btn-voice-clear" class="btn btn-ghost">Clear conversation</button>
+              </div>
+              <div id="voice-audio-out" class="voice-audio-out"></div>
+            </div>
           </div>
-          <p id="voice-record-status" class="status-text"></p>
-          <button type="button" id="btn-voice-send" class="btn btn-secondary btn-block">Send text</button>
-          <button type="button" id="btn-voice-audio-send" class="btn btn-primary btn-block">Send voice turn</button>
-        </label>
-        <p id="voice-turn-status" class="status-text"></p>
-        <div class="voice-replay-row">
-          <button type="button" id="btn-voice-speak-full" class="btn btn-secondary">Speak full reply</button>
-          <button type="button" id="btn-voice-speak-quick" class="btn btn-secondary">Speak first sentence</button>
-          <button type="button" id="btn-voice-clear" class="btn btn-ghost">Clear conversation</button>
         </div>
-        <div id="voice-audio-out" class="voice-audio-out"></div>
       </div>
-      <div class="card coach-panel-wrap">
-        <h2 class="section-label">EchoCoach Feedback</h2>
-        <div class="recording-row">
-          <button type="button" id="btn-coach-record-start" class="btn btn-secondary">Start mic</button>
-          <button type="button" id="btn-coach-record-stop" class="btn btn-secondary" disabled>Stop mic</button>
+      <div class="card coach-panel-wrap view-coach-only">
+        <div class="coach-card-head">
+          <h2 class="section-label">EchoCoach · Pitch analysis</h2>
+          <p class="coach-card-desc">Record or upload a short pitch for pace, filler highlights, and spoken feedback.</p>
+        </div>
+        <div class="coach-capture-row">
+          <div class="coach-capture-controls">
+            <div class="recording-row coach-recording-row">
+              <button type="button" id="btn-coach-record-start" class="btn btn-secondary">Start mic</button>
+              <button type="button" id="btn-coach-record-stop" class="btn btn-secondary" disabled>Stop mic</button>
+              <button type="button" id="btn-coach-sample" class="btn btn-ghost">Load sample</button>
+            </div>
+            <p id="coach-record-status" class="status-text coach-record-status"></p>
+          </div>
+          <label class="field coach-upload-field">
+            <span>Upload pitch (WAV)</span>
+            <input id="coach-audio" type="file" accept="audio/*" />
+          </label>
         </div>
-        <p id="coach-record-status" class="status-text"></p>
-        <button type="button" id="btn-coach-sample" class="btn btn-ghost btn-block">Load sample clip</button>
-        <label class="field">
-          <span>Or upload pitch (WAV)</span>
-          <input id="coach-audio" type="file" accept="audio/*" />
-        </label>
         <div class="controls-grid coach-presets">
           <label class="field">
             <span>Language</span>
@@ -378,19 +402,22 @@
             <select id="coach-asr" class="input"></select>
           </label>
         </div>
-        <label class="toggle-row">
+        <label class="toggle-row coach-voiceout-toggle">
           <span>Speak full rewrite (VoiceOut)</span>
           <input id="coach-speak-rewrite" type="checkbox" />
         </label>
-        <button type="button" id="btn-analyze" class="btn btn-secondary btn-block">Analyze pitch</button>
-        <div id="coach-panel"></div>
+        <button type="button" id="btn-analyze" class="btn btn-primary btn-block coach-analyze-btn">Analyze pitch</button>
+        <div id="coach-panel" class="coach-results-panel"></div>
       </div>
     </section>
 
     <section class="col col-debug">
-      <div class="card card-tall">
-        <h2 class="section-label">Chat (debug)</h2>
-        <p class="status-text">Plain chat or corpus-grounded answers — traces appear below when RAG is on.</p>
+      <div class="card card-tall coach-debug-card">
+        <div class="coach-card-head">
+          <h2 class="section-label">Chat (debug)</h2>
+          <p class="coach-card-desc view-coach-only">Plain chat or corpus-grounded answers — traces appear below when RAG is on.</p>
+          <p class="status-text view-debug-only">Plain chat or corpus-grounded answers — traces appear below when RAG is on.</p>
+        </div>
         <label class="toggle-row">
           <span>Use ResearchMind RAG</span>
           <input id="debug-use-rag" type="checkbox" />
@@ -418,11 +445,13 @@
         <div id="debug-chat-messages" class="research-chat-messages debug-chat-messages">
           <p class="research-chat-empty">Send a message to test the active local model.</p>
         </div>
-        <label class="field">
-          <span>Message</span>
-          <textarea id="debug-message" class="input" rows="3" placeholder="Hello, model…"></textarea>
-        </label>
-        <button type="button" id="btn-debug-send" class="btn btn-primary btn-block">Send</button>
+        <div class="coach-debug-compose">
+          <label class="field">
+            <span>Message</span>
+            <textarea id="debug-message" class="input" rows="2" placeholder="Hello, model…"></textarea>
+          </label>
+          <button type="button" id="btn-debug-send" class="btn btn-primary btn-block">Send</button>
+        </div>
         <details class="studio-debug-trace" id="debug-trace-details">
           <summary>Debug trace</summary>
           <div id="debug-trace-panel"></div>

apps/gradio-space/static/studio/studio.css

@@ -960,11 +960,288 @@ body {
 
 .workspace[data-view="voice"] .col-research,
 .workspace[data-view="voice"] .col-slides { display: none; }
-.workspace[data-view="voice"] { grid-template-columns: 1fr; max-width: 520px; margin-left: auto; margin-right: auto; }
+
+.workspace[data-view="voice"] .col-debug,
+.workspace[data-view="voice"] .view-coach-only { display: none; }
+
+.view-voice-only { display: none; }
+
+.workspace[data-view="voice"] {
+  grid-template-columns: minmax(0, 1fr);
+  max-width: 1280px;
+  gap: 1.25rem;
+}
+
+.workspace[data-view="voice"] .col-studio {
+  grid-column: 1 / -1;
+  width: 100%;
+  min-width: 0;
+}
+
+.workspace[data-view="voice"] .voice-layout {
+  display: grid;
+  grid-template-columns: minmax(260px, 0.78fr) minmax(0, 1.22fr);
+  gap: 1.25rem;
+  align-items: start;
+  width: 100%;
+}
+
+.workspace[data-view="voice"] .voice-rail {
+  display: flex;
+  flex-direction: column;
+  gap: 1rem;
+  min-width: 0;
+}
+
+.workspace[data-view="voice"] .voice-main {
+  min-width: 0;
+}
+
+.workspace[data-view="voice"] .voice-main-card {
+  display: flex;
+  flex-direction: column;
+}
+
+.workspace[data-view="voice"] .voice-compose {
+  display: flex;
+  flex-direction: column;
+  gap: 0.5rem;
+}
+
+.workspace[data-view="voice"] .voice-compose .field {
+  margin: 0;
+}
+
+.workspace[data-view="voice"] .voice-compose textarea {
+  min-height: 3.25rem;
+  resize: vertical;
+}
+
+.workspace[data-view="voice"] .voice-rail .voice-mode-cards {
+  flex-direction: row;
+  flex-wrap: wrap;
+  gap: 0.35rem;
+  margin-bottom: 0.75rem;
+}
+
+.workspace[data-view="voice"] .voice-rail .voice-mode-cards .mode-card {
+  flex: 1 1 calc(33.333% - 0.35rem);
+  text-align: center;
+  justify-content: center;
+  min-width: 0;
+  padding-left: 0.5rem;
+  padding-right: 0.5rem;
+}
+
+.workspace[data-view="voice"] .voice-rail-controls .voice-topic-wrap {
+  margin: 0 0 0.75rem;
+}
+
+.workspace[data-view="voice"] .voice-rag-sources {
+  margin: 0;
+}
+
+.workspace[data-view="voice"] .voice-rag-sources summary {
+  cursor: pointer;
+  font-weight: 600;
+  font-size: 0.82rem;
+}
+
+.workspace[data-view="voice"] .voice-chat-messages {
+  min-height: 160px;
+  max-height: min(260px, 32vh);
+  margin: 0 0 0.75rem;
+}
+
+.workspace[data-view="voice"] .voice-input-toolbar {
+  padding: 0.65rem 0.75rem;
+  border: 1px solid var(--outline-variant);
+  border-radius: var(--radius-lg);
+  background: var(--surface-container-low);
+  margin-bottom: 0.65rem;
+}
+
+.workspace[data-view="voice"] .voice-recording-row {
+  margin: 0;
+}
+
+.workspace[data-view="voice"] .voice-record-status {
+  margin: 0.35rem 0 0;
+  min-height: 1.1rem;
+}
+
+.workspace[data-view="voice"] .voice-send-row {
+  display: grid;
+  grid-template-columns: 1fr 1fr;
+  gap: 0.5rem;
+  margin-bottom: 0.35rem;
+}
+
+.workspace[data-view="voice"] .voice-card-head {
+  margin-bottom: 0.85rem;
+}
+
+.workspace[data-view="voice"] .voice-card-head .section-label {
+  margin-bottom: 0.35rem;
+}
+
+.voice-card-desc {
+  margin: 0;
+  font-size: 0.84rem;
+  line-height: 1.45;
+  color: var(--secondary);
+}
+
+@media (max-width: 960px) {
+  .workspace[data-view="voice"] .voice-layout {
+    grid-template-columns: 1fr;
+    max-width: 640px;
+    margin-left: auto;
+    margin-right: auto;
+  }
+
+  .workspace[data-view="voice"] .voice-rail .voice-mode-cards {
+    flex-direction: column;
+  }
+
+  .workspace[data-view="voice"] .voice-rail .voice-mode-cards .mode-card {
+    flex: 1 1 auto;
+    text-align: left;
+    justify-content: space-between;
+  }
+
+  .workspace[data-view="voice"] .voice-send-row {
+    grid-template-columns: 1fr;
+  }
+}
 
 .workspace[data-view="coach"] .col-research,
 .workspace[data-view="coach"] .col-slides { display: none; }
-.workspace[data-view="coach"] { grid-template-columns: 1fr; max-width: 520px; margin-left: auto; margin-right: auto; }
+
+.workspace[data-view="coach"] .view-voice-only { display: none; }
+
+.workspace[data-view="slides"] .col-studio,
+.workspace[data-view="research"] .col-debug { display: none; }
+
+.workspace[data-view="coach"] {
+  grid-template-columns: minmax(0, 1.05fr) minmax(0, 0.95fr);
+  max-width: 1280px;
+  gap: 1.25rem;
+  align-items: start;
+}
+
+.workspace[data-view="coach"] .coach-panel-wrap,
+.workspace[data-view="coach"] .coach-debug-card {
+  display: flex;
+  flex-direction: column;
+}
+
+.workspace[data-view="coach"] .coach-results-panel {
+  flex: 1;
+  min-height: 120px;
+  margin-top: 0.75rem;
+  overflow-y: auto;
+}
+
+.workspace[data-view="coach"] .coach-results-panel:not(:empty) {
+  border-top: 1px solid var(--outline-variant);
+  padding-top: 0.75rem;
+}
+
+.workspace[data-view="coach"] .debug-chat-messages {
+  min-height: 140px;
+  max-height: min(240px, 30vh);
+  margin-bottom: 0.5rem;
+}
+
+.workspace[data-view="coach"] .coach-debug-compose {
+  padding-top: 0;
+  border-top: none;
+}
+
+.workspace[data-view="coach"] .coach-debug-compose textarea {
+  min-height: 3.5rem;
+  resize: vertical;
+}
+
+.workspace[data-view="coach"] .coach-debug-card .studio-debug-trace {
+  flex-shrink: 0;
+  margin-top: 0.5rem;
+}
+
+.workspace[data-view="coach"] .coach-debug-card .toggle-row,
+.workspace[data-view="coach"] .coach-debug-card .debug-rag-scope {
+  flex-shrink: 0;
+}
+
+.view-coach-only { display: none; }
+.workspace[data-view="coach"] .view-coach-only:not(.coach-panel-wrap) { display: block; }
+.workspace[data-view="coach"] .view-debug-only { display: none; }
+
+.coach-card-head {
+  margin-bottom: 0.85rem;
+}
+
+.coach-card-head .section-label {
+  margin-bottom: 0.35rem;
+}
+
+.coach-card-desc {
+  margin: 0;
+  font-size: 0.84rem;
+  line-height: 1.45;
+  color: var(--secondary);
+}
+
+.coach-capture-row {
+  display: grid;
+  grid-template-columns: minmax(0, 1.2fr) minmax(0, 1fr);
+  gap: 0.75rem;
+  align-items: start;
+  margin-bottom: 0.75rem;
+  padding: 0.75rem;
+  border: 1px solid var(--outline-variant);
+  border-radius: var(--radius-lg);
+  background: var(--surface-container-low);
+}
+
+.coach-recording-row {
+  margin: 0;
+}
+
+.coach-record-status {
+  margin: 0.35rem 0 0;
+  min-height: 1.25rem;
+}
+
+.coach-upload-field {
+  margin: 0;
+}
+
+.coach-upload-field input[type="file"] {
+  font-size: 0.78rem;
+}
+
+.coach-voiceout-toggle {
+  margin: 0.75rem 0;
+}
+
+.coach-analyze-btn {
+  margin-top: 0.25rem;
+}
+
+@media (max-width: 960px) {
+  .workspace[data-view="coach"] {
+    grid-template-columns: 1fr;
+    max-width: 640px;
+    margin-left: auto;
+    margin-right: auto;
+  }
+
+  .coach-capture-row {
+    grid-template-columns: 1fr;
+  }
+}
 
 @media (max-width: 768px) {
   :root { --sidebar-w: 0px; }
@@ -1229,9 +1506,12 @@ body {
 }
 
 .coach-presets {
-  margin-top: 0.5rem;
+  margin-top: 0;
 }
 
+.workspace[data-view="debug"] .view-coach-only { display: none; }
+.workspace[data-view="debug"] .view-debug-only { display: block; }
+
 .workspace[data-view="debug"] .col-research,
 .workspace[data-view="debug"] .col-slides,
 .workspace[data-view="debug"] .col-studio { display: none; }

models.yaml

@@ -67,9 +67,3 @@ models:
     backend: transformers
     model_id: ./models/finetuned/minicpm5-1b-lora-merged
     trust_remote_code: true
-
-  jepa-ensemble-lesson:
-    label: JEPA ensemble (LLM + emb + JEPA) lesson pretrain
-    backend: transformers
-    model_id: ./models/ensemble/jepa-lesson-pretrain
-    trust_remote_code: true

packages.txt

@@ -0,0 +1,2 @@
+ffmpeg
+libsndfile1

pyproject.toml

@@ -7,7 +7,6 @@ requires-python = ">=3.12"
 dependencies = [
     "agent",
     "echocoach",
-    "ensemble",
     "gradio-space",
     "inference",
     "researchmind",
@@ -23,11 +22,6 @@ finetune = [
     "datasets>=3.0.0",
     "peft>=0.14.0",
 ]
-ensemble = [
-    "accelerate>=1.2.0",
-    "peft>=0.14.0",
-    "transformers>=5.7.0",
-]
 evals = [
     "slm-evals",
 ]
@@ -39,14 +33,12 @@ lm-eval = [
 members = [
     "apps/*",
     "libs/*",
-    "research/ensemble",
     "research/evals",
 ]
 
 [tool.uv.sources]
 agent = { workspace = true }
 echocoach = { workspace = true }
-ensemble = { workspace = true }
 gradio-space = { workspace = true }
 inference = { workspace = true }
 researchmind = { workspace = true }

requirements.txt

@@ -0,0 +1,32 @@
+# Workspace packages (HF clones full repo)
+-e ./libs/inference
+-e ./libs/researchmind
+-e ./libs/agent
+-e ./libs/echocoach[piper,whisper]
+-e ./apps/gradio-space
+
+# Pinned runtime deps (do not pin gradio, spaces, or huggingface_hub — HF preinstalls them)
+accelerate==1.13.0
+torch==2.12.0
+torchvision==0.27.0
+transformers==5.10.2
+peft==0.19.1
+llama-cpp-python==0.3.26
+sentence-transformers==5.5.1
+pydantic>=2.0.0
+pyyaml>=6.0.2
+pillow>=10.0.0
+python-pptx>=1.0.0
+python-docx>=1.1.0
+httpx>=0.28.0
+numpy>=2.0.0
+ddgs==9.14.4
+googlesearch-python>=1.3.0
+pypdf>=5.0.0
+trafilatura==2.1.0
+matplotlib==3.11.0
+soundfile>=0.12.0
+sounddevice>=0.5.0
+librosa==0.11.0
+piper-tts==1.4.2
+pywhispercpp==1.5.0

research/README.md

@@ -1,27 +1,25 @@
 # Research
 
-Experimental code for **fine-tuning**, **ensemble architectures**, and **agentic benchmarks**. Nothing here is wired into the Gradio Lesson Agent by default — use it to train models, probe JEPA/world-model ideas, and score checkpoints against public benchmarks.
+Experimental code for **fine-tuning** and **agentic benchmarks**. Nothing here is wired into the Gradio Lesson Agent by default — use it to train models and score checkpoints against public benchmarks.
 
 | Path | Purpose |
 | ---- | ------- |
 | [`finetune.py`](finetune.py) | LoRA / QLoRA / full fine-tune on chat or instruction data |
-| [`ensemble/`](ensemble/) | JEPA + world-model ensemble experiments (uv package `ensemble`) |
 | [`evals/`](evals/) | SLM agentic benchmark suite — BFCL, τ-bench, GAIA, SWE-bench (uv package `slm-evals`) |
-| [`data/`](data/) | Shared JSONL datasets for finetune and ensemble harnesses |
+| [`data/`](data/) | Shared JSONL datasets for finetune and evals |
 
 ## Quick links
 
 - **[USAGE.md](USAGE.md)** — install groups, commands, and typical workflows
 - **[docs/overview.md](docs/overview.md)** — how the pieces fit together
-- **[ensemble/README.md](ensemble/README.md)** — ensemble smoke tests and harnesses
 - **[evals/USAGE.md](evals/USAGE.md)** — benchmark CLI, configs, and results
 - **[evals/docs/benchmarks.md](evals/docs/benchmarks.md)** — what each benchmark measures
 
 ## Install (from repo root)
 
 ```bash
-# Everything you need for research scripts
-uv sync --group finetune --group ensemble --group evals
+# All research tooling
+uv sync --group finetune --group evals --group lm-eval
 ```
 
 Individual groups:
@@ -29,8 +27,8 @@ Individual groups:
 | Group | Command | Enables |
 | ----- | ------- | ------- |
 | `finetune` | `uv sync --group finetune` | `research/finetune.py` (LoRA, QLoRA, merge) |
-| `ensemble` | `uv sync --group ensemble` | `research/ensemble/` package |
 | `evals` | `uv sync --group evals` | `research/evals/` package (`slm-benchmark`) |
+| `lm-eval` | `uv sync --group lm-eval` | `slm-lm-eval` CLI (GSM8K, ARC, HellaSwag, …) |
 
 ## Typical workflow
 
@@ -40,9 +38,7 @@ research/data/education-lesson-chat.jsonl
         ▼
   research/finetune.py  ──►  models/finetuned/<preset>-lora/
         │
-        ├──► research/evals/  (BFCL, τ-bench, GAIA, SWE-bench)
-        │
-        └──► research/ensemble/  (JEPA / world-model ablations)
+        └──► research/evals/  (BFCL, τ-bench, GAIA, SWE-bench, lm-eval)
 ```
 
 See [USAGE.md](USAGE.md) for copy-paste commands.

research/USAGE.md

@@ -1,24 +1,23 @@
 # Research usage
 
-How to run fine-tuning, ensemble experiments, and agentic benchmarks under `research/`. All commands assume the **repo root** as the working directory unless noted.
+How to run fine-tuning and agentic benchmarks under `research/`. All commands assume the **repo root** as the working directory unless noted.
 
 The Lesson Agent app lives in `apps/gradio-space/` — see root [USAGE.md](../USAGE.md). Research code is optional and isolated here.
 
 ## Prerequisites
 
 - [uv](https://docs.astral.sh/uv/) and Python 3.12
-- GPU recommended for real-model runs (CPU works for smoke tests and `tiny` backends)
+- GPU recommended for real-model runs (CPU works for smoke tests)
 - Hugging Face Hub access for model downloads and some benchmark datasets
 
 ## Install dependency groups
 
 ```bash
 # All research tooling
-uv sync --group finetune --group ensemble --group evals --group lm-eval
+uv sync --group finetune --group evals --group lm-eval
 
 # Or one at a time
 uv sync --group finetune
-uv sync --group ensemble
 uv sync --group evals
 uv sync --group lm-eval
 ```
@@ -26,7 +25,6 @@ uv sync --group lm-eval
 | Group | Package / script | What it adds |
 | ----- | ---------------- | ------------ |
 | `finetune` | `research/finetune.py` | `peft`, `datasets`, `bitsandbytes` (QLoRA) |
-| `ensemble` | `ensemble` workspace member | JEPA / world-model ensemble + harnesses |
 | `evals` | `slm-evals` workspace member | `slm-benchmark` CLI |
 | `lm-eval` | `slm-evals[lm-eval]` | `slm-lm-eval` CLI (GSM8K, ARC, HellaSwag, …) |
 
@@ -94,83 +92,7 @@ Training writes to `<out>/` (default `./models/finetuned/<preset>-<mode>/`):
 
 ---
 
-## 2. Ensemble experiments (`research/ensemble/`)
-
-JEPA and world-model ensemble prototypes: small LLM + embedding memory + latent predictors + energy-based draft selection. **Not connected to the Gradio app.**
-
-Install: `uv sync --group ensemble`
-
-### Tier 1 — CPU smoke (no Hub download)
-
-```bash
-uv run --package ensemble python -m ensemble.jepa_ensemble tiny
-uv run --package ensemble python -m ensemble.world_ensemble tiny
-bash research/ensemble/scripts/smoke.sh
-```
-
-### Tier 2 — Real small model
-
-```bash
-uv run --package ensemble python -m ensemble.jepa_ensemble Qwen/Qwen2.5-0.5B-Instruct
-uv run --package ensemble python -m ensemble.world_ensemble Qwen/Qwen2.5-0.5B-Instruct
-```
-
-### Pretrain + save (LLM + emb + JEPA)
-
-```bash
-# Default LLM: ENSEMBLE_LLM → LLM_PATH → BASE → MODEL_ID → ACTIVE_MODEL (models.yaml)
-uv run --package ensemble ensemble-pretrain --steps 200
-
-# Or override
-uv run --package ensemble ensemble-pretrain \
-  --llm Qwen/Qwen2.5-0.5B-Instruct \
-  --steps 200
-
-# Benchmark saved ensemble with slm-evals (compare to base HF model)
-uv run --package slm-evals slm-benchmark \
-  --model ./models/ensemble/jepa-lesson-pretrain \
-  --model-type ensemble \
-  --benchmarks bfcl tau_bench --max-samples 20
-```
-
-Checkpoint files: `manifest.json`, `aux.pt`, `llm/` (PEFT adapters), optional `store.pt`.
-
-### Tier 3 — Benchmark harnesses
-
-Uses `research/data/benchmark-qa.jsonl` (questions) and `benchmark-kb.jsonl` (retrieval snippets).
-
-```bash
-# JEPA track — toy
-uv run --package ensemble python -m ensemble.eval.jepa_harness \
-  --llm tiny --toy --limit 20 --n_drafts 8
-
-# JEPA track — education QA
-uv run --package ensemble python -m ensemble.eval.jepa_harness \
-  --llm Qwen/Qwen2.5-0.5B-Instruct \
-  --qa research/data/benchmark-qa.jsonl \
-  --kb research/data/benchmark-kb.jsonl \
-  --limit 50 --n_drafts 8
-
-# World-model track
-uv run --package ensemble python -m ensemble.eval.world_harness \
-  --llm tiny --toy --limit 20 --n_drafts 8
-```
-
-More detail: [ensemble/README.md](ensemble/README.md), [docs/overview.md](docs/overview.md).
-
-### Legacy shims
-
-Top-level files re-export the package for old scripts:
-
-- `research/llm_emb_jepa_ensemble_pluggable.py` → `ensemble.jepa_ensemble`
-- `research/world_model_ensemble.py` → `ensemble.world_ensemble`
-- `research/eval_harness.py` → `ensemble.eval.jepa_harness`
-
-Prefer `uv run --package ensemble python -m ensemble.<module>`.
-
----
-
-## 3. Agentic benchmarks (`research/evals/`)
+## 2. Agentic benchmarks (`research/evals/`)
 
 Evaluate local HuggingFace checkpoints on BFCL, τ-bench, GAIA, and SWE-bench Verified.
 
@@ -192,9 +114,9 @@ Full reference: [evals/USAGE.md](evals/USAGE.md).
 
 ---
 
-## 4. Academic benchmarks (`slm-lm-eval`)
+## 3. Academic benchmarks (`slm-lm-eval`)
 
-Standard lm-evaluation-harness tasks (ARC, HellaSwag, GSM8K, …) for base presets, LoRA adapters, merged checkpoints, and ensemble manifests.
+Standard lm-evaluation-harness tasks (ARC, HellaSwag, GSM8K, …) for base presets, LoRA adapters, and merged checkpoints.
 
 Install: `uv sync --group lm-eval`
 
@@ -222,12 +144,6 @@ uv run --package slm-evals slm-lm-eval \
   --preset minicpm5-1b-lesson-lora \
   --experiment-name minicpm5-1b-lora__v1 \
   --compare-to results/lm_eval/minicpm5-1b__baseline/results.json
-
-# Ensemble checkpoint
-uv run --package slm-evals slm-lm-eval \
-  --config research/evals/configs/lm_eval_smoke.yaml \
-  --model ./models/ensemble/jepa-lesson-pretrain \
-  --experiment-name ensemble-jepa__lm-eval
 ```
 
 Post-training hook:
@@ -248,8 +164,8 @@ Full reference: [evals/USAGE.md](evals/USAGE.md#lm-evaluation-harness-slm-lm-eva
 | File | Used by | Format |
 | ---- | ------- | ------ |
 | `education-lesson-chat.jsonl` | `finetune.py` default | Chat messages for lesson agent |
-| `benchmark-qa.jsonl` | Ensemble harnesses | `question`, `answer`, `domain` |
-| `benchmark-kb.jsonl` | Ensemble harnesses | Retrieval snippets for memory routing |
+| `benchmark-qa.jsonl` | Optional domain QA evals | `question`, `answer`, `domain` |
+| `benchmark-kb.jsonl` | Optional retrieval snippets | KB entries for domain QA |
 
 ---
 
@@ -283,18 +199,12 @@ Full reference: [evals/USAGE.md](evals/USAGE.md#lm-evaluation-harness-slm-lm-eva
      --compare-to results/lm_eval/minicpm5-1b__baseline/results.json
    ```
 
-5. **Optional** — probe ensemble ideas on the same QA/KB files:
-   ```bash
-   bash research/ensemble/scripts/smoke.sh
-   ```
-
 ### Verification checklist
 
 - Use the **same** lm-eval YAML (`tasks`, `num_fewshot`, `limit`, `seed`) for baseline and candidate runs.
 - Compare lm-eval `results.json` files with `--compare-to`; do not compare `training_results.json` `result_score` to lm-eval accuracy.
 - For LoRA checkpoints, prefer `--preset minicpm5-1b-lesson-lora` (base + adapter) over passing the adapter dir alone to `--model`.
 - Report mean ± std only after multiple training seeds; single-seed deltas are indicative, not conclusive.
-- Ensemble `loglikelihood` tasks score the underlying LLM head; generative tasks (`gsm8k`) use the full JEPA+RAG stack.
 
 ---
 
@@ -302,7 +212,6 @@ Full reference: [evals/USAGE.md](evals/USAGE.md#lm-evaluation-harness-slm-lm-eva
 
 | Symptom | Fix |
 | ------- | --- |
-| `No module named 'ensemble'` | `uv sync --group ensemble` |
 | `slm-benchmark: command not found` | `uv sync --group evals` |
 | `slm-lm-eval: command not found` | `uv sync --group lm-eval` |
 | CUDA OOM during finetune | Use `--mode qlora` or reduce batch size in script args |

research/docs/overview.md

@@ -13,13 +13,12 @@ small-model-hackathon/
 └── research/              ← experiments (this tree)
     ├── finetune.py
     ├── data/
-    ├── ensemble/          ← uv workspace package
     └── evals/             ← uv workspace package
 ```
 
-Research code is a **uv workspace sibling** of `apps/*` and `libs/*`. Root `pyproject.toml` declares optional dependency groups (`finetune`, `ensemble`, `evals`) so the Docker Space image does not need to install torch-heavy extras unless you opt in locally.
+Research code is a **uv workspace sibling** of `apps/*` and `libs/*`. Root `pyproject.toml` declares optional dependency groups (`finetune`, `evals`, `lm-eval`) so the Docker Space image does not need to install torch-heavy extras unless you opt in locally.
 
-## Three tracks
+## Two tracks
 
 ### Fine-tuning
 
@@ -27,38 +26,12 @@ Research code is a **uv workspace sibling** of `apps/*` and `libs/*`. Root `pypr
 
 Outputs land in `models/finetuned/` — you can register a new preset in `models.yaml` pointing at merged weights for the **Well-Tuned** hackathon badge.
 
-### Ensemble (JEPA / world model)
+### Agentic and academic evals
 
-`research/ensemble/` explores a modular stack inspired by LeCun-style architectures:
+`research/evals/` (`slm-evals` package) scores **whole models** on:
 
-```text
-Input ──► Embedder + VectorStore (retrieval memory)
-              │
-              ▼
-         JEPA encoder ──► latent state
-              │
-              ├──► World model (multi-step latent rollout)
-              │
-              └──► Energy model (scores LLM draft continuations)
-                        │
-                        ▼
-              Small LLM generates N drafts → pick lowest energy
-```
-
-Two entry ensembles:
-
-| Module | File | Critic |
-| ------ | ---- | ------ |
-| JEPA track | `ensemble.jepa_ensemble` | JEPA latent prediction |
-| World track | `ensemble.world_ensemble` | Energy model over world-model rollouts |
-
-`TinyBackend` runs on CPU with random weights for smoke tests. `HFBackend` loads real Hub models via `transformers` + optional `peft` LoRA banks.
-
-Eval harnesses (`ensemble.eval.jepa_harness`, `ensemble.eval.world_harness`) measure draft-selection accuracy on `research/data/benchmark-qa.jsonl` with optional KB retrieval from `benchmark-kb.jsonl`.
-
-### Agentic evals
-
-`research/evals/` (`slm-evals` package) scores **whole models** on public agent benchmarks — function calling, multi-turn tool use, GAIA tasks, and SWE-bench patches. This complements ensemble harnesses: evals test end-to-end model behavior; ensemble harnesses test internal selection mechanisms on a small custom QA set.
+- **Agentic benchmarks** — BFCL, τ-bench, GAIA, SWE-bench (`slm-benchmark`)
+- **Academic benchmarks** — GSM8K, ARC, HellaSwag, etc. via lm-evaluation-harness (`slm-lm-eval`)
 
 ## Data flow
 
@@ -79,20 +52,12 @@ flowchart LR
     tau[tau-bench]
     gaia[GAIA]
     swe[SWE-bench]
-  end
-
-  subgraph ens [ensemble]
-    jepa[JEPA harness]
-    world[World harness]
+    lmeval[lm-eval tasks]
   end
 
   lesson --> train
   train --> ckpt
   ckpt --> evals
-  qa --> jepa
-  kb --> jepa
-  qa --> world
-  kb --> world
 ```
 
 ## When to use which tool
@@ -101,14 +66,11 @@ flowchart LR
 | ---- | ---- |
 | Improve lesson slide quality on your data | `finetune.py` + optional eval before/after |
 | Compare base vs LoRA on public agent tasks | `slm-benchmark` |
-| Prototype latent draft selection | `ensemble` smoke → harness |
+| Compare base vs LoRA on academic tasks | `slm-lm-eval` |
 | Ship in Gradio Space | `apps/gradio-space` only — wire new weights via `models.yaml` |
 
-## Workspace packages
-
-Both subpackages are listed in root `[tool.uv.workspace] members`:
+## Workspace package
 
-- `research/ensemble` → import name `ensemble`
-- `research/evals` → import name `slm_evals`, CLI `slm-benchmark`
+`research/evals` is listed in root `[tool.uv.workspace] members` as import name `slm_evals`, CLI `slm-benchmark` and `slm-lm-eval`.
 
-Run with `uv run --package <name>` from the repo root so uv resolves workspace paths and shared lockfile versions.
+Run with `uv run --package slm-evals ...` from the repo root so uv resolves workspace paths and shared lockfile versions.

research/ensemble/README.md

@@ -1,113 +0,0 @@
-# Ensemble research package
-
-JEPA and world-model ensemble experiments. Stays under `research/` — not wired into the Gradio agent.
-
-See also: [../USAGE.md](../USAGE.md) · [../docs/overview.md](../docs/overview.md)
-
-## Install
-
-```bash
-uv sync --group ensemble
-```
-
-## Tier 1 — Smoke (CPU, no HF download)
-
-```bash
-uv run --package ensemble python -m ensemble.jepa_ensemble tiny
-uv run --package ensemble python -m ensemble.world_ensemble tiny
-bash research/ensemble/scripts/smoke.sh
-```
-
-## Tier 2 — Micro demo (real small model)
-
-```bash
-uv run --package ensemble python -m ensemble.jepa_ensemble Qwen/Qwen2.5-0.5B-Instruct
-uv run --package ensemble python -m ensemble.world_ensemble Qwen/Qwen2.5-0.5B-Instruct
-```
-
-## Pretrain + save (LLM + emb + JEPA)
-
-Joint training writes a full checkpoint to `models/ensemble/<name>/`:
-
-```bash
-# CPU smoke (tiny backend, no HF download)
-uv run --package ensemble ensemble-pretrain \
-  --llm tiny --steps 50 --no-kb \
-  --out models/ensemble/jepa-smoke
-
-# Uses ACTIVE_MODEL / BASE / LLM_PATH from .env + models.yaml by default
-uv run --package ensemble ensemble-pretrain \
-  --data research/data/education-lesson-chat.jsonl \
-  --kb research/data/benchmark-kb.jsonl \
-  --steps 200
-
-# Override base LLM explicitly
-uv run --package ensemble ensemble-pretrain \
-  --llm Qwen/Qwen2.5-0.5B-Instruct --steps 200
-```
-
-Checkpoint layout: `manifest.json`, `aux.pt` (emb/jepa/bridge/router), `llm/` (PEFT adapters).
-
-Benchmark the saved ensemble with **slm-evals** (auto-detects `manifest.json`):
-
-```bash
-uv run --package slm-evals slm-benchmark \
-  --model ./models/ensemble/jepa-lesson-pretrain \
-  --model-type ensemble \
-  --benchmarks bfcl tau_bench --max-samples 20
-
-# Or use the template config
-uv run --package slm-evals slm-benchmark \
-  --config research/evals/configs/ensemble_jepa_lesson.yaml
-```
-
-Compare against a base HF model by running the same config with `model_type: hf` and `model_path: openbmb/MiniCPM5-1B`.
-
-## Tier 3 — Benchmark
-
-### JEPA ablation ladder
-
-```bash
-# Toy (no download)
-uv run --package ensemble python -m ensemble.eval.jepa_harness \
-  --llm tiny --toy --limit 20 --n_drafts 8
-
-# Education QA set
-uv run --package ensemble python -m ensemble.eval.jepa_harness \
-  --llm Qwen/Qwen2.5-0.5B-Instruct \
-  --qa research/data/benchmark-qa.jsonl \
-  --kb research/data/benchmark-kb.jsonl \
-  --limit 50 --n_drafts 8
-```
-
-### World-model energy selector
-
-```bash
-uv run --package ensemble python -m ensemble.eval.world_harness \
-  --llm tiny --toy --limit 20 --n_drafts 8
-
-uv run --package ensemble python -m ensemble.eval.world_harness \
-  --llm Qwen/Qwen2.5-0.5B-Instruct \
-  --qa research/data/benchmark-qa.jsonl \
-  --kb research/data/benchmark-kb.jsonl \
-  --limit 50 --n_drafts 8
-```
-
-## Layout
-
-```
-research/ensemble/
-  src/ensemble/
-    backends.py       # TinyBackend, HFBackend, TinyLLM, HFLLM
-    memory.py         # Embedder, VectorStore, Router
-    jepa.py           # JEPA latent predictor
-    bridge.py         # LLM hidden -> latent alignment
-    world_model.py    # Latent dynamics + rollout
-    energy.py         # Energy-based critic
-    jepa_ensemble.py  # Ensemble (JEPA track)
-    world_ensemble.py # WorldEnsemble
-    eval/
-      metrics.py
-      jepa_harness.py
-      world_harness.py
-```

research/ensemble/pyproject.toml

@@ -1,16 +0,0 @@
-[project]
-name = "ensemble"
-version = "0.1.0"
-description = "JEPA and world-model ensemble research package"
-readme = "README.md"
-requires-python = ">=3.12"
-dependencies = [
-    "torch>=2.5.0",
-]
-
-[project.scripts]
-ensemble-pretrain = "ensemble.pretrain:main"
-
-[build-system]
-requires = ["uv_build>=0.8.13,<0.9.0"]
-build-backend = "uv_build"

research/ensemble/scripts/smoke.sh

@@ -1,35 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-ROOT="$(cd "$(dirname "$0")/../../.." && pwd)"
-cd "$ROOT"
-
-echo "== JEPA ensemble demo (tiny) =="
-uv run --package ensemble python -m ensemble.jepa_ensemble tiny
-
-echo ""
-echo "== World ensemble demo (tiny) =="
-uv run --package ensemble python -m ensemble.world_ensemble tiny
-
-echo ""
-echo "== JEPA harness (toy) =="
-uv run --package ensemble python -m ensemble.eval.jepa_harness \
-  --llm tiny --toy --limit 10 --n_drafts 4
-
-echo "== Pretrain smoke + checkpoint roundtrip =="
-uv run --package ensemble ensemble-pretrain \
-  --llm tiny --steps 20 --no-kb \
-  --out models/ensemble/jepa-smoke
-uv run --package ensemble python -c "
-from ensemble.checkpoint import load_checkpoint
-ens = load_checkpoint('models/ensemble/jepa-smoke')
-print('loaded ensemble, adapters:', ens.adapter_names)
-"
-
-echo ""
-echo "== World harness (toy) =="
-uv run --package ensemble python -m ensemble.eval.world_harness \
-  --llm tiny --toy --limit 10 --n_drafts 4
-
-echo ""
-echo "All smoke checks passed."

research/ensemble/src/ensemble/__init__.py

@@ -1,15 +0,0 @@
-"""Research ensemble package: JEPA and world-model tracks."""
-
-__all__ = ["Ensemble", "WorldEnsemble"]
-
-
-def __getattr__(name: str):
-    if name == "Ensemble":
-        from ensemble.jepa_ensemble import Ensemble
-
-        return Ensemble
-    if name == "WorldEnsemble":
-        from ensemble.world_ensemble import WorldEnsemble
-
-        return WorldEnsemble
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

research/ensemble/src/ensemble/backends.py

@@ -1,418 +0,0 @@
-"""LLM backends: toy fallbacks and HuggingFace + LoRA loaders."""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class LLMBackend(nn.Module):
-    """Contract for JEPA ensemble backends."""
-
-    vocab_size: int
-    hidden_size: int
-
-
-class HFBackend(LLMBackend):
-    """HuggingFace causal LM with PEFT LoRA adapter bank."""
-
-    def __init__(
-        self,
-        model_path: str,
-        *,
-        load_in_4bit: bool = False,
-        lora_r: int = 16,
-        lora_alpha: int = 32,
-        target_modules=("q_proj", "v_proj"),
-        device: str | None = None,
-        torch_dtype=None,
-    ):
-        super().__init__()
-        from peft import LoraConfig, get_peft_model
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-
-        self.device_ = torch.device(
-            device or ("cuda" if torch.cuda.is_available() else "cpu")
-        )
-
-        kwargs = {}
-        if load_in_4bit:
-            from transformers import BitsAndBytesConfig
-
-            kwargs["quantization_config"] = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_compute_dtype=torch.bfloat16,
-                bnb_4bit_quant_type="nf4",
-            )
-        if torch_dtype is not None:
-            kwargs["torch_dtype"] = torch_dtype
-
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-        base = AutoModelForCausalLM.from_pretrained(model_path, **kwargs)
-        if not load_in_4bit:
-            base.to(self.device_)
-
-        for p in base.parameters():
-            p.requires_grad_(False)
-
-        self._lora_cfg = LoraConfig(
-            r=lora_r,
-            lora_alpha=lora_alpha,
-            lora_dropout=0.05,
-            target_modules=list(target_modules),
-            task_type="CAUSAL_LM",
-        )
-        self.model = get_peft_model(base, self._lora_cfg, adapter_name="general")
-        self._adapters = {"general"}
-
-        self.vocab_size = self.model.config.vocab_size
-        self.hidden_size = self.model.config.hidden_size
-
-    def add_adapter(self, name: str):
-        if name not in self._adapters:
-            self.model.add_adapter(name, self._lora_cfg)
-            self._adapters.add(name)
-
-    def set_adapter(self, name: str):
-        self.model.set_adapter(name)
-
-    def trainable_parameters(self):
-        return (p for p in self.model.parameters() if p.requires_grad)
-
-    def forward(self, ids):
-        out = self.model(
-            input_ids=ids.to(self.device_), output_hidden_states=True
-        )
-        return out.logits, out.hidden_states[-1]
-
-    @torch.no_grad()
-    def generate(self, ids, n_new=64, temperature=0.8):
-        gen_kwargs: dict = dict(
-            input_ids=ids.to(self.device_),
-            max_new_tokens=n_new,
-            pad_token_id=self.tokenizer.pad_token_id,
-        )
-        if temperature <= 0:
-            gen_kwargs["do_sample"] = False
-        else:
-            gen_kwargs.update(do_sample=True, temperature=temperature)
-        out = self.model.generate(**gen_kwargs)
-        return out
-
-    def encode_text(self, text: str):
-        return self.tokenizer(text, return_tensors="pt").input_ids.to(self.device_)
-
-    def decode(self, ids):
-        return self.tokenizer.decode(ids[0], skip_special_tokens=True)
-
-    @property
-    def device(self):
-        return self.device_
-
-
-class TinyBackend(LLMBackend):
-    """Toy transformer with LoRA adapters (no transformers dependency)."""
-
-    VOCAB, D_MODEL, N_LAYERS, N_HEADS, SEQ_LEN, LORA_R = 1000, 128, 2, 4, 32, 8
-
-    class _LoRALinear(nn.Module):
-        def __init__(self, d_in, d_out, r):
-            super().__init__()
-            self.base = nn.Linear(d_in, d_out)
-            self.base.weight.requires_grad_(False)
-            self.base.bias.requires_grad_(False)
-            self.adapters, self.active, self.r = nn.ModuleDict(), None, r
-
-        def add_adapter(self, name):
-            A = nn.Linear(self.base.in_features, self.r, bias=False)
-            B = nn.Linear(self.r, self.base.out_features, bias=False)
-            nn.init.zeros_(B.weight)
-            self.adapters[name] = nn.Sequential(A, B)
-
-        def forward(self, x):
-            y = self.base(x)
-            if self.active and self.active in self.adapters:
-                y = y + self.adapters[self.active](x)
-            return y
-
-    class _Block(nn.Module):
-        def __init__(self, D, H, R):
-            super().__init__()
-            L = TinyBackend._LoRALinear
-            self.ln1 = nn.LayerNorm(D)
-            self.attn = nn.MultiheadAttention(D, H, batch_first=True)
-            self.ln2 = nn.LayerNorm(D)
-            self.up, self.down = L(D, 4 * D, R), L(4 * D, D, R)
-
-        def forward(self, x, mask):
-            h = self.ln1(x)
-            a, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
-            x = x + a
-            return x + self.down(F.gelu(self.up(self.ln2(x))))
-
-    def __init__(self):
-        super().__init__()
-        D, V = self.D_MODEL, self.VOCAB
-        self.tok = nn.Embedding(V, D)
-        self.pos = nn.Embedding(self.SEQ_LEN * 4, D)
-        self.blocks = nn.ModuleList(
-            [self._Block(D, self.N_HEADS, self.LORA_R) for _ in range(self.N_LAYERS)]
-        )
-        self.ln_f, self.head = nn.LayerNorm(D), nn.Linear(D, V, bias=False)
-        self.vocab_size, self.hidden_size = V, D
-        self.add_adapter("general")
-        self.set_adapter("general")
-
-    def add_adapter(self, name):
-        for b in self.blocks:
-            b.up.add_adapter(name)
-            b.down.add_adapter(name)
-
-    def set_adapter(self, name):
-        for b in self.blocks:
-            b.up.active = name
-            b.down.active = name
-
-    def trainable_parameters(self):
-        return (p for p in self.parameters() if p.requires_grad)
-
-    def forward(self, ids):
-        B, T = ids.shape
-        x = self.tok(ids) + self.pos(torch.arange(T, device=ids.device))
-        mask = torch.triu(
-            torch.full((T, T), float("-inf"), device=ids.device), 1
-        )
-        for b in self.blocks:
-            x = b(x, mask)
-        h = self.ln_f(x)
-        return self.head(h), h
-
-    @torch.no_grad()
-    def generate(self, ids, n_new=16, temperature=1.0):
-        for _ in range(n_new):
-            logits, _ = self(ids[:, -self.SEQ_LEN :])
-            if temperature <= 0:
-                nxt = logits[:, -1].argmax(dim=-1, keepdim=True)
-            else:
-                nxt = torch.multinomial(
-                    F.softmax(logits[:, -1] / temperature, -1), 1
-                )
-            ids = torch.cat([ids, nxt], dim=1)
-        return ids
-
-    def encode_text(self, text: str):
-        vals = [ord(c) % self.vocab_size for c in text[: self.SEQ_LEN]]
-        if not vals:
-            vals = [0]
-        return torch.tensor([vals], dtype=torch.long)
-
-    def decode(self, ids):
-        return " ".join(str(int(t)) for t in ids[0].tolist())
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
-
-def make_backend(llm: str, **kw) -> LLMBackend:
-    """'tiny' -> toy model; anything else -> HF hub id or local path."""
-    return TinyBackend() if llm == "tiny" else HFBackend(llm, **kw)
-
-
-def load_hf_backend_from_checkpoint(
-    base_llm: str,
-    adapter_dir: str | None,
-    *,
-    adapter_names: tuple[str, ...] = ("general",),
-    device: str | None = None,
-    load_in_4bit: bool = False,
-    lora_r: int = 16,
-    lora_alpha: int = 32,
-) -> HFBackend:
-    """Load a frozen base LM + saved PEFT adapters (ensemble checkpoint llm/)."""
-    from pathlib import Path
-
-    from peft import LoraConfig, PeftModel, get_peft_model
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-
-    def _discover_adapter_dirs(root: Path) -> dict[str, Path]:
-        if (root / "adapter_config.json").is_file():
-            return {"general": root}
-        discovered: dict[str, Path] = {}
-        for child in sorted(root.iterdir()):
-            if child.is_dir() and (child / "adapter_config.json").is_file():
-                discovered[child.name] = child
-        return discovered
-
-    resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-    tokenizer = AutoTokenizer.from_pretrained(adapter_dir or base_llm)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    kwargs: dict = {}
-    if load_in_4bit:
-        from transformers import BitsAndBytesConfig
-
-        kwargs["quantization_config"] = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_quant_type="nf4",
-        )
-    elif resolved_device != "cpu":
-        kwargs["torch_dtype"] = torch.bfloat16
-
-    base = AutoModelForCausalLM.from_pretrained(base_llm, **kwargs)
-    if not load_in_4bit and resolved_device != "cpu":
-        base.to(resolved_device)
-    for p in base.parameters():
-        p.requires_grad_(False)
-
-    if adapter_dir:
-        adapter_dirs = _discover_adapter_dirs(Path(adapter_dir))
-        if not adapter_dirs:
-            raise ValueError(
-                f"No PEFT adapters found under {adapter_dir} "
-                "(expected adapter_config.json or <name>/adapter_config.json)"
-            )
-        preferred = [name for name in adapter_names if name in adapter_dirs]
-        load_order = preferred + [
-            name for name in adapter_dirs if name not in preferred
-        ]
-        first_name = load_order[0]
-        model = PeftModel.from_pretrained(
-            base,
-            str(adapter_dirs[first_name]),
-            adapter_name=first_name,
-            is_trainable=False,
-        )
-        for name in load_order[1:]:
-            model.load_adapter(str(adapter_dirs[name]), adapter_name=name)
-        adapters = set(load_order)
-    else:
-        lora_cfg = LoraConfig(
-            r=lora_r,
-            lora_alpha=lora_alpha,
-            lora_dropout=0.05,
-            target_modules=["q_proj", "v_proj"],
-            task_type="CAUSAL_LM",
-        )
-        model = get_peft_model(base, lora_cfg, adapter_name="general")
-        adapters = {"general"}
-
-    backend = HFBackend.__new__(HFBackend)
-    nn.Module.__init__(backend)
-    backend.device_ = torch.device(resolved_device)
-    backend.tokenizer = tokenizer
-    backend.model = model
-    backend._lora_cfg = None
-    backend._adapters = adapters
-    backend.vocab_size = model.config.vocab_size
-    backend.hidden_size = model.config.hidden_size
-    if adapter_names:
-        backend.set_adapter(adapter_names[0])
-    return backend
-
-
-class TinyLLM(nn.Module):
-    """Simpler toy LLM for the world-model track (no adapter bank)."""
-
-    VOCAB, D, L, H, T = 1000, 128, 2, 4, 32
-
-    def __init__(self):
-        super().__init__()
-        self.tok = nn.Embedding(self.VOCAB, self.D)
-        self.pos = nn.Embedding(self.T * 4, self.D)
-        layer = nn.TransformerEncoderLayer(
-            self.D, self.H, 4 * self.D, batch_first=True, norm_first=True
-        )
-        self.blocks = nn.TransformerEncoder(layer, self.L)
-        self.head = nn.Linear(self.D, self.VOCAB, bias=False)
-        self.vocab_size, self.hidden_size = self.VOCAB, self.D
-
-    def forward(self, ids):
-        Tn = ids.size(1)
-        x = self.tok(ids) + self.pos(torch.arange(Tn, device=ids.device))
-        mask = torch.triu(
-            torch.full((Tn, Tn), float("-inf"), device=ids.device), 1
-        )
-        h = self.blocks(x, mask=mask)
-        return self.head(h), h
-
-    @torch.no_grad()
-    def generate(self, ids, n_new=16, temperature=1.0):
-        for _ in range(n_new):
-            logits, _ = self(ids[:, -self.T :])
-            nxt = torch.multinomial(
-                F.softmax(logits[:, -1] / temperature, -1), 1
-            )
-            ids = torch.cat([ids, nxt], 1)
-        return ids
-
-    def trainable_parameters(self):
-        return self.parameters()
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
-
-class HFLLM(nn.Module):
-    """Small HF model with single LoRA stack (world-model track)."""
-
-    def __init__(self, path, lora_r=16):
-        super().__init__()
-        from peft import LoraConfig, get_peft_model
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-        base = AutoModelForCausalLM.from_pretrained(
-            path,
-            torch_dtype=torch.bfloat16
-            if torch.cuda.is_available()
-            else torch.float32,
-            device_map="auto" if torch.cuda.is_available() else None,
-        )
-        for p in base.parameters():
-            p.requires_grad_(False)
-        cfg = LoraConfig(
-            r=lora_r,
-            lora_alpha=2 * lora_r,
-            lora_dropout=0.05,
-            target_modules=["q_proj", "v_proj"],
-            task_type="CAUSAL_LM",
-        )
-        self.model = get_peft_model(base, cfg)
-        self.vocab_size = self.model.config.vocab_size
-        self.hidden_size = self.model.config.hidden_size
-
-    def forward(self, ids):
-        out = self.model(
-            input_ids=ids.to(self.device), output_hidden_states=True
-        )
-        return out.logits, out.hidden_states[-1]
-
-    @torch.no_grad()
-    def generate(self, ids, n_new=32, temperature=0.8):
-        return self.model.generate(
-            input_ids=ids.to(self.device),
-            max_new_tokens=n_new,
-            do_sample=True,
-            temperature=temperature,
-            pad_token_id=self.tokenizer.pad_token_id,
-        )
-
-    def trainable_parameters(self):
-        return (p for p in self.model.parameters() if p.requires_grad)
-
-    @property
-    def device(self):
-        return next(self.model.parameters()).device
-
-
-def load_llm(spec: str):
-    return TinyLLM() if spec == "tiny" else HFLLM(spec)

research/ensemble/src/ensemble/bridge.py

@@ -1,28 +0,0 @@
-"""Bridge: align LLM hidden states with JEPA latent space."""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class Bridge(nn.Module):
-    def __init__(self, d_llm_hidden: int, d_latent: int):
-        super().__init__()
-        self.proj = nn.Sequential(
-            nn.Linear(d_llm_hidden, d_latent),
-            nn.GELU(),
-            nn.Linear(d_latent, d_latent),
-        )
-
-    def forward(self, llm_hidden):
-        return self.proj(llm_hidden.float().mean(dim=1))
-
-    def info_nce(self, z1, z2, tau=0.07):
-        z1, z2 = F.normalize(z1, dim=-1), F.normalize(z2, dim=-1)
-        logits = z1 @ z2.t() / tau
-        labels = torch.arange(z1.size(0), device=z1.device)
-        return 0.5 * (
-            F.cross_entropy(logits, labels) + F.cross_entropy(logits.t(), labels)
-        )

research/ensemble/src/ensemble/checkpoint.py

@@ -1,149 +0,0 @@
-"""Save and load JEPA ensemble checkpoints under models/ensemble/."""
-
-from __future__ import annotations
-
-import json
-from pathlib import Path
-from typing import Any
-
-import torch
-
-from ensemble.backends import TinyBackend, load_hf_backend_from_checkpoint
-from ensemble.jepa_ensemble import Ensemble
-
-MANIFEST_FILE = "manifest.json"
-AUX_FILE = "aux.pt"
-STORE_FILE = "store.pt"
-LLM_DIR = "llm"
-TINY_LLM_FILE = "tiny_llm.pt"
-
-CHECKPOINT_VERSION = 1
-
-
-def _aux_state_dict(ens: Ensemble) -> dict[str, torch.Tensor]:
-    return {
-        "emb": ens.emb.state_dict(),
-        "jepa": ens.jepa.state_dict(),
-        "bridge": ens.bridge.state_dict(),
-        "router": ens.router.state_dict(),
-    }
-
-
-def _store_payload(ens: Ensemble) -> dict[str, Any]:
-    return {
-        "keys": [k for k in ens.store.keys],
-        "values": [v for v in ens.store.values],
-    }
-
-
-def save_checkpoint(
-    ens: Ensemble,
-    out_dir: str | Path,
-    *,
-    base_llm: str,
-    training_meta: dict[str, Any] | None = None,
-) -> Path:
-    """Persist ensemble (LLM adapters + emb + JEPA + bridge + router + store)."""
-    root = Path(out_dir).resolve()
-    root.mkdir(parents=True, exist_ok=True)
-
-    backend = "tiny" if isinstance(ens.llm, TinyBackend) else "hf"
-    manifest: dict[str, Any] = {
-        "version": CHECKPOINT_VERSION,
-        "track": "jepa",
-        "backend": backend,
-        "base_llm": base_llm,
-        "adapter_names": list(ens.adapter_names),
-        "d_emb": ens.emb.d_emb,
-        "d_jepa": ens.jepa.d_latent,
-        "training": training_meta or {},
-    }
-
-    torch.save(_aux_state_dict(ens), root / AUX_FILE)
-    store = _store_payload(ens)
-    if store["keys"]:
-        torch.save(store, root / STORE_FILE)
-
-    if backend == "hf":
-        llm_path = root / LLM_DIR
-        llm_path.mkdir(exist_ok=True)
-        ens.llm.model.save_pretrained(llm_path)
-        ens.llm.tokenizer.save_pretrained(llm_path)
-    else:
-        torch.save(ens.llm.state_dict(), root / TINY_LLM_FILE)
-
-    with open(root / MANIFEST_FILE, "w") as f:
-        json.dump(manifest, f, indent=2)
-
-    return root
-
-
-def is_ensemble_checkpoint(path: str | Path) -> bool:
-    return (Path(path) / MANIFEST_FILE).is_file()
-
-
-def load_checkpoint(
-    ckpt_dir: str | Path,
-    *,
-    device: str | None = None,
-    load_in_4bit: bool = False,
-) -> Ensemble:
-    """Restore a saved JEPA ensemble from models/ensemble/<name>/."""
-    root = Path(ckpt_dir).resolve()
-    manifest_path = root / MANIFEST_FILE
-    if not manifest_path.is_file():
-        raise FileNotFoundError(
-            f"Not an ensemble checkpoint (missing {MANIFEST_FILE}): {root}"
-        )
-
-    with open(manifest_path) as f:
-        manifest = json.load(f)
-
-    base_llm = manifest["base_llm"]
-    backend = manifest.get("backend", "hf")
-    adapter_names = tuple(manifest.get("adapter_names", ["general"]))
-    d_emb = manifest.get("d_emb", 64)
-    d_jepa = manifest.get("d_jepa", 64)
-
-    if backend == "tiny":
-        ens = Ensemble(
-            llm="tiny",
-            adapter_names=adapter_names,
-            d_emb=d_emb,
-            d_jepa=d_jepa,
-        )
-        tiny_state = torch.load(
-            root / TINY_LLM_FILE, map_location="cpu", weights_only=True
-        )
-        ens.llm.load_state_dict(tiny_state)
-    else:
-        llm_dir = root / LLM_DIR
-        llm_backend = load_hf_backend_from_checkpoint(
-            base_llm,
-            str(llm_dir) if llm_dir.is_dir() else None,
-            adapter_names=adapter_names,
-            device=device,
-            load_in_4bit=load_in_4bit,
-        )
-        ens = Ensemble(
-            llm=base_llm,
-            adapter_names=adapter_names,
-            d_emb=d_emb,
-            d_jepa=d_jepa,
-            llm_backend=llm_backend,
-        )
-
-    aux = torch.load(root / AUX_FILE, map_location="cpu", weights_only=True)
-    ens.emb.load_state_dict(aux["emb"])
-    ens.jepa.load_state_dict(aux["jepa"])
-    ens.bridge.load_state_dict(aux["bridge"])
-    ens.router.load_state_dict(aux["router"])
-
-    store_path = root / STORE_FILE
-    if store_path.is_file():
-        store = torch.load(store_path, map_location="cpu", weights_only=True)
-        ens.store.keys = list(store["keys"])
-        ens.store.values = list(store["values"])
-
-    ens.eval()
-    return ens

research/ensemble/src/ensemble/config.py

@@ -1,163 +0,0 @@
-"""Resolve base LLM for ensemble from .env and models.yaml (same order as finetune)."""
-
-from __future__ import annotations
-
-import os
-import sys
-from pathlib import Path
-
-_REPO_ROOT = Path(__file__).resolve().parents[4]
-_FALLBACK_PRESET = "minicpm5-1b"
-
-_ENV_LLM_KEYS = (
-    "ENSEMBLE_LLM",
-    "LLM_PATH",
-    "BASE",
-    "FINETUNE_MODEL",
-    "MODEL_ID",
-)
-
-
-def repo_root() -> Path:
-    return _REPO_ROOT
-
-
-def load_dotenv() -> None:
-    """Load KEY=VALUE pairs from repo .env without overriding existing env vars."""
-    path = _REPO_ROOT / ".env"
-    if not path.is_file():
-        return
-    for line in path.read_text().splitlines():
-        line = line.strip()
-        if not line or line.startswith("#") or "=" not in line:
-            continue
-        key, _, value = line.partition("=")
-        key = key.strip()
-        value = value.strip().strip('"').strip("'")
-        if key:
-            os.environ.setdefault(key, value)
-
-
-def _ensure_inference_on_path() -> None:
-    libs = _REPO_ROOT / "libs" / "inference" / "src"
-    if str(libs) not in sys.path:
-        sys.path.insert(0, str(libs))
-
-
-def _is_ensemble_llm_preset(model) -> bool:
-    return model.backend == "transformers" and not model.multimodal and bool(
-        model.model_id
-    )
-
-
-def _llm_from_local_path(raw: str) -> str | None:
-    path = Path(raw)
-    if not path.is_absolute():
-        path = (_REPO_ROOT / path).resolve()
-    if path.suffix == ".gguf":
-        return None
-    if path.is_dir() and (path / "config.json").is_file():
-        return str(path)
-    if path.is_file():
-        return None
-    return None
-
-
-def _llm_from_env_paths() -> str | None:
-    for key in ("LLM_PATH", "MODEL_PATH"):
-        raw = os.environ.get(key)
-        if raw:
-            resolved = _llm_from_local_path(raw)
-            if resolved:
-                return resolved
-    return None
-
-
-def resolve_llm(
-    *,
-    llm_arg: str | None = None,
-    preset_arg: str | None = None,
-) -> tuple[str, str | None]:
-    """
-    Return (hub_id_or_local_path, preset_key) for ensemble HF backends.
-
-    Priority when llm_arg is None or ``auto``:
-      1. ENSEMBLE_LLM, LLM_PATH (local HF dir), BASE, FINETUNE_MODEL, MODEL_ID
-      2. MODEL_PATH if it points at a HuggingFace model directory (not .gguf)
-      3. ENSEMBLE_PRESET, FINETUNE_PRESET, or ACTIVE_MODEL from models.yaml
-      4. First fine-tunable transformers preset (default minicpm5-1b)
-    """
-    if llm_arg and llm_arg not in ("auto",):
-        return llm_arg, preset_arg
-
-    for env_name in _ENV_LLM_KEYS:
-        raw = os.environ.get(env_name)
-        if raw:
-            local = _llm_from_local_path(raw)
-            return local or raw, preset_arg
-
-    local = _llm_from_env_paths()
-    if local:
-        return local, preset_arg
-
-    _ensure_inference_on_path()
-    from inference.config import get_app_config, get_model_config
-
-    app_config = get_app_config(reload=True)
-    preset_key = (
-        preset_arg
-        or os.environ.get("ENSEMBLE_PRESET")
-        or os.environ.get("FINETUNE_PRESET")
-        or os.environ.get("ACTIVE_MODEL")
-    )
-
-    if preset_key and preset_key in app_config.models:
-        model = get_model_config(preset_key)
-        if not _is_ensemble_llm_preset(model):
-            preset_key = None
-
-    if preset_key is None:
-        for candidate in (_FALLBACK_PRESET, *app_config.models):
-            if candidate not in app_config.models:
-                continue
-            model = get_model_config(candidate)
-            if _is_ensemble_llm_preset(model):
-                preset_key = candidate
-                break
-
-    if not preset_key:
-        raise SystemExit(
-            "No transformers LLM found for ensemble. Pass --llm, set LLM_PATH/BASE/"
-            "MODEL_ID in .env, or ACTIVE_MODEL in models.yaml."
-        )
-
-    model = get_model_config(preset_key)
-    if not _is_ensemble_llm_preset(model):
-        raise SystemExit(
-            f"Preset {preset_key!r} cannot back an ensemble "
-            f"(backend={model.backend}, multimodal={model.multimodal})."
-        )
-    return model.model_id, preset_key
-
-
-def default_ensemble_out(preset_key: str | None) -> str:
-    label = preset_key or "custom"
-    return str((_REPO_ROOT / "models" / "ensemble" / f"{label}-jepa-pretrain").resolve())
-
-
-def resolve_llm_cli(
-    llm: str | None,
-    *,
-    toy: bool = False,
-    preset: str | None = None,
-) -> str:
-    """CLI helper: explicit tiny, else .env / models.yaml unless --toy without --llm."""
-    if llm == "tiny":
-        return "tiny"
-    if llm is None or llm == "auto":
-        if toy:
-            return "tiny"
-        load_dotenv()
-        resolved, _ = resolve_llm(preset_arg=preset)
-        return resolved
-    return llm

research/ensemble/src/ensemble/energy.py

@@ -1,45 +0,0 @@
-"""Energy model: score candidate latents against world state."""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class EnergyModel(nn.Module):
-    def __init__(self, d_latent: int):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(2 * d_latent, 2 * d_latent),
-            nn.GELU(),
-            nn.Linear(2 * d_latent, d_latent),
-            nn.GELU(),
-            nn.Linear(d_latent, 1),
-        )
-        self.d_latent = d_latent
-
-    def energy(self, s, z):
-        return self.net(torch.cat([s, z], -1)).squeeze(-1)
-
-    def contrastive_loss(self, s, z_pos, z_negs=None, tau=0.5):
-        B = s.size(0)
-        s_rep = s.unsqueeze(1).expand(B, B, self.d_latent).reshape(
-            B * B, self.d_latent
-        )
-        z_rep = z_pos.unsqueeze(0).expand(B, B, self.d_latent).reshape(
-            B * B, self.d_latent
-        )
-        E = self.energy(s_rep, z_rep).view(B, B)
-        if z_negs is not None:
-            En = self.energy(
-                s.repeat_interleave(z_negs.size(1), 0),
-                z_negs.reshape(-1, self.d_latent),
-            ).view(B, -1)
-            E = torch.cat([E, En], dim=1)
-        labels = torch.arange(B, device=s.device)
-        return F.cross_entropy(-E / tau, labels)
-
-    @torch.no_grad()
-    def rank(self, s, candidates):
-        return self.energy(s.expand(candidates.size(0), -1), candidates)

research/ensemble/src/ensemble/eval/__init__.py

@@ -1 +0,0 @@
-"""Evaluation harnesses for JEPA and world-model ensembles."""

research/ensemble/src/ensemble/eval/jepa_harness.py

@@ -1,266 +0,0 @@
-"""Ablation ladder + JEPA best-of-N benchmark for the ensemble."""
-
-from __future__ import annotations
-
-import argparse
-import json
-import random
-import time
-from collections import defaultdict
-
-import torch
-import torch.nn.functional as F
-
-from ensemble.eval.metrics import em_score, f1_score, paired_bootstrap
-from ensemble.backends import TinyBackend
-from ensemble.checkpoint import load_checkpoint
-from ensemble.config import load_dotenv, resolve_llm_cli
-from ensemble.jepa_ensemble import Ensemble
-
-
-@torch.no_grad()
-def generate_plain(ens, q_ids, n_new):
-    ens.llm.set_adapter(ens.adapter_names[0])
-    t0 = time.time()
-    out = ens.llm.generate(q_ids.to(ens.llm.device), n_new=n_new, temperature=0.7)
-    return out[:, q_ids.size(1) :], time.time() - t0
-
-
-@torch.no_grad()
-def generate_config(
-    ens, q_ids, n_new, *, use_rag, use_router, use_jepa, n_drafts=1, tau=0.0
-):
-    q_emb = ens.emb(q_ids.cpu())
-
-    if use_router:
-        a_idx = ens.router(q_emb).item()
-        ens.llm.set_adapter(ens.adapter_names[a_idx])
-    else:
-        ens.llm.set_adapter(ens.adapter_names[0])
-
-    ctx = q_ids.cpu()
-    if use_rag:
-        mems = ens.store.search(q_emb, k=1)
-        if mems:
-            ctx = torch.cat([mems[0], ctx], dim=1)
-
-    t0 = time.time()
-    if not use_jepa:
-        out = ens.llm.generate(
-            ctx.to(ens.llm.device), n_new=n_new, temperature=0.7
-        )
-        return out[:, ctx.size(1) :], time.time() - t0, None
-
-    z_exp = ens.jepa.predict_next_latent(ctx)
-    drafts, scores = [], []
-    for _ in range(n_drafts):
-        out = ens.llm.generate(
-            ctx.to(ens.llm.device), n_new=n_new, temperature=0.9
-        )
-        new = out[:, ctx.size(1) :].cpu()
-        drafts.append(new)
-        scores.append(
-            F.cosine_similarity(z_exp, ens.jepa.encode(new)).item()
-        )
-    best = max(range(n_drafts), key=lambda i: scores[i])
-    return drafts[best], time.time() - t0, (drafts, scores)
-
-
-def selector_comparison(drafts_scores_gold, decode_fn, rng):
-    res = defaultdict(list)
-    for drafts, scores, gold in drafts_scores_gold:
-        texts = [decode_fn(d) for d in drafts]
-        ems = [em_score(t, gold) for t in texts]
-        res["first"].append(ems[0])
-        res["random"].append(ems[rng.randrange(len(ems))])
-        res["jepa"].append(ems[max(range(len(ems)), key=lambda i: scores[i])])
-        res["oracle"].append(max(ems))
-    return {k: sum(v) / len(v) for k, v in res.items()}, res
-
-
-def load_jsonl(path):
-    with open(path) as f:
-        return [json.loads(line) for line in f if line.strip()]
-
-
-def make_toy_data(ens, n_qa=20, vocab=None):
-    vocab = vocab or ens.llm.vocab_size
-    qa, kb = [], []
-    for _ in range(n_qa):
-        key = torch.randint(0, vocab, (1, 6))
-        ans = torch.randint(0, vocab, (1, 4))
-        kb.append(torch.cat([key, ans], dim=1))
-        qa.append({"q_ids": key, "answer_ids": ans})
-    return qa, kb
-
-
-def run(args):
-    torch.manual_seed(args.seed)
-    rng = random.Random(args.seed)
-
-    if args.ckpt:
-        ens = load_checkpoint(args.ckpt)
-        print(f"loaded ensemble checkpoint: {args.ckpt}")
-        is_text = not isinstance(ens.llm, TinyBackend)
-    else:
-        load_dotenv()
-        args.llm = resolve_llm_cli(
-            args.llm, toy=args.toy, preset=getattr(args, "preset", None)
-        )
-        print(f"Resolved LLM: {args.llm}")
-        ens = Ensemble(llm=args.llm)
-        is_text = args.llm != "tiny"
-
-    if args.toy or not is_text:
-        qa, kb = make_toy_data(ens)
-        for mem in kb:
-            ens.memorize_ids(mem)
-
-        def to_ids(item):
-            return item["q_ids"]
-
-        def gold_text(item):
-            return " ".join(map(str, item["answer_ids"][0].tolist()))
-
-        def decode(ids):
-            return " ".join(map(str, ids[0].tolist()))
-    else:
-        qa = load_jsonl(args.qa)
-        if args.kb:
-            for row in load_jsonl(args.kb):
-                ens.memorize_text(row["text"])
-
-        def to_ids(item):
-            return ens.llm.encode_text(
-                f"Answer briefly.\nQ: {item['question']}\nA:"
-            )
-
-        def gold_text(item):
-            return item["answer"]
-
-        def decode(ids):
-            return ens.llm.decode(ids)
-
-    qa = qa[: args.limit]
-    print(
-        f"eval set: {len(qa)} questions | store: {len(ens.store.keys)} memories\n"
-    )
-
-    configs = {
-        "C1_base": dict(use_rag=False, use_router=False, use_jepa=False),
-        "C2_rag": dict(use_rag=True, use_router=False, use_jepa=False),
-        "C3_rag_router": dict(use_rag=True, use_router=True, use_jepa=False),
-        "C4_full_jepa": dict(
-            use_rag=True,
-            use_router=True,
-            use_jepa=True,
-            n_drafts=args.n_drafts,
-        ),
-    }
-
-    per_q = {}
-    summary = {}
-    jepa_material = []
-
-    for name, cfg in configs.items():
-        ems, f1s, lats = [], [], []
-        for item in qa:
-            ids = to_ids(item)
-            if name == "C1_base":
-                out, dt = generate_plain(ens, ids, args.n_new)
-                extra = None
-            else:
-                out, dt, extra = generate_config(ens, ids, args.n_new, **cfg)
-            pred, gold = decode(out), gold_text(item)
-            ems.append(em_score(pred, gold))
-            f1s.append(f1_score(pred, gold))
-            lats.append(dt)
-            if name == "C4_full_jepa" and extra is not None:
-                jepa_material.append((extra[0], extra[1], gold))
-        per_q[name] = ems
-        summary[name] = (
-            sum(ems) / len(ems),
-            sum(f1s) / len(f1s),
-            sum(lats) / len(lats),
-        )
-
-    print(f"{'config':<16}{'EM':>8}{'F1':>8}{'lat(s)':>9}")
-    for k, (em, f1, lat) in summary.items():
-        print(f"{k:<16}{em:>8.3f}{f1:>8.3f}{lat:>9.3f}")
-
-    print("\ncomponent contributions (paired bootstrap, P(B>A)):")
-    ladder = list(configs.keys())
-    for a, b in zip(ladder, ladder[1:]):
-        d = summary[b][0] - summary[a][0]
-        p = paired_bootstrap(per_q[a], per_q[b])
-        print(f"  {b} - {a}: ΔEM={d:+.3f}   P(better)={p:.2f}")
-
-    if jepa_material:
-        sel, sel_per_q = selector_comparison(jepa_material, decode, rng)
-        print(
-            f"\nbest-of-N selector comparison (same drafts, N={args.n_drafts}):"
-        )
-        for k in ("first", "random", "jepa", "oracle"):
-            print(f"  {k:<8}EM={sel[k]:.3f}")
-        p = paired_bootstrap(sel_per_q["random"], sel_per_q["jepa"])
-        verdict = (
-            "JEPA critic WORKS"
-            if p > 0.95
-            else "inconclusive — critic ~ random"
-        )
-        print(f"  P(jepa > random) = {p:.2f}   {verdict}")
-        print(f"  headroom to oracle: {sel['oracle'] - sel['jepa']:.3f}")
-
-    if args.continual:
-        print(
-            "\ncontinual test: accuracy on task-A questions "
-            "before vs after adding adapters B and C"
-        )
-        ems_before = per_q["C3_rag_router"]
-        ens.new_task_adapter("task_B")
-        ens.new_task_adapter("task_C")
-        ems_after = []
-        for item in qa:
-            out, _, _ = generate_config(
-                ens,
-                to_ids(item),
-                args.n_new,
-                use_rag=True,
-                use_router=True,
-                use_jepa=False,
-            )
-            ems_after.append(em_score(decode(out), gold_text(item)))
-        bt = sum(ems_after) / len(ems_after) - sum(ems_before) / len(
-            ems_before
-        )
-        print(f"  backward transfer (≈0 is ideal): {bt:+.3f}")
-
-    return summary
-
-
-def parse_args():
-    p = argparse.ArgumentParser()
-    p.add_argument(
-        "--llm",
-        default=None,
-        help="HF id / path, 'tiny', or omit for LLM_PATH / ACTIVE_MODEL from .env",
-    )
-    p.add_argument("--preset", default=None, help="models.yaml preset override")
-    p.add_argument("--qa", default=None, help="jsonl with question/answer")
-    p.add_argument("--kb", default=None, help="jsonl with text -> vector store")
-    p.add_argument(
-        "--ckpt",
-        default=None,
-        help="saved ensemble directory (models/ensemble/... with manifest.json)",
-    )
-    p.add_argument("--toy", action="store_true", help="synthetic data smoke test")
-    p.add_argument("--limit", type=int, default=100)
-    p.add_argument("--n_new", type=int, default=24)
-    p.add_argument("--n_drafts", type=int, default=8)
-    p.add_argument("--continual", action="store_true")
-    p.add_argument("--seed", type=int, default=0)
-    return p.parse_args()
-
-
-if __name__ == "__main__":
-    run(parse_args())

research/ensemble/src/ensemble/eval/metrics.py

@@ -1,42 +0,0 @@
-"""QA metrics and paired bootstrap significance."""
-
-from __future__ import annotations
-
-import random
-import re
-import string
-from collections import Counter
-
-
-def normalize(s: str) -> str:
-    s = s.lower()
-    s = "".join(c for c in s if c not in string.punctuation)
-    s = re.sub(r"\b(a|an|the)\b", " ", s)
-    return " ".join(s.split())
-
-
-def em_score(pred: str, gold: str) -> float:
-    return float(normalize(gold) in normalize(pred))
-
-
-def f1_score(pred: str, gold: str) -> float:
-    p, g = normalize(pred).split(), normalize(gold).split()
-    if not p or not g:
-        return float(p == g)
-    common = Counter(p) & Counter(g)
-    overlap = sum(common.values())
-    if overlap == 0:
-        return 0.0
-    prec, rec = overlap / len(p), overlap / len(g)
-    return 2 * prec * rec / (prec + rec)
-
-
-def paired_bootstrap(scores_a, scores_b, iters=2000, seed=0):
-    rng = random.Random(seed)
-    n, wins = len(scores_a), 0
-    for _ in range(iters):
-        idx = [rng.randrange(n) for _ in range(n)]
-        da = sum(scores_a[i] for i in idx) / n
-        db = sum(scores_b[i] for i in idx) / n
-        wins += db > da
-    return wins / iters

research/ensemble/src/ensemble/eval/world_harness.py

@@ -1,174 +0,0 @@
-"""Energy-based draft selector benchmark for the world-model ensemble."""
-
-from __future__ import annotations
-
-import argparse
-import json
-import random
-import time
-from collections import defaultdict
-
-import torch
-
-from ensemble.eval.metrics import em_score, f1_score, paired_bootstrap
-from ensemble.world_ensemble import WorldEnsemble
-
-
-@torch.no_grad()
-def generate_drafts(ens, q_ids, n_new, n_drafts, use_rag=True):
-    q_emb = ens.emb(q_ids.cpu())
-    mems = ens.store.search(q_emb, k=1) if use_rag else []
-    segments = (mems + [q_ids.cpu()]) if mems else [q_ids.cpu()]
-    ctx = torch.cat(segments, dim=1)
-
-    s = ens.world_state(segments)
-    ens.world.rollout(s, horizon=3)
-
-    drafts, energies = [], []
-    t0 = time.time()
-    for _ in range(n_drafts):
-        out = ens.llm.generate(
-            ctx.to(ens.llm.device), n_new=n_new, temperature=0.9
-        )
-        new = out[:, ctx.size(1) :].cpu()
-        drafts.append(new)
-        z = ens.jepa.encode(new)
-        energies.append(ens.energy.rank(s, z).item())
-    return drafts, energies, time.time() - t0
-
-
-def selector_comparison(drafts_energy_gold, decode_fn, rng):
-    res = defaultdict(list)
-    for drafts, energies, gold in drafts_energy_gold:
-        texts = [decode_fn(d) for d in drafts]
-        ems = [em_score(t, gold) for t in texts]
-        res["first"].append(ems[0])
-        res["random"].append(ems[rng.randrange(len(ems))])
-        res["energy"].append(
-            ems[min(range(len(ems)), key=lambda i: energies[i])]
-        )
-        res["oracle"].append(max(ems))
-    return {k: sum(v) / len(v) for k, v in res.items()}, res
-
-
-def load_jsonl(path):
-    with open(path) as f:
-        return [json.loads(line) for line in f if line.strip()]
-
-
-def make_toy_data(ens, n_qa=20, vocab=None):
-    vocab = vocab or ens.llm.vocab_size
-    qa, kb = [], []
-    for _ in range(n_qa):
-        key = torch.randint(0, vocab, (1, 6))
-        ans = torch.randint(0, vocab, (1, 4))
-        kb.append(torch.cat([key, ans], dim=1))
-        qa.append({"q_ids": key, "answer_ids": ans})
-    return qa, kb
-
-
-def run(args):
-    from ensemble.config import load_dotenv, resolve_llm_cli
-
-    torch.manual_seed(args.seed)
-    rng = random.Random(args.seed)
-
-    load_dotenv()
-    args.llm = resolve_llm_cli(
-        args.llm, toy=args.toy, preset=getattr(args, "preset", None)
-    )
-    print(f"Resolved LLM: {args.llm}")
-    ens = WorldEnsemble(args.llm)
-    if args.ckpt:
-        state = torch.load(args.ckpt, map_location="cpu")
-        ens.load_state_dict(state, strict=False)
-        print(f"loaded world ensemble checkpoint: {args.ckpt}")
-
-    is_text = args.llm != "tiny"
-
-    if args.toy or not is_text:
-        qa, kb = make_toy_data(ens)
-        for mem in kb:
-            ens.memorize(mem)
-
-        def to_ids(item):
-            return item["q_ids"]
-
-        def gold_text(item):
-            return " ".join(map(str, item["answer_ids"][0].tolist()))
-
-        def decode(ids):
-            return " ".join(map(str, ids[0].tolist()))
-    else:
-        qa = load_jsonl(args.qa)
-        if args.kb:
-            for row in load_jsonl(args.kb):
-                ids = ens.llm.tokenizer(
-                    row["text"], return_tensors="pt"
-                ).input_ids
-                ens.memorize(ids)
-
-        def to_ids(item):
-            return ens.llm.tokenizer(
-                f"Answer briefly.\nQ: {item['question']}\nA:",
-                return_tensors="pt",
-            ).input_ids
-
-        def gold_text(item):
-            return item["answer"]
-
-        def decode(ids):
-            return ens.llm.tokenizer.decode(ids[0], skip_special_tokens=True)
-
-    qa = qa[: args.limit]
-    print(
-        f"eval set: {len(qa)} questions | store: {len(ens.store.keys)} memories\n"
-    )
-
-    material = []
-    lats = []
-    for item in qa:
-        drafts, energies, dt = generate_drafts(
-            ens, to_ids(item), args.n_new, args.n_drafts
-        )
-        material.append((drafts, energies, gold_text(item)))
-        lats.append(dt)
-
-    sel, sel_per_q = selector_comparison(material, decode, rng)
-    print(f"best-of-N selector comparison (same drafts, N={args.n_drafts}):")
-    for k in ("first", "random", "energy", "oracle"):
-        print(f"  {k:<8}EM={sel[k]:.3f}")
-    p = paired_bootstrap(sel_per_q["random"], sel_per_q["energy"])
-    verdict = (
-        "Energy critic WORKS"
-        if p > 0.95
-        else "inconclusive — critic ~ random"
-    )
-    print(f"  P(energy > random) = {p:.2f}   {verdict}")
-    print(f"  headroom to oracle: {sel['oracle'] - sel['energy']:.3f}")
-    print(f"  mean latency: {sum(lats) / len(lats):.3f}s")
-
-    return sel
-
-
-def parse_args():
-    p = argparse.ArgumentParser()
-    p.add_argument(
-        "--llm",
-        default=None,
-        help="HF id / path, 'tiny', or omit for LLM_PATH / ACTIVE_MODEL from .env",
-    )
-    p.add_argument("--preset", default=None, help="models.yaml preset override")
-    p.add_argument("--qa", default=None, help="jsonl with question/answer")
-    p.add_argument("--kb", default=None, help="jsonl with text -> vector store")
-    p.add_argument("--ckpt", default=None, help="trained world ensemble .pt")
-    p.add_argument("--toy", action="store_true", help="synthetic data smoke test")
-    p.add_argument("--limit", type=int, default=100)
-    p.add_argument("--n_new", type=int, default=24)
-    p.add_argument("--n_drafts", type=int, default=8)
-    p.add_argument("--seed", type=int, default=0)
-    return p.parse_args()
-
-
-if __name__ == "__main__":
-    run(parse_args())

research/ensemble/src/ensemble/eval_harness.py

@@ -1,309 +0,0 @@
-"""
-eval_harness.py — Ablation ladder + JEPA best-of-N test for the ensemble
-========================================================================
-Companion to `llm_emb_jepa_ensemble_pluggable.py` (must be importable,
-i.e. in the same directory).
-
-What it runs
-------------
-1. ABLATION LADDER on a QA set:
-     C1  base LLM alone
-     C2  C1 + RAG (embedding retrieval)
-     C3  C2 + router/adapters
-     C4  C3 + JEPA best-of-N critic
-   (C5 = C4 with a bridge-trained checkpoint — just pass --ckpt)
-
-2. BEST-OF-N SELECTOR comparison (the decisive JEPA experiment):
-     first-sample | random-pick | JEPA-score pick | oracle pick
-   All on the SAME N drafts per question, so differences are pure selection.
-
-3. CONTINUAL FORGETTING test (optional, --continual):
-     accuracy on task A before vs after training adapters for B and C.
-
-4. PAIRED BOOTSTRAP significance between any two configs.
-
-Usage
------
-# Smoke test, no GPU/deps beyond torch (toy backend, synthetic QA):
-python eval_harness.py --llm tiny --toy
-
-# Real model + your QA file (jsonl: {"question": ..., "answer": ..., "context": optional}):
-python eval_harness.py --llm Qwen/Qwen2.5-0.5B-Instruct \
-    --qa ./domain_qa.jsonl --kb ./knowledge.jsonl --n_drafts 8
-
-# With a bridge-trained ensemble checkpoint (C5):
-python eval_harness.py --llm /models/llama-3.2-1b --qa ./qa.jsonl \
-    --kb ./kb.jsonl --ckpt ./ensemble_bridge.pt
-
-QA file:  {"question": str, "answer": str, "domain": optional str}
-KB file:  {"text": str}   (each line becomes one memory in the vector store)
-"""
-
-import argparse
-import json
-import random
-import re
-import string
-import time
-from collections import Counter, defaultdict
-
-import torch
-
-from llm_emb_jepa_ensemble_pluggable import Ensemble  # same directory
-
-# ----------------------------------------------------------------------------
-# Metrics: normalized exact match + token F1 (SQuAD-style)
-# ----------------------------------------------------------------------------
-def normalize(s: str) -> str:
-    s = s.lower()
-    s = "".join(c for c in s if c not in string.punctuation)
-    s = re.sub(r"\b(a|an|the)\b", " ", s)
-    return " ".join(s.split())
-
-
-def em_score(pred: str, gold: str) -> float:
-    return float(normalize(gold) in normalize(pred))   # containment EM
-
-
-def f1_score(pred: str, gold: str) -> float:
-    p, g = normalize(pred).split(), normalize(gold).split()
-    if not p or not g:
-        return float(p == g)
-    common = Counter(p) & Counter(g)
-    overlap = sum(common.values())
-    if overlap == 0:
-        return 0.0
-    prec, rec = overlap / len(p), overlap / len(g)
-    return 2 * prec * rec / (prec + rec)
-
-
-# ----------------------------------------------------------------------------
-# Paired bootstrap: P(config B beats config A)
-# ----------------------------------------------------------------------------
-def paired_bootstrap(scores_a, scores_b, iters=2000, seed=0):
-    rng = random.Random(seed)
-    n, wins = len(scores_a), 0
-    for _ in range(iters):
-        idx = [rng.randrange(n) for _ in range(n)]
-        da = sum(scores_a[i] for i in idx) / n
-        db = sum(scores_b[i] for i in idx) / n
-        wins += db > da
-    return wins / iters
-
-
-# ----------------------------------------------------------------------------
-# Config runners — each returns per-question dicts
-# ----------------------------------------------------------------------------
-@torch.no_grad()
-def generate_plain(ens, q_ids, n_new):
-    """C1: base adapter, no retrieval, single sample."""
-    ens.llm.set_adapter(ens.adapter_names[0])
-    t0 = time.time()
-    out = ens.llm.generate(q_ids.to(ens.llm.device), n_new=n_new, temperature=0.7)
-    return out[:, q_ids.size(1):], time.time() - t0
-
-
-@torch.no_grad()
-def generate_config(ens, q_ids, n_new, *, use_rag, use_router, use_jepa,
-                    n_drafts=1, tau=0.0):
-    """Unified runner for C2/C3/C4."""
-    q_emb = ens.emb(q_ids.cpu())
-
-    if use_router:
-        a_idx = ens.router(q_emb).item()
-        ens.llm.set_adapter(ens.adapter_names[a_idx])
-    else:
-        ens.llm.set_adapter(ens.adapter_names[0])
-
-    ctx = q_ids.cpu()
-    if use_rag:
-        mems = ens.store.search(q_emb, k=1)
-        if mems:
-            ctx = torch.cat([mems[0], ctx], dim=1)
-
-    t0 = time.time()
-    if not use_jepa:
-        out = ens.llm.generate(ctx.to(ens.llm.device), n_new=n_new, temperature=0.7)
-        return out[:, ctx.size(1):], time.time() - t0, None
-
-    # JEPA best-of-N: sample drafts, keep the one closest to predicted latent
-    z_exp = ens.jepa.predict_next_latent(ctx)
-    drafts, scores = [], []
-    for _ in range(n_drafts):
-        out = ens.llm.generate(ctx.to(ens.llm.device), n_new=n_new, temperature=0.9)
-        new = out[:, ctx.size(1):].cpu()
-        drafts.append(new)
-        scores.append(torch.nn.functional.cosine_similarity(
-            z_exp, ens.jepa.encode(new)).item())
-    best = max(range(n_drafts), key=lambda i: scores[i])
-    return drafts[best], time.time() - t0, (drafts, scores)
-
-
-# ----------------------------------------------------------------------------
-# Best-of-N selector comparison on shared drafts
-# ----------------------------------------------------------------------------
-def selector_comparison(drafts_scores_gold, decode_fn, rng):
-    """drafts_scores_gold: list of (drafts, jepa_scores, gold_answer).
-    Returns EM for: first | random | jepa | oracle — all on the SAME drafts."""
-    res = defaultdict(list)
-    for drafts, scores, gold in drafts_scores_gold:
-        texts = [decode_fn(d) for d in drafts]
-        ems = [em_score(t, gold) for t in texts]
-        res["first"].append(ems[0])
-        res["random"].append(ems[rng.randrange(len(ems))])
-        res["jepa"].append(ems[max(range(len(ems)), key=lambda i: scores[i])])
-        res["oracle"].append(max(ems))     # upper bound of selection
-    return {k: sum(v) / len(v) for k, v in res.items()}, res
-
-
-# ----------------------------------------------------------------------------
-# Data loading
-# ----------------------------------------------------------------------------
-def load_jsonl(path):
-    with open(path) as f:
-        return [json.loads(l) for l in f if l.strip()]
-
-
-def make_toy_data(ens, n_qa=20, vocab=None):
-    """Synthetic QA for the tiny backend: 'answer' token sequence is planted
-    in the KB so RAG can genuinely help even with random weights."""
-    vocab = vocab or ens.llm.vocab_size
-    qa, kb = [], []
-    for i in range(n_qa):
-        key = torch.randint(0, vocab, (1, 6))
-        ans = torch.randint(0, vocab, (1, 4))
-        kb.append(torch.cat([key, ans], dim=1))            # memory = key+answer
-        qa.append({"q_ids": key, "answer_ids": ans})
-    return qa, kb
-
-
-# ----------------------------------------------------------------------------
-# Main evaluation
-# ----------------------------------------------------------------------------
-def run(args):
-    torch.manual_seed(args.seed)
-    rng = random.Random(args.seed)
-
-    ens = Ensemble(llm=args.llm)
-    if args.ckpt:
-        state = torch.load(args.ckpt, map_location="cpu")
-        ens.load_state_dict(state, strict=False)
-        print(f"loaded ensemble checkpoint: {args.ckpt}")
-
-    is_text = args.llm != "tiny"
-
-    # ---- load data and fill the vector store -------------------------------
-    if args.toy or not is_text:
-        qa, kb = make_toy_data(ens)
-        for mem in kb:
-            ens.memorize_ids(mem)
-        def to_ids(item):  return item["q_ids"]
-        def gold_of(item): return item["answer_ids"]
-        def decode(ids):   return " ".join(map(str, ids[0].tolist()))
-        def gold_text(item): return decode(item["answer_ids"])
-    else:
-        qa = load_jsonl(args.qa)
-        if args.kb:
-            for row in load_jsonl(args.kb):
-                ens.memorize_text(row["text"])
-        def to_ids(item):  return ens.llm.encode_text(
-            f"Answer briefly.\nQ: {item['question']}\nA:")
-        def gold_text(item): return item["answer"]
-        def decode(ids):   return ens.llm.decode(ids)
-
-    qa = qa[: args.limit]
-    print(f"eval set: {len(qa)} questions | store: {len(ens.store.keys)} memories\n")
-
-    # ---- ablation ladder ----------------------------------------------------
-    configs = {
-        "C1_base":        dict(use_rag=False, use_router=False, use_jepa=False),
-        "C2_rag":         dict(use_rag=True,  use_router=False, use_jepa=False),
-        "C3_rag_router":  dict(use_rag=True,  use_router=True,  use_jepa=False),
-        "C4_full_jepa":   dict(use_rag=True,  use_router=True,  use_jepa=True,
-                               n_drafts=args.n_drafts),
-    }
-
-    per_q = {}            # config -> list of EM scores (for bootstrap)
-    summary = {}
-    jepa_material = []    # (drafts, scores, gold) for selector comparison
-
-    for name, cfg in configs.items():
-        ems, f1s, lats = [], [], []
-        for item in qa:
-            ids = to_ids(item)
-            if name == "C1_base":
-                out, dt = generate_plain(ens, ids, args.n_new)
-                extra = None
-            else:
-                out, dt, extra = generate_config(ens, ids, args.n_new, **cfg)
-            pred, gold = decode(out), gold_text(item)
-            ems.append(em_score(pred, gold))
-            f1s.append(f1_score(pred, gold))
-            lats.append(dt)
-            if name == "C4_full_jepa" and extra is not None:
-                jepa_material.append((extra[0], extra[1], gold))
-        per_q[name] = ems
-        summary[name] = (sum(ems) / len(ems), sum(f1s) / len(f1s),
-                         sum(lats) / len(lats))
-
-    print(f"{'config':<16}{'EM':>8}{'F1':>8}{'lat(s)':>9}")
-    for k, (em, f1, lat) in summary.items():
-        print(f"{k:<16}{em:>8.3f}{f1:>8.3f}{lat:>9.3f}")
-
-    # deltas + significance
-    print("\ncomponent contributions (paired bootstrap, P(B>A)):")
-    ladder = list(configs.keys())
-    for a, b in zip(ladder, ladder[1:]):
-        d = summary[b][0] - summary[a][0]
-        p = paired_bootstrap(per_q[a], per_q[b])
-        print(f"  {b} - {a}: ΔEM={d:+.3f}   P(better)={p:.2f}")
-
-    # ---- decisive JEPA selector experiment ----------------------------------
-    if jepa_material:
-        sel, sel_per_q = selector_comparison(jepa_material, decode, rng)
-        print("\nbest-of-N selector comparison (same drafts, N="
-              f"{args.n_drafts}):")
-        for k in ("first", "random", "jepa", "oracle"):
-            print(f"  {k:<8}EM={sel[k]:.3f}")
-        p = paired_bootstrap(sel_per_q["random"], sel_per_q["jepa"])
-        print(f"  P(jepa > random) = {p:.2f}   "
-              f"{'JEPA critic WORKS' if p > 0.95 else 'inconclusive — critic ~ random'}")
-        gap = sel["oracle"] - sel["jepa"]
-        print(f"  headroom to oracle: {gap:.3f}")
-
-    # ---- continual forgetting (optional) ------------------------------------
-    if args.continual:
-        print("\ncontinual test: accuracy on task-A questions "
-              "before vs after adding adapters B and C")
-        ems_before = per_q["C3_rag_router"]
-        ens.new_task_adapter("task_B")
-        ens.new_task_adapter("task_C")
-        ems_after = []
-        for item in qa:
-            out, _, _ = generate_config(ens, to_ids(item), args.n_new,
-                                        use_rag=True, use_router=True,
-                                        use_jepa=False)
-            ems_after.append(em_score(decode(out), gold_text(item)))
-        bt = sum(ems_after) / len(ems_after) - sum(ems_before) / len(ems_before)
-        print(f"  backward transfer (≈0 is ideal): {bt:+.3f}")
-
-    return summary
-
-
-def parse_args():
-    p = argparse.ArgumentParser()
-    p.add_argument("--llm", default="tiny", help="'tiny' | HF id | local path")
-    p.add_argument("--qa", default=None, help="jsonl with question/answer")
-    p.add_argument("--kb", default=None, help="jsonl with text -> vector store")
-    p.add_argument("--ckpt", default=None, help="bridge-trained ensemble .pt (C5)")
-    p.add_argument("--toy", action="store_true", help="synthetic data smoke test")
-    p.add_argument("--limit", type=int, default=100)
-    p.add_argument("--n_new", type=int, default=24)
-    p.add_argument("--n_drafts", type=int, default=8)
-    p.add_argument("--continual", action="store_true")
-    p.add_argument("--seed", type=int, default=0)
-    return p.parse_args()
-
-
-if __name__ == "__main__":
-    run(parse_args())

research/ensemble/src/ensemble/jepa.py

@@ -1,75 +0,0 @@
-"""JEPA latent predictor with EMA target encoder."""
-
-from __future__ import annotations
-
-import copy
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class _SegEncoder(nn.Module):
-    def __init__(self, vocab_size, d):
-        super().__init__()
-        self.tok = nn.Embedding(vocab_size, d)
-        self.enc = nn.GRU(d, d, batch_first=True)
-        self.out = nn.Linear(d, d)
-
-    def forward(self, ids):
-        h, _ = self.enc(self.tok(ids))
-        return self.out(h.mean(dim=1))
-
-
-class JEPA(nn.Module):
-    def __init__(self, vocab_size: int, d_latent: int = 64, ema_m: float = 0.996):
-        super().__init__()
-        self.ctx_enc = _SegEncoder(vocab_size, d_latent)
-        self.tgt_enc = copy.deepcopy(self.ctx_enc)
-        for p in self.tgt_enc.parameters():
-            p.requires_grad_(False)
-        self.predictor = nn.Sequential(
-            nn.Linear(d_latent, 2 * d_latent),
-            nn.GELU(),
-            nn.Linear(2 * d_latent, d_latent),
-        )
-        self.m = ema_m
-        self.d_latent = d_latent
-
-    @property
-    def enc(self):
-        """Alias used by world-model track."""
-        return self.ctx_enc
-
-    @property
-    def tgt(self):
-        return self.tgt_enc
-
-    @property
-    def pred(self):
-        return self.predictor
-
-    @torch.no_grad()
-    def ema_update(self):
-        for p_t, p_c in zip(self.tgt_enc.parameters(), self.ctx_enc.parameters()):
-            p_t.mul_(self.m).add_(p_c.detach(), alpha=1 - self.m)
-
-    def ema(self):
-        """Alias used by world-model track."""
-        self.ema_update()
-
-    def loss(self, seg_ctx, seg_tgt):
-        z_hat = self.predictor(self.ctx_enc(seg_ctx))
-        with torch.no_grad():
-            z_tgt = self.tgt_enc(seg_tgt)
-        pred = F.mse_loss(z_hat, z_tgt)
-        var_reg = F.relu(1.0 - z_hat.std(dim=0)).mean()
-        return pred + 0.5 * var_reg
-
-    @torch.no_grad()
-    def predict_next_latent(self, seg_ctx):
-        return self.predictor(self.ctx_enc(seg_ctx))
-
-    @torch.no_grad()
-    def encode(self, seg):
-        return self.tgt_enc(seg)

research/ensemble/src/ensemble/jepa_ensemble.py

@@ -1,232 +0,0 @@
-"""JEPA ensemble: route -> retrieve -> generate -> JEPA-verify."""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from ensemble.backends import HFBackend, make_backend
-from ensemble.bridge import Bridge
-from ensemble.jepa import JEPA
-from ensemble.memory import Embedder, Router, VectorStore
-
-torch.manual_seed(0)
-
-
-class Ensemble(nn.Module):
-    def __init__(
-        self,
-        llm: str = "tiny",
-        adapter_names=("general",),
-        d_emb: int = 64,
-        d_jepa: int = 64,
-        llm_backend: HFBackend | None = None,
-        **backend_kw,
-    ):
-        super().__init__()
-        self.llm = llm_backend if llm_backend is not None else make_backend(llm, **backend_kw)
-        V, H = self.llm.vocab_size, self.llm.hidden_size
-
-        self.emb = Embedder(V, d_emb)
-        self.jepa = JEPA(V, d_jepa)
-        self.bridge = Bridge(H, d_jepa)
-        self.store = VectorStore()
-
-        self.adapter_names = list(adapter_names)
-        for n in self.adapter_names:
-            self.llm.add_adapter(n)
-        self.llm.set_adapter(self.adapter_names[0])
-        self.router = Router(d_emb, len(self.adapter_names))
-
-    @torch.no_grad()
-    def answer_ids(
-        self,
-        query_ids,
-        n_new=32,
-        tau_consistency=0.0,
-        max_retries=2,
-        temperature: float = 0.7,
-    ):
-        q_emb = self.emb(query_ids.cpu())
-        a_idx = self.router(q_emb).item()
-        self.llm.set_adapter(self.adapter_names[a_idx])
-
-        mems = self.store.search(q_emb, k=1)
-        ctx = (
-            torch.cat([mems[0], query_ids.cpu()], dim=1)
-            if mems
-            else query_ids.cpu()
-        )
-
-        z_expected = self.jepa.predict_next_latent(ctx)
-
-        best = None
-        for attempt in range(max_retries + 1):
-            temp = temperature if attempt == 0 else max(temperature, 0.8 + 0.3 * attempt)
-            draft = self.llm.generate(
-                ctx.to(self.llm.device),
-                n_new=n_new,
-                temperature=temp,
-            )
-            new_part = draft[:, ctx.size(1) :].cpu()
-            score = F.cosine_similarity(
-                z_expected, self.jepa.encode(new_part)
-            ).item()
-            if best is None or score > best[1]:
-                best = (draft, score, attempt)
-            if score >= tau_consistency:
-                break
-        draft, score, attempt = best
-        return draft, score, self.adapter_names[a_idx], attempt
-
-    def answer_text(self, prompt: str, **kw):
-        ids = self.llm.encode_text(prompt)
-        out, score, adapter, retries = self.answer_ids(ids, **kw)
-        return self.llm.decode(out), score, adapter, retries
-
-    def generate_text(
-        self,
-        prompt: str,
-        *,
-        max_new_tokens: int = 512,
-        temperature: float = 0.0,
-    ) -> str:
-        """Greedy or sampled generation through the full ensemble stack."""
-        ids = self.llm.encode_text(prompt)
-        out, _, _, _ = self.answer_ids(
-            ids,
-            n_new=max_new_tokens,
-            tau_consistency=-1.0,
-            max_retries=0 if temperature <= 0 else 1,
-            temperature=temperature,
-        )
-        return self.llm.decode(out)
-
-    def memorize_ids(self, ids):
-        self.store.add(self.emb(ids.cpu()), ids.cpu())
-
-    def memorize_text(self, text: str):
-        self.memorize_ids(self.llm.encode_text(text))
-
-    def new_task_adapter(self, name: str):
-        self.adapter_names.append(name)
-        self.llm.add_adapter(name)
-        old = self.router
-        self.router = Router(self.emb.d_emb, len(self.adapter_names))
-        with torch.no_grad():
-            self.router.fc.weight[: old.fc.out_features] = old.fc.weight
-            self.router.fc.bias[: old.fc.out_features] = old.fc.bias
-
-    def train_step(self, seg_a, seg_b, opt, w_bridge=0.1):
-        logits, hidden = self.llm(seg_a.to(self.llm.device))
-        lm_loss = F.cross_entropy(
-            logits[:, :-1].reshape(-1, self.llm.vocab_size).float(),
-            seg_a[:, 1:].reshape(-1).to(logits.device),
-        )
-
-        jepa_loss = self.jepa.loss(seg_a.cpu(), seg_b.cpu())
-
-        z_llm = self.bridge(
-            hidden.cpu() if hidden.device.type != "cpu" else hidden
-        )
-        z_jepa = self.jepa.ctx_enc(seg_a.cpu()).detach()
-        bridge_loss = self.bridge.info_nce(z_llm, z_jepa.to(z_llm.device))
-
-        loss = lm_loss.cpu() + jepa_loss + w_bridge * bridge_loss
-        opt.zero_grad()
-        loss.backward()
-        opt.step()
-        self.jepa.ema_update()
-        return {
-            "lm": lm_loss.item(),
-            "jepa": jepa_loss.item(),
-            "bridge": bridge_loss.item(),
-        }
-
-    def make_optimizer(self, lr_lora=2e-4, lr_aux=1e-3):
-        return torch.optim.AdamW(
-            [
-                {"params": list(self.llm.trainable_parameters()), "lr": lr_lora},
-                {
-                    "params": list(self.jepa.ctx_enc.parameters())
-                    + list(self.jepa.predictor.parameters()),
-                    "lr": lr_aux,
-                },
-                {
-                    "params": list(self.bridge.parameters())
-                    + list(self.emb.parameters())
-                    + list(self.router.parameters()),
-                    "lr": lr_aux,
-                },
-            ]
-        )
-
-
-def segment_pairs_from_texts(backend: HFBackend, texts, seg_len=64):
-    a_list, b_list = [], []
-    for t in texts:
-        ids = backend.tokenizer(t, return_tensors="pt").input_ids[0]
-        for i in range(0, len(ids) - 2 * seg_len, seg_len):
-            a_list.append(ids[i : i + seg_len])
-            b_list.append(ids[i + seg_len : i + 2 * seg_len])
-    if not a_list:
-        raise ValueError("texts too short for the chosen seg_len")
-    return torch.stack(a_list), torch.stack(b_list)
-
-
-def demo_tiny(steps=50):
-    ens = Ensemble(llm="tiny")
-    opt = ens.make_optimizer()
-    for s in range(steps):
-        seg_a = torch.randint(0, ens.llm.vocab_size, (8, 32))
-        seg_b = torch.randint(0, ens.llm.vocab_size, (8, 32))
-        logs = ens.train_step(seg_a, seg_b, opt)
-        if s % 10 == 0:
-            print(
-                f"step {s:3d} | "
-                + " | ".join(f"{k} {v:.3f}" for k, v in logs.items())
-            )
-
-    for _ in range(5):
-        ens.memorize_ids(torch.randint(0, ens.llm.vocab_size, (1, 32)))
-    ens.new_task_adapter("medical")
-
-    q = torch.randint(0, ens.llm.vocab_size, (1, 8))
-    out, score, adapter, retries = ens.answer_ids(q, tau_consistency=-1.0)
-    print(f"\nadapter={adapter} jepa_consistency={score:.3f} retries={retries}")
-
-
-def demo_hf(model_path="Qwen/Qwen2.5-0.5B-Instruct"):
-    ens = Ensemble(llm=model_path, load_in_4bit=False)
-    opt = ens.make_optimizer()
-
-    texts = ["Replace this with your real corpus. " * 50]
-    seg_a, seg_b = segment_pairs_from_texts(ens.llm, texts, seg_len=32)
-    for s in range(10):
-        logs = ens.train_step(seg_a[:4], seg_b[:4], opt)
-        print(f"step {s} | " + " | ".join(f"{k} {v:.3f}" for k, v in logs.items()))
-
-    ens.memorize_text("The project codename is AURORA and it ships in Q3.")
-    ens.new_task_adapter("project_aurora")
-
-    text, score, adapter, retries = ens.answer_text(
-        "What is the project codename?", n_new=24, tau_consistency=-1.0
-    )
-    print(f"\n[{adapter} | jepa={score:.3f} | retries={retries}]\n{text}")
-
-
-if __name__ == "__main__":
-    import sys
-
-    from ensemble.config import load_dotenv, resolve_llm
-
-    load_dotenv()
-    arg = sys.argv[1] if len(sys.argv) > 1 else None
-    if arg is None or arg == "auto":
-        arg, preset = resolve_llm()
-        print(f"Resolved LLM: {arg} (preset {preset})")
-    if arg == "tiny":
-        demo_tiny()
-    else:
-        demo_hf(arg)

research/ensemble/src/ensemble/llm_emb_jepa_ensemble_pluggable.py

@@ -1,507 +0,0 @@
-"""
-LLM + Embedding + JEPA Ensemble — pluggable base-model edition
-==============================================================
-Now the LLM is a swappable BACKEND. Three ways to load it:
-
-    # 1. HuggingFace Hub id
-    ens = Ensemble(llm="Qwen/Qwen2.5-0.5B-Instruct")
-
-    # 2. Local path (e.g. downloaded Llama / converted checkpoint)
-    ens = Ensemble(llm="/models/llama-3.2-1b")
-
-    # 3. Toy fallback (no transformers needed, runs on CPU in seconds)
-    ens = Ensemble(llm="tiny")
-
-Requirements for real models:
-    pip install torch transformers peft accelerate
-    (optional 4-bit: pip install bitsandbytes -> load_in_4bit=True)
-
-Everything else (Embedder, JEPA, Bridge, VectorStore, Router, the
-JEPA-critic inference loop, continual-learning hooks) only touches
-token ids / hidden states / latents, so it works with ANY backend.
-"""
-
-from __future__ import annotations
-import copy
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-torch.manual_seed(0)
-
-# ----------------------------------------------------------------------------
-# 0. Backend interface — everything the ensemble needs from "an LLM"
-# ----------------------------------------------------------------------------
-class LLMBackend(nn.Module):
-    """Contract:
-        vocab_size : int
-        hidden_size: int
-        device     : torch.device
-        forward(ids)            -> (logits [B,T,V], hidden [B,T,H])
-        generate(ids, n_new)    -> ids [B, T+n_new]
-        add_adapter(name) / set_adapter(name)
-        trainable_parameters()  -> iterable of params to optimize
-        encode_text(str) / decode(ids)   (real backends only)
-    """
-    vocab_size: int
-    hidden_size: int
-
-
-# ----------------------------------------------------------------------------
-# 0a. HuggingFace backend (local path OR hub id) with PEFT LoRA adapters
-# ----------------------------------------------------------------------------
-class HFBackend(LLMBackend):
-    def __init__(self, model_path: str, *, load_in_4bit: bool = False,
-                 lora_r: int = 16, lora_alpha: int = 32,
-                 target_modules=("q_proj", "v_proj"),
-                 device: str | None = None, torch_dtype=None):
-        super().__init__()
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        from peft import LoraConfig, get_peft_model
-
-        self.device_ = torch.device(
-            device or ("cuda" if torch.cuda.is_available() else "cpu"))
-
-        kwargs = {}
-        if load_in_4bit:
-            from transformers import BitsAndBytesConfig
-            kwargs["quantization_config"] = BitsAndBytesConfig(
-                load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
-                bnb_4bit_quant_type="nf4")
-        if torch_dtype is not None:
-            kwargs["torch_dtype"] = torch_dtype
-
-        # `model_path` may be "Qwen/Qwen2.5-0.5B-Instruct", "meta-llama/...",
-        # or a local directory like "/models/llama-3.2-1b".
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-        base = AutoModelForCausalLM.from_pretrained(model_path, **kwargs)
-        if not load_in_4bit:
-            base.to(self.device_)
-
-        # Freeze the base; all learning happens in LoRA adapters.
-        for p in base.parameters():
-            p.requires_grad_(False)
-
-        self._lora_cfg = LoraConfig(
-            r=lora_r, lora_alpha=lora_alpha, lora_dropout=0.05,
-            target_modules=list(target_modules), task_type="CAUSAL_LM")
-        self.model = get_peft_model(base, self._lora_cfg, adapter_name="general")
-        self._adapters = {"general"}
-
-        self.vocab_size = self.model.config.vocab_size
-        self.hidden_size = self.model.config.hidden_size
-
-    # ---- adapters -----------------------------------------------------------
-    def add_adapter(self, name: str):
-        if name not in self._adapters:
-            self.model.add_adapter(name, self._lora_cfg)
-            self._adapters.add(name)
-
-    def set_adapter(self, name: str):
-        self.model.set_adapter(name)
-
-    def trainable_parameters(self):
-        return (p for p in self.model.parameters() if p.requires_grad)
-
-    # ---- core ops -----------------------------------------------------------
-    def forward(self, ids):
-        out = self.model(input_ids=ids.to(self.device_),
-                         output_hidden_states=True)
-        return out.logits, out.hidden_states[-1]      # last layer hidden
-
-    @torch.no_grad()
-    def generate(self, ids, n_new=64, temperature=0.8):
-        out = self.model.generate(
-            input_ids=ids.to(self.device_),
-            max_new_tokens=n_new, do_sample=True, temperature=temperature,
-            pad_token_id=self.tokenizer.pad_token_id)
-        return out
-
-    # ---- text helpers -------------------------------------------------------
-    def encode_text(self, text: str):
-        return self.tokenizer(text, return_tensors="pt").input_ids.to(self.device_)
-
-    def decode(self, ids):
-        return self.tokenizer.decode(ids[0], skip_special_tokens=True)
-
-    @property
-    def device(self):
-        return self.device_
-
-
-# ----------------------------------------------------------------------------
-# 0b. Tiny fallback backend (no transformers; same toy model as before)
-# ----------------------------------------------------------------------------
-class TinyBackend(LLMBackend):
-    VOCAB, D_MODEL, N_LAYERS, N_HEADS, SEQ_LEN, LORA_R = 1000, 128, 2, 4, 32, 8
-
-    class _LoRALinear(nn.Module):
-        def __init__(self, d_in, d_out, r):
-            super().__init__()
-            self.base = nn.Linear(d_in, d_out)
-            self.base.weight.requires_grad_(False)
-            self.base.bias.requires_grad_(False)
-            self.adapters, self.active, self.r = nn.ModuleDict(), None, r
-
-        def add_adapter(self, name):
-            A = nn.Linear(self.base.in_features, self.r, bias=False)
-            B = nn.Linear(self.r, self.base.out_features, bias=False)
-            nn.init.zeros_(B.weight)
-            self.adapters[name] = nn.Sequential(A, B)
-
-        def forward(self, x):
-            y = self.base(x)
-            if self.active and self.active in self.adapters:
-                y = y + self.adapters[self.active](x)
-            return y
-
-    class _Block(nn.Module):
-        def __init__(self, D, H, R):
-            super().__init__()
-            L = TinyBackend._LoRALinear
-            self.ln1 = nn.LayerNorm(D)
-            self.attn = nn.MultiheadAttention(D, H, batch_first=True)
-            self.ln2 = nn.LayerNorm(D)
-            self.up, self.down = L(D, 4 * D, R), L(4 * D, D, R)
-
-        def forward(self, x, mask):
-            h = self.ln1(x)
-            a, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
-            x = x + a
-            return x + self.down(F.gelu(self.up(self.ln2(x))))
-
-    def __init__(self):
-        super().__init__()
-        D, V = self.D_MODEL, self.VOCAB
-        self.tok = nn.Embedding(V, D)
-        self.pos = nn.Embedding(self.SEQ_LEN * 4, D)
-        self.blocks = nn.ModuleList(
-            [self._Block(D, self.N_HEADS, self.LORA_R) for _ in range(self.N_LAYERS)])
-        self.ln_f, self.head = nn.LayerNorm(D), nn.Linear(D, V, bias=False)
-        self.vocab_size, self.hidden_size = V, D
-        self.add_adapter("general")
-        self.set_adapter("general")
-
-    def add_adapter(self, name):
-        for b in self.blocks:
-            b.up.add_adapter(name); b.down.add_adapter(name)
-
-    def set_adapter(self, name):
-        for b in self.blocks:
-            b.up.active = name; b.down.active = name
-
-    def trainable_parameters(self):
-        return (p for p in self.parameters() if p.requires_grad)
-
-    def forward(self, ids):
-        B, T = ids.shape
-        x = self.tok(ids) + self.pos(torch.arange(T, device=ids.device))
-        mask = torch.triu(torch.full((T, T), float("-inf"), device=ids.device), 1)
-        for b in self.blocks:
-            x = b(x, mask)
-        h = self.ln_f(x)
-        return self.head(h), h
-
-    @torch.no_grad()
-    def generate(self, ids, n_new=16, temperature=1.0):
-        for _ in range(n_new):
-            logits, _ = self(ids[:, -self.SEQ_LEN:])
-            nxt = torch.multinomial(F.softmax(logits[:, -1] / temperature, -1), 1)
-            ids = torch.cat([ids, nxt], dim=1)
-        return ids
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
-
-def make_backend(llm: str, **kw) -> LLMBackend:
-    """'tiny' -> toy model; anything else -> HF hub id or local path."""
-    return TinyBackend() if llm == "tiny" else HFBackend(llm, **kw)
-
-
-# ----------------------------------------------------------------------------
-# 1. Embedder — vocab-agnostic (sized from the backend's tokenizer)
-#    Swap for a real model: pass embed_fn=lambda txt: sbert.encode(...)
-# ----------------------------------------------------------------------------
-class Embedder(nn.Module):
-    def __init__(self, vocab_size: int, d_emb: int = 64):
-        super().__init__()
-        self.tok = nn.Embedding(vocab_size, d_emb)
-        self.enc = nn.GRU(d_emb, d_emb, batch_first=True, bidirectional=True)
-        self.proj = nn.Linear(2 * d_emb, d_emb)
-        self.d_emb = d_emb
-
-    def forward(self, ids):
-        h, _ = self.enc(self.tok(ids))
-        return F.normalize(self.proj(h.mean(dim=1)), dim=-1)
-
-
-# ----------------------------------------------------------------------------
-# 2. JEPA — vocab-agnostic latent predictor with EMA target encoder
-# ----------------------------------------------------------------------------
-class _JEPAEncoder(nn.Module):
-    def __init__(self, vocab_size, d):
-        super().__init__()
-        self.tok = nn.Embedding(vocab_size, d)
-        self.enc = nn.GRU(d, d, batch_first=True)
-        self.out = nn.Linear(d, d)
-
-    def forward(self, ids):
-        h, _ = self.enc(self.tok(ids))
-        return self.out(h.mean(dim=1))
-
-
-class JEPA(nn.Module):
-    def __init__(self, vocab_size: int, d_jepa: int = 64, ema_m: float = 0.996):
-        super().__init__()
-        self.ctx_enc = _JEPAEncoder(vocab_size, d_jepa)
-        self.tgt_enc = copy.deepcopy(self.ctx_enc)
-        for p in self.tgt_enc.parameters():
-            p.requires_grad_(False)
-        self.predictor = nn.Sequential(
-            nn.Linear(d_jepa, 2 * d_jepa), nn.GELU(), nn.Linear(2 * d_jepa, d_jepa))
-        self.m, self.d_jepa = ema_m, d_jepa
-
-    @torch.no_grad()
-    def ema_update(self):
-        for p_t, p_c in zip(self.tgt_enc.parameters(), self.ctx_enc.parameters()):
-            p_t.mul_(self.m).add_(p_c.detach(), alpha=1 - self.m)
-
-    def loss(self, seg_ctx, seg_tgt):
-        z_hat = self.predictor(self.ctx_enc(seg_ctx))
-        with torch.no_grad():
-            z_tgt = self.tgt_enc(seg_tgt)
-        pred = F.mse_loss(z_hat, z_tgt)
-        var_reg = F.relu(1.0 - z_hat.std(dim=0)).mean()   # anti-collapse
-        return pred + 0.5 * var_reg
-
-    @torch.no_grad()
-    def predict_next_latent(self, seg_ctx):
-        return self.predictor(self.ctx_enc(seg_ctx))
-
-    @torch.no_grad()
-    def encode(self, seg):
-        return self.tgt_enc(seg)
-
-
-# ----------------------------------------------------------------------------
-# 3. Bridge — sized from backend.hidden_size at construction
-# ----------------------------------------------------------------------------
-class Bridge(nn.Module):
-    def __init__(self, d_llm_hidden: int, d_jepa: int):
-        super().__init__()
-        self.proj = nn.Sequential(
-            nn.Linear(d_llm_hidden, d_jepa), nn.GELU(), nn.Linear(d_jepa, d_jepa))
-
-    def forward(self, llm_hidden):                       # [B,T,H] -> [B,d_jepa]
-        return self.proj(llm_hidden.float().mean(dim=1))
-
-    def info_nce(self, z1, z2, tau=0.07):
-        z1, z2 = F.normalize(z1, dim=-1), F.normalize(z2, dim=-1)
-        logits = z1 @ z2.t() / tau
-        labels = torch.arange(z1.size(0), device=z1.device)
-        return 0.5 * (F.cross_entropy(logits, labels) +
-                      F.cross_entropy(logits.t(), labels))
-
-
-# ----------------------------------------------------------------------------
-# 4. Memory + Router
-# ----------------------------------------------------------------------------
-class VectorStore:
-    def __init__(self):
-        self.keys, self.values = [], []
-
-    def add(self, emb, payload):
-        self.keys.append(emb.squeeze(0).detach().cpu())
-        self.values.append(payload)
-
-    def search(self, q, k=2):
-        if not self.keys:
-            return []
-        K = torch.stack(self.keys)
-        sims = (q.detach().cpu() @ K.t()).squeeze(0)
-        top = sims.topk(min(k, len(self.keys))).indices
-        return [self.values[i] for i in top]
-
-
-class Router(nn.Module):
-    def __init__(self, d_emb, n_adapters):
-        super().__init__()
-        self.fc = nn.Linear(d_emb, n_adapters)
-
-    def forward(self, emb):
-        return self.fc(emb).argmax(dim=-1)
-
-
-# ----------------------------------------------------------------------------
-# 5. Ensemble — backend-agnostic
-# ----------------------------------------------------------------------------
-class Ensemble(nn.Module):
-    def __init__(self, llm: str = "tiny", adapter_names=("general",),
-                 d_emb: int = 64, d_jepa: int = 64, **backend_kw):
-        super().__init__()
-        self.llm = make_backend(llm, **backend_kw)
-        V, H = self.llm.vocab_size, self.llm.hidden_size
-
-        self.emb = Embedder(V, d_emb)
-        self.jepa = JEPA(V, d_jepa)
-        self.bridge = Bridge(H, d_jepa)
-        self.store = VectorStore()
-
-        self.adapter_names = list(adapter_names)
-        for n in self.adapter_names:
-            self.llm.add_adapter(n)
-        self.llm.set_adapter(self.adapter_names[0])
-        self.router = Router(d_emb, len(self.adapter_names))
-
-    # -------- inference: route -> retrieve -> generate -> JEPA-verify -------
-    @torch.no_grad()
-    def answer_ids(self, query_ids, n_new=32, tau_consistency=0.0, max_retries=2):
-        q_emb = self.emb(query_ids.cpu())
-        a_idx = self.router(q_emb).item()
-        self.llm.set_adapter(self.adapter_names[a_idx])
-
-        mems = self.store.search(q_emb, k=1)
-        ctx = (torch.cat([mems[0], query_ids.cpu()], dim=1)
-               if mems else query_ids.cpu())
-
-        z_expected = self.jepa.predict_next_latent(ctx)
-
-        best = None
-        for attempt in range(max_retries + 1):
-            draft = self.llm.generate(ctx.to(self.llm.device), n_new=n_new,
-                                      temperature=0.8 + 0.3 * attempt)
-            new_part = draft[:, ctx.size(1):].cpu()
-            score = F.cosine_similarity(
-                z_expected, self.jepa.encode(new_part)).item()
-            if best is None or score > best[1]:
-                best = (draft, score, attempt)
-            if score >= tau_consistency:
-                break
-        draft, score, attempt = best
-        return draft, score, self.adapter_names[a_idx], attempt
-
-    def answer_text(self, prompt: str, **kw):
-        """Convenience wrapper for HF backends (uses the real tokenizer)."""
-        ids = self.llm.encode_text(prompt)
-        out, score, adapter, retries = self.answer_ids(ids, **kw)
-        return self.llm.decode(out), score, adapter, retries
-
-    # -------- continual learning hooks ---------------------------------------
-    def memorize_ids(self, ids):
-        self.store.add(self.emb(ids.cpu()), ids.cpu())
-
-    def memorize_text(self, text: str):
-        self.memorize_ids(self.llm.encode_text(text))
-
-    def new_task_adapter(self, name: str):
-        self.adapter_names.append(name)
-        self.llm.add_adapter(name)
-        old = self.router
-        self.router = Router(self.emb.d_emb, len(self.adapter_names))
-        with torch.no_grad():
-            self.router.fc.weight[: old.fc.out_features] = old.fc.weight
-            self.router.fc.bias[: old.fc.out_features] = old.fc.bias
-
-    # -------- one joint training step (LM + JEPA + Bridge) -------------------
-    def train_step(self, seg_a, seg_b, opt, w_bridge=0.1):
-        """seg_a, seg_b: consecutive token-id segments [B, T] (same tokenizer
-        as the backend!). For HF backends build them with backend.tokenizer."""
-        logits, hidden = self.llm(seg_a.to(self.llm.device))
-        lm_loss = F.cross_entropy(
-            logits[:, :-1].reshape(-1, self.llm.vocab_size).float(),
-            seg_a[:, 1:].reshape(-1).to(logits.device))
-
-        jepa_loss = self.jepa.loss(seg_a.cpu(), seg_b.cpu())
-
-        z_llm = self.bridge(hidden.cpu() if hidden.device.type != "cpu" else hidden)
-        z_jepa = self.jepa.ctx_enc(seg_a.cpu()).detach()
-        bridge_loss = self.bridge.info_nce(z_llm, z_jepa.to(z_llm.device))
-
-        loss = lm_loss.cpu() + jepa_loss + w_bridge * bridge_loss
-        opt.zero_grad(); loss.backward(); opt.step()
-        self.jepa.ema_update()
-        return {"lm": lm_loss.item(), "jepa": jepa_loss.item(),
-                "bridge": bridge_loss.item()}
-
-    def make_optimizer(self, lr_lora=2e-4, lr_aux=1e-3):
-        return torch.optim.AdamW([
-            {"params": list(self.llm.trainable_parameters()), "lr": lr_lora},
-            {"params": list(self.jepa.ctx_enc.parameters())
-                     + list(self.jepa.predictor.parameters()), "lr": lr_aux},
-            {"params": list(self.bridge.parameters())
-                     + list(self.emb.parameters())
-                     + list(self.router.parameters()), "lr": lr_aux},
-        ])
-
-
-# ----------------------------------------------------------------------------
-# 6. Helpers: turn raw text into (seg_a, seg_b) pairs with the HF tokenizer
-# ----------------------------------------------------------------------------
-def segment_pairs_from_texts(backend: HFBackend, texts, seg_len=64):
-    """Yields consecutive-segment id pairs for the JEPA + LM losses."""
-    a_list, b_list = [], []
-    for t in texts:
-        ids = backend.tokenizer(t, return_tensors="pt").input_ids[0]
-        for i in range(0, len(ids) - 2 * seg_len, seg_len):
-            a_list.append(ids[i:i + seg_len])
-            b_list.append(ids[i + seg_len:i + 2 * seg_len])
-    if not a_list:
-        raise ValueError("texts too short for the chosen seg_len")
-    return torch.stack(a_list), torch.stack(b_list)
-
-
-# ----------------------------------------------------------------------------
-# 7. Demos
-# ----------------------------------------------------------------------------
-def demo_tiny(steps=50):
-    """No-dependency smoke test."""
-    ens = Ensemble(llm="tiny")
-    opt = ens.make_optimizer()
-    for s in range(steps):
-        seg_a = torch.randint(0, ens.llm.vocab_size, (8, 32))
-        seg_b = torch.randint(0, ens.llm.vocab_size, (8, 32))
-        logs = ens.train_step(seg_a, seg_b, opt)
-        if s % 10 == 0:
-            print(f"step {s:3d} | " + " | ".join(f"{k} {v:.3f}" for k, v in logs.items()))
-
-    for _ in range(5):
-        ens.memorize_ids(torch.randint(0, ens.llm.vocab_size, (1, 32)))
-    ens.new_task_adapter("medical")
-
-    q = torch.randint(0, ens.llm.vocab_size, (1, 8))
-    out, score, adapter, retries = ens.answer_ids(q, tau_consistency=-1.0)
-    print(f"\nadapter={adapter} jepa_consistency={score:.3f} retries={retries}")
-
-
-def demo_hf(model_path="Qwen/Qwen2.5-0.5B-Instruct"):
-    """Real model from hub id OR local path, e.g. '/models/llama-3.2-1b'.
-    For gated Llama repos: huggingface-cli login first."""
-    ens = Ensemble(llm=model_path, load_in_4bit=False)   # 4bit needs bitsandbytes
-    opt = ens.make_optimizer()
-
-    texts = ["Replace this with your real corpus. " * 50]
-    seg_a, seg_b = segment_pairs_from_texts(ens.llm, texts, seg_len=32)
-    for s in range(10):                                   # tiny demo run
-        logs = ens.train_step(seg_a[:4], seg_b[:4], opt)
-        print(f"step {s} | " + " | ".join(f"{k} {v:.3f}" for k, v in logs.items()))
-
-    ens.memorize_text("The project codename is AURORA and it ships in Q3.")
-    ens.new_task_adapter("project_aurora")
-
-    text, score, adapter, retries = ens.answer_text(
-        "What is the project codename?", n_new=24, tau_consistency=-1.0)
-    print(f"\n[{adapter} | jepa={score:.3f} | retries={retries}]\n{text}")
-
-
-if __name__ == "__main__":
-    import sys
-    arg = sys.argv[1] if len(sys.argv) > 1 else "tiny"
-    if arg == "tiny":
-        demo_tiny()
-    else:
-        demo_hf(arg)   # python ensemble.py /models/llama-3.2-1b
-                       # python ensemble.py Qwen/Qwen2.5-0.5B-Instruct

research/ensemble/src/ensemble/memory.py

@@ -1,46 +0,0 @@
-"""Retrieval memory: embedder, vector store, and adapter router."""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class Embedder(nn.Module):
-    def __init__(self, vocab_size: int, d_emb: int = 64):
-        super().__init__()
-        self.tok = nn.Embedding(vocab_size, d_emb)
-        self.enc = nn.GRU(d_emb, d_emb, batch_first=True, bidirectional=True)
-        self.proj = nn.Linear(2 * d_emb, d_emb)
-        self.d_emb = d_emb
-
-    def forward(self, ids):
-        h, _ = self.enc(self.tok(ids))
-        return F.normalize(self.proj(h.mean(dim=1)), dim=-1)
-
-
-class VectorStore:
-    def __init__(self):
-        self.keys, self.values = [], []
-
-    def add(self, emb, payload):
-        self.keys.append(emb.squeeze(0).detach().cpu())
-        self.values.append(payload)
-
-    def search(self, q, k=2):
-        if not self.keys:
-            return []
-        K = torch.stack(self.keys)
-        sims = (q.detach().cpu() @ K.t()).squeeze(0)
-        top = sims.topk(min(k, len(self.keys))).indices
-        return [self.values[i] for i in top]
-
-
-class Router(nn.Module):
-    def __init__(self, d_emb, n_adapters):
-        super().__init__()
-        self.fc = nn.Linear(d_emb, n_adapters)
-
-    def forward(self, emb):
-        return self.fc(emb).argmax(dim=-1)

research/ensemble/src/ensemble/pretrain.py

@@ -1,198 +0,0 @@
-"""Joint pretrain: LLM (LoRA) + embedder + JEPA + bridge, saved to models/ensemble/."""
-
-from __future__ import annotations
-
-import argparse
-import json
-import os
-import random
-import time
-from pathlib import Path
-
-import torch
-
-from ensemble.checkpoint import save_checkpoint
-from ensemble.config import default_ensemble_out, load_dotenv, resolve_llm
-from ensemble.jepa_ensemble import Ensemble, segment_pairs_from_texts
-
-_REPO_ROOT = Path(__file__).resolve().parents[4]
-_DEFAULT_DATA = _REPO_ROOT / "research/data/education-lesson-chat.jsonl"
-_DEFAULT_KB = _REPO_ROOT / "research/data/benchmark-kb.jsonl"
-
-
-def _load_jsonl(path: Path) -> list[dict]:
-    rows = []
-    with open(path) as f:
-        for line in f:
-            line = line.strip()
-            if line:
-                rows.append(json.loads(line))
-    return rows
-
-
-def _chat_to_text(row: dict) -> str:
-    messages = row.get("messages", [])
-    parts = [f"{m.get('role', 'user')}: {m.get('content', '')}" for m in messages]
-    return "\n".join(parts)
-
-
-def _collect_texts(data_path: Path, max_samples: int | None) -> list[str]:
-    rows = _load_jsonl(data_path)
-    if max_samples is not None:
-        rows = rows[:max_samples]
-    return [_chat_to_text(r) for r in rows if _chat_to_text(r).strip()]
-
-
-def _seed_memory(ens: Ensemble, kb_path: Path | None) -> int:
-    if kb_path is None or not kb_path.is_file():
-        return 0
-    count = 0
-    for row in _load_jsonl(kb_path):
-        text = row.get("text", "").strip()
-        if text:
-            ens.memorize_text(text)
-            count += 1
-    return count
-
-
-def pretrain(args) -> Path:
-    torch.manual_seed(args.seed)
-    random.seed(args.seed)
-
-    data_path = Path(args.data).resolve()
-    out_dir = Path(args.out).resolve()
-    kb_path = Path(args.kb).resolve() if args.kb else None
-
-    print(f"Loading ensemble backend: {args.llm}")
-    ens = Ensemble(llm=args.llm, load_in_4bit=args.load_in_4bit)
-    opt = ens.make_optimizer(lr_lora=args.lr_lora, lr_aux=args.lr_aux)
-
-    texts = _collect_texts(data_path, args.max_samples)
-    if not texts and args.llm != "tiny":
-        raise SystemExit(f"No training texts found in {data_path}")
-
-    mem_count = _seed_memory(ens, kb_path)
-    print(f"Training texts: {len(texts)} | memory snippets: {mem_count}")
-
-    if args.llm == "tiny":
-        n_pairs = max(args.steps * args.batch_size, args.batch_size)
-        v = ens.llm.vocab_size
-        seg_a = torch.randint(0, v, (n_pairs, args.seg_len))
-        seg_b = torch.randint(0, v, (n_pairs, args.seg_len))
-    else:
-        seg_a, seg_b = segment_pairs_from_texts(
-            ens.llm, texts, seg_len=args.seg_len
-        )
-    n_pairs = seg_a.size(0)
-    batch = min(args.batch_size, n_pairs)
-    print(f"Segment pairs: {n_pairs} | batch={batch} | steps={args.steps}")
-
-    t0 = time.time()
-    for step in range(args.steps):
-        idx = torch.randint(0, n_pairs, (batch,))
-        logs = ens.train_step(seg_a[idx], seg_b[idx], opt, w_bridge=args.w_bridge)
-        if step % max(1, args.log_every) == 0 or step == args.steps - 1:
-            parts = " | ".join(f"{k} {v:.4f}" for k, v in logs.items())
-            print(f"step {step:4d}/{args.steps} | {parts}")
-
-    elapsed = time.time() - t0
-    meta = {
-        "steps": args.steps,
-        "batch_size": batch,
-        "seg_len": args.seg_len,
-        "data": str(data_path),
-        "kb": str(kb_path) if kb_path else None,
-        "memory_count": mem_count,
-        "text_count": len(texts),
-        "elapsed_s": round(elapsed, 1),
-        "lr_lora": args.lr_lora,
-        "lr_aux": args.lr_aux,
-        "w_bridge": args.w_bridge,
-        "seed": args.seed,
-        "preset": getattr(args, "preset", None),
-    }
-
-    saved = save_checkpoint(
-        ens,
-        out_dir,
-        base_llm=args.llm,
-        training_meta=meta,
-    )
-    print(f"\nSaved ensemble checkpoint → {saved}")
-    print("Benchmark with slm-evals:")
-    print(
-        f"  uv run --package slm-evals slm-benchmark "
-        f"--model {saved} --model-type ensemble "
-        f"--benchmarks bfcl --max-samples 5"
-    )
-    return saved
-
-
-def parse_args():
-    p = argparse.ArgumentParser(
-        description="Pretrain JEPA ensemble (LLM+emb+JEPA) and save to models/ensemble/"
-    )
-    p.add_argument(
-        "--llm",
-        default=None,
-        help=(
-            "HF hub id / local path, 'tiny' for CPU smoke, or omit to use "
-            "LLM_PATH / BASE / MODEL_ID / ACTIVE_MODEL from .env + models.yaml"
-        ),
-    )
-    p.add_argument(
-        "--preset",
-        default=None,
-        help="models.yaml preset key (default: ENSEMBLE_PRESET or ACTIVE_MODEL)",
-    )
-    p.add_argument(
-        "--data",
-        default=str(_DEFAULT_DATA),
-        help="Chat JSONL (messages[]) for segment-pair training",
-    )
-    p.add_argument(
-        "--kb",
-        default=str(_DEFAULT_KB),
-        help="Optional KB JSONL (text field) loaded into vector store",
-    )
-    p.add_argument(
-        "--out",
-        default=None,
-        help="Output dir (default: ENSEMBLE_OUT or models/ensemble/<preset>-jepa-pretrain)",
-    )
-    p.add_argument("--steps", type=int, default=100)
-    p.add_argument("--batch-size", type=int, default=4)
-    p.add_argument("--seg-len", type=int, default=32)
-    p.add_argument("--max-samples", type=int, default=None)
-    p.add_argument("--lr-lora", type=float, default=2e-4)
-    p.add_argument("--lr-aux", type=float, default=1e-3)
-    p.add_argument("--w-bridge", type=float, default=0.1)
-    p.add_argument("--log-every", type=int, default=10)
-    p.add_argument("--seed", type=int, default=0)
-    p.add_argument("--load-in-4bit", action="store_true")
-    p.add_argument("--no-kb", action="store_true", help="Skip loading KB into memory")
-    return p.parse_args()
-
-
-def main():
-    load_dotenv()
-    args = parse_args()
-    if args.no_kb:
-        args.kb = None
-
-    preset_key = args.preset
-    if args.llm is None or args.llm == "auto":
-        args.llm, preset_key = resolve_llm(preset_arg=args.preset)
-    elif args.llm != "tiny" and not args.preset:
-        _, preset_key = resolve_llm(llm_arg=args.llm)
-
-    if not args.out:
-        args.out = os.environ.get("ENSEMBLE_OUT") or default_ensemble_out(preset_key)
-
-    args.preset = preset_key
-    print(f"Resolved LLM: {args.llm}" + (f" (preset {preset_key})" if preset_key else ""))
-    pretrain(args)
-
-
-if __name__ == "__main__":
-    main()

research/ensemble/src/ensemble/world_ensemble.py

@@ -1,228 +0,0 @@
-"""World-model ensemble: plan -> generate -> energy-rank."""
-
-from __future__ import annotations
-
-import sys
-import time
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from ensemble.backends import HFLLM, load_llm
-from ensemble.bridge import Bridge
-from ensemble.energy import EnergyModel
-from ensemble.jepa import JEPA
-from ensemble.memory import Embedder, VectorStore
-from ensemble.world_model import WorldModel
-
-torch.manual_seed(0)
-
-D_LAT = 96
-D_EMB = 64
-
-
-class WorldEnsemble(nn.Module):
-    def __init__(self, llm_spec="tiny"):
-        super().__init__()
-        self.llm = load_llm(llm_spec)
-        V, H = self.llm.vocab_size, self.llm.hidden_size
-        self.emb = Embedder(V, D_EMB)
-        self.jepa = JEPA(V, D_LAT)
-        self.world = WorldModel(D_LAT)
-        self.energy = EnergyModel(D_LAT)
-        self.bridge = Bridge(H, D_LAT)
-        self.store = VectorStore()
-
-    @torch.no_grad()
-    def world_state(self, segments):
-        s = self.world.init_state(1, "cpu")
-        for seg in segments:
-            z = self.jepa.encode(seg.cpu())
-            s, _ = self.world.step(s, z)
-        return s
-
-    @torch.no_grad()
-    def answer(self, query_ids, n_new=24, n_drafts=6, horizon=3):
-        q_emb = self.emb(query_ids.cpu())
-        mems = self.store.search(q_emb, k=1)
-        segments = (mems + [query_ids.cpu()]) if mems else [query_ids.cpu()]
-        ctx = torch.cat(segments, dim=1)
-
-        s = self.world_state(segments)
-        plan, _ = self.world.rollout(s, horizon)
-
-        drafts, lat = [], []
-        for _ in range(n_drafts):
-            out = self.llm.generate(
-                ctx.to(self.llm.device), n_new=n_new, temperature=0.9
-            )
-            new = out[:, ctx.size(1) :].cpu()
-            drafts.append(new)
-            lat.append(self.jepa.encode(new))
-        Z = torch.cat(lat, 0)
-        E = self.energy.rank(s, Z)
-        best = E.argmin().item()
-        return {
-            "output": drafts[best],
-            "energy": E[best].item(),
-            "all_energies": E.tolist(),
-            "plan_alignment": F.cosine_similarity(
-                plan[:, 0], Z[best : best + 1]
-            ).item(),
-        }
-
-    def memorize(self, ids):
-        self.store.add(self.emb(ids.cpu()), ids.cpu())
-
-    def train_step(
-        self,
-        seg_seq,
-        opt,
-        w=None,
-        hard_negs=True,
-    ):
-        if w is None:
-            w = dict(lm=1.0, jepa=1.0, world=1.0, ebm=1.0, bridge=0.1)
-
-        B, T, L = seg_seq.shape
-        dev = self.llm.device
-
-        flat = seg_seq[:, 0].to(dev)
-        logits, hidden = self.llm(flat)
-        lm = F.cross_entropy(
-            logits[:, :-1].reshape(-1, self.llm.vocab_size).float(),
-            flat[:, 1:].reshape(-1),
-        )
-
-        jepa = self.jepa.loss(seg_seq[:, 0], seg_seq[:, 1])
-
-        z_seq = torch.stack(
-            [self.jepa.enc(seg_seq[:, t]) for t in range(T)], 1
-        )
-        world = self.world.sequence_loss(z_seq)
-
-        s = self.world.init_state(B, z_seq.device)
-        s, _ = self.world.step(s, z_seq[:, 0].detach())
-        z_pos = z_seq[:, 1].detach()
-        z_negs = None
-        if hard_negs:
-            with torch.no_grad():
-                gen = self.llm.generate(seg_seq[:, 0].to(dev), n_new=L)
-                gen_new = gen[:, seg_seq.size(2) :].cpu()
-                z_negs = self.jepa.encode(gen_new).unsqueeze(1)
-        ebm = self.energy.contrastive_loss(s, z_pos, z_negs)
-
-        bridge = self.bridge.info_nce(
-            self.bridge(
-                hidden.cpu() if hidden.device.type != "cpu" else hidden
-            ),
-            self.jepa.enc(seg_seq[:, 0]).detach(),
-        )
-
-        loss = (
-            w["lm"] * lm.cpu()
-            + w["jepa"] * jepa
-            + w["world"] * world
-            + w["ebm"] * ebm
-            + w["bridge"] * bridge
-        )
-        opt.zero_grad()
-        loss.backward()
-        opt.step()
-        self.jepa.ema()
-        return dict(
-            lm=lm.item(),
-            jepa=jepa.item(),
-            world=world.item(),
-            ebm=ebm.item(),
-            bridge=bridge.item(),
-        )
-
-    def make_optimizer(self, lr_lora=2e-4, lr_aux=1e-3):
-        return torch.optim.AdamW(
-            [
-                {"params": list(self.llm.trainable_parameters()), "lr": lr_lora},
-                {
-                    "params": list(self.jepa.enc.parameters())
-                    + list(self.jepa.pred.parameters()),
-                    "lr": lr_aux,
-                },
-                {"params": list(self.world.parameters()), "lr": lr_aux},
-                {"params": list(self.energy.parameters()), "lr": lr_aux},
-                {
-                    "params": list(self.bridge.parameters())
-                    + list(self.emb.parameters()),
-                    "lr": lr_aux,
-                },
-            ]
-        )
-
-
-def toy_segment_sequences(B=8, T=4, L=24, vocab=1000):
-    return torch.randint(0, vocab, (B, T, L))
-
-
-def hf_segment_sequences(llm: HFLLM, texts, T=4, L=64):
-    seqs = []
-    for t in texts:
-        ids = llm.tokenizer(t, return_tensors="pt").input_ids[0]
-        n = (len(ids) // (T * L)) * T * L
-        if n:
-            seqs.append(ids[:n].view(-1, T, L))
-    if not seqs:
-        raise ValueError("corpus too short for T*L window")
-    return torch.cat(seqs, 0)
-
-
-def demo(spec="tiny", steps=60):
-    ens = WorldEnsemble(spec)
-    opt = ens.make_optimizer()
-
-    if spec == "tiny":
-        get_batch = lambda: toy_segment_sequences(vocab=ens.llm.vocab_size)
-    else:
-        corpus = ["Replace with your real documents. " * 200]
-        data = hf_segment_sequences(ens.llm, corpus, T=4, L=32)
-        get_batch = lambda: data[torch.randperm(len(data))[:4]]
-        steps = min(steps, 10)
-
-    t0 = time.time()
-    for s in range(steps):
-        logs = ens.train_step(
-            get_batch(), opt, hard_negs=(s > steps // 2)
-        )
-        if s % 10 == 0:
-            print(
-                f"step {s:3d} | "
-                + " | ".join(f"{k} {v:.3f}" for k, v in logs.items())
-            )
-    print(f"trained {steps} steps in {time.time() - t0:.1f}s")
-
-    for _ in range(4):
-        if spec == "tiny":
-            ens.memorize(torch.randint(0, ens.llm.vocab_size, (1, 24)))
-    q = (
-        torch.randint(0, ens.llm.vocab_size, (1, 12))
-        if spec == "tiny"
-        else ens.llm.tokenizer(
-            "What is this document about?", return_tensors="pt"
-        ).input_ids
-    )
-    res = ens.answer(q, n_drafts=6, horizon=3)
-    print(
-        f"\nselected draft energy={res['energy']:.3f} "
-        f"(all: {[f'{e:.2f}' for e in res['all_energies']]})"
-    )
-    print(f"plan↔output alignment: {res['plan_alignment']:.3f}")
-
-
-if __name__ == "__main__":
-    from ensemble.config import load_dotenv, resolve_llm
-
-    load_dotenv()
-    spec = sys.argv[1] if len(sys.argv) > 1 else None
-    if spec is None or spec == "auto":
-        spec, preset = resolve_llm()
-        print(f"Resolved LLM: {spec} (preset {preset})")
-    demo(spec or "tiny")

research/ensemble/src/ensemble/world_model.py

@@ -1,40 +0,0 @@
-"""Latent world model: multi-step rollout in JEPA space."""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class WorldModel(nn.Module):
-    def __init__(self, d_latent: int):
-        super().__init__()
-        self.cell = nn.GRUCell(d_latent, d_latent)
-        self.head = nn.Linear(d_latent, d_latent)
-        self.s0 = nn.Parameter(torch.zeros(d_latent))
-        self.d_latent = d_latent
-
-    def init_state(self, B, device):
-        return self.s0.unsqueeze(0).expand(B, -1).contiguous().to(device)
-
-    def step(self, s, z):
-        s = self.cell(z, s)
-        return s, self.head(s)
-
-    def rollout(self, s, horizon):
-        preds = []
-        for _ in range(horizon):
-            z_hat = self.head(s)
-            preds.append(z_hat)
-            s = self.cell(z_hat, s)
-        return torch.stack(preds, 1), s
-
-    def sequence_loss(self, z_seq):
-        B, T, _ = z_seq.shape
-        s = self.init_state(B, z_seq.device)
-        loss = 0.0
-        for t in range(T - 1):
-            s, z_hat = self.step(s, z_seq[:, t])
-            loss = loss + F.mse_loss(z_hat, z_seq[:, t + 1])
-        return loss / (T - 1)

research/ensemble/src/ensemble/world_model_ensemble.py

@@ -1,499 +0,0 @@
-"""
-World-Model Ensemble: EMB + EBM + JEPA + World Model + small LLM (from path)
-=============================================================================
-A LeCun-style modular agent built around a small language model.
-
-ARCHITECTURE
-------------
-                          ┌────────────────────────────┐
-   input tokens ──► EMB ──┤ VectorStore (retrieval/CL) │──► context
-        │                 └────────────────────────────┘      │
-        │                                                     │
-        ▼                                                     ▼
-   JEPA encoder ──► latent state s_t ──► WORLD MODEL ──► ŝ_{t+1..t+H}
-        │                 (GRU dynamics, multi-step rollout)   │
-        │                                                      │
-        │            ┌────────────────────────────────────┐   │
-        └──────────► │ ENERGY MODEL  E(s_ctx, z_candidate)│ ◄─┘
-                     │ low energy = compatible/plausible  │
-                     └────────────────┬───────────────────┘
-                                      │ scores drafts / plans
-                                      ▼
-   LLM (small, loaded from path, LoRA bank) ──► N drafts ──► pick argmin E
-
-ROLES
------
-EMB         perception for retrieval + routing (non-parametric memory)
-JEPA        learns the latent space: predict z(next segment) from z(context)
-            (EMA target encoder + variance reg, no token reconstruction)
-WORLD MODEL deterministic latent dynamics  s_{t+1} = f(s_t, z_t):
-            rolls the conversation/document state forward H steps in
-            LATENT space — cheap lookahead without decoding tokens
-ENERGY      E(s, z) ∈ R, trained so true continuations have LOW energy and
-            negatives (shuffled / model-generated) have HIGH energy.
-            At inference it is the critic: rank LLM drafts, reject bad plans.
-LLM         the only token-level generator. Loaded from a local path or HF id;
-            frozen base + LoRA adapters (continual learning by isolation).
-
-WHY EBM *and* JEPA?  JEPA gives a point prediction ẑ of the future latent;
-the EBM gives a *compatibility landscape* E(s, z) — it can say "both A and B
-are plausible" where a point predictor must average them. JEPA trains the
-representation; the EBM scores hypotheses in it. World model chains JEPA
-one-step predictions into multi-step rollouts that the EBM can evaluate.
-
-USAGE
------
-    pip install torch            # toy mode
-    pip install transformers peft accelerate   # real LLM mode
-
-    python world_model_ensemble.py tiny                 # smoke test
-    python world_model_ensemble.py /models/llama-3.2-1b # local weights
-    python world_model_ensemble.py Qwen/Qwen2.5-0.5B-Instruct
-"""
-
-from __future__ import annotations
-import copy
-import math
-import sys
-import time
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-torch.manual_seed(0)
-
-D_LAT = 96          # shared latent dimension (JEPA / world / energy)
-D_EMB = 64          # retrieval embedding dim
-
-
-# ============================================================================
-# 1. LLM backend — load small model from path / hub, or toy fallback
-#    (same contract as before: forward -> (logits, hidden), generate, adapters)
-# ============================================================================
-class TinyLLM(nn.Module):
-    VOCAB, D, L, H, T = 1000, 128, 2, 4, 32
-
-    def __init__(self):
-        super().__init__()
-        self.tok = nn.Embedding(self.VOCAB, self.D)
-        self.pos = nn.Embedding(self.T * 4, self.D)
-        layer = nn.TransformerEncoderLayer(self.D, self.H, 4 * self.D,
-                                           batch_first=True, norm_first=True)
-        self.blocks = nn.TransformerEncoder(layer, self.L)
-        self.head = nn.Linear(self.D, self.VOCAB, bias=False)
-        self.vocab_size, self.hidden_size = self.VOCAB, self.D
-
-    def forward(self, ids):
-        Tn = ids.size(1)
-        x = self.tok(ids) + self.pos(torch.arange(Tn, device=ids.device))
-        mask = torch.triu(torch.full((Tn, Tn), float("-inf"),
-                                     device=ids.device), 1)
-        h = self.blocks(x, mask=mask)
-        return self.head(h), h
-
-    @torch.no_grad()
-    def generate(self, ids, n_new=16, temperature=1.0):
-        for _ in range(n_new):
-            logits, _ = self(ids[:, -self.T:])
-            nxt = torch.multinomial(
-                F.softmax(logits[:, -1] / temperature, -1), 1)
-            ids = torch.cat([ids, nxt], 1)
-        return ids
-
-    def trainable_parameters(self):
-        return self.parameters()
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
-
-class HFLLM(nn.Module):
-    """Small model from a local path or HF id, frozen base + LoRA."""
-    def __init__(self, path, lora_r=16):
-        super().__init__()
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        from peft import LoraConfig, get_peft_model
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-        base = AutoModelForCausalLM.from_pretrained(
-            path, torch_dtype=torch.bfloat16
-            if torch.cuda.is_available() else torch.float32,
-            device_map="auto" if torch.cuda.is_available() else None)
-        for p in base.parameters():
-            p.requires_grad_(False)
-        cfg = LoraConfig(r=lora_r, lora_alpha=2 * lora_r, lora_dropout=0.05,
-                         target_modules=["q_proj", "v_proj"],
-                         task_type="CAUSAL_LM")
-        self.model = get_peft_model(base, cfg)
-        self.vocab_size = self.model.config.vocab_size
-        self.hidden_size = self.model.config.hidden_size
-
-    def forward(self, ids):
-        out = self.model(input_ids=ids.to(self.device),
-                         output_hidden_states=True)
-        return out.logits, out.hidden_states[-1]
-
-    @torch.no_grad()
-    def generate(self, ids, n_new=32, temperature=0.8):
-        return self.model.generate(
-            input_ids=ids.to(self.device), max_new_tokens=n_new,
-            do_sample=True, temperature=temperature,
-            pad_token_id=self.tokenizer.pad_token_id)
-
-    def trainable_parameters(self):
-        return (p for p in self.model.parameters() if p.requires_grad)
-
-    @property
-    def device(self):
-        return next(self.model.parameters()).device
-
-
-def load_llm(spec: str):
-    return TinyLLM() if spec == "tiny" else HFLLM(spec)
-
-
-# ============================================================================
-# 2. Embedder (retrieval) — vocab-agnostic
-# ============================================================================
-class Embedder(nn.Module):
-    def __init__(self, vocab):
-        super().__init__()
-        self.tok = nn.Embedding(vocab, D_EMB)
-        self.gru = nn.GRU(D_EMB, D_EMB, batch_first=True, bidirectional=True)
-        self.out = nn.Linear(2 * D_EMB, D_EMB)
-
-    def forward(self, ids):
-        h, _ = self.gru(self.tok(ids))
-        return F.normalize(self.out(h.mean(1)), dim=-1)
-
-
-class VectorStore:
-    def __init__(self):
-        self.K, self.V = [], []
-
-    def add(self, k, v):
-        self.K.append(k.squeeze(0).detach().cpu()); self.V.append(v)
-
-    def search(self, q, k=1):
-        if not self.K:
-            return []
-        sims = (q.detach().cpu() @ torch.stack(self.K).t()).squeeze(0)
-        return [self.V[i] for i in sims.topk(min(k, len(self.K))).indices]
-
-
-# ============================================================================
-# 3. JEPA — owns the latent space (EMA target encoder, variance-regularized)
-# ============================================================================
-class SegEncoder(nn.Module):
-    def __init__(self, vocab):
-        super().__init__()
-        self.tok = nn.Embedding(vocab, D_LAT)
-        self.gru = nn.GRU(D_LAT, D_LAT, batch_first=True)
-        self.out = nn.Linear(D_LAT, D_LAT)
-
-    def forward(self, ids):
-        h, _ = self.gru(self.tok(ids))
-        return self.out(h.mean(1))                       # [B, D_LAT]
-
-
-class JEPA(nn.Module):
-    def __init__(self, vocab, m=0.996):
-        super().__init__()
-        self.enc = SegEncoder(vocab)                     # context/online enc
-        self.tgt = copy.deepcopy(self.enc)               # EMA target
-        for p in self.tgt.parameters():
-            p.requires_grad_(False)
-        self.pred = nn.Sequential(nn.Linear(D_LAT, 2 * D_LAT), nn.GELU(),
-                                  nn.Linear(2 * D_LAT, D_LAT))
-        self.m = m
-
-    @torch.no_grad()
-    def ema(self):
-        for pt, pc in zip(self.tgt.parameters(), self.enc.parameters()):
-            pt.mul_(self.m).add_(pc.detach(), alpha=1 - self.m)
-
-    def loss(self, seg_a, seg_b):
-        z_hat = self.pred(self.enc(seg_a))
-        with torch.no_grad():
-            z_tgt = self.tgt(seg_b)
-        var = F.relu(1.0 - z_hat.std(0)).mean()          # anti-collapse
-        return F.mse_loss(z_hat, z_tgt) + 0.5 * var
-
-    @torch.no_grad()
-    def encode(self, seg):                               # target space
-        return self.tgt(seg)
-
-
-# ============================================================================
-# 4. WORLD MODEL — latent dynamics s_{t+1} = f(s_t, z_t), multi-step rollout
-#    Trained on SEQUENCES of segments: predict each next latent from state.
-# ============================================================================
-class WorldModel(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.cell = nn.GRUCell(D_LAT, D_LAT)             # state update
-        self.head = nn.Linear(D_LAT, D_LAT)              # state -> ẑ_{t+1}
-        self.s0 = nn.Parameter(torch.zeros(D_LAT))
-
-    def init_state(self, B, device):
-        return self.s0.unsqueeze(0).expand(B, -1).contiguous().to(device)
-
-    def step(self, s, z):
-        """Consume observed latent z_t, return (new state, prediction ẑ_{t+1})."""
-        s = self.cell(z, s)
-        return s, self.head(s)
-
-    def rollout(self, s, horizon):
-        """Imagine H future latents feeding its own predictions back in."""
-        preds = []
-        for _ in range(horizon):
-            z_hat = self.head(s)
-            preds.append(z_hat)
-            s = self.cell(z_hat, s)
-        return torch.stack(preds, 1), s                  # [B, H, D_LAT]
-
-    def sequence_loss(self, z_seq):
-        """z_seq: [B, T, D_LAT] observed segment latents (teacher forcing)."""
-        B, T, _ = z_seq.shape
-        s = self.init_state(B, z_seq.device)
-        loss = 0.0
-        for t in range(T - 1):
-            s, z_hat = self.step(s, z_seq[:, t])
-            loss = loss + F.mse_loss(z_hat, z_seq[:, t + 1])
-        return loss / (T - 1)
-
-
-# ============================================================================
-# 5. ENERGY MODEL — E(state, candidate latent) ∈ R, low = plausible
-#    Trained with InfoNCE-style contrastive: positives = true next latent,
-#    negatives = (a) other batch items, (b) LLM-generated drafts (optional).
-# ============================================================================
-class EnergyModel(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(2 * D_LAT, 2 * D_LAT), nn.GELU(),
-            nn.Linear(2 * D_LAT, D_LAT), nn.GELU(),
-            nn.Linear(D_LAT, 1))
-
-    def energy(self, s, z):
-        """s: [B, D_LAT] context state; z: [B, D_LAT] candidate. -> [B]"""
-        return self.net(torch.cat([s, z], -1)).squeeze(-1)
-
-    def contrastive_loss(self, s, z_pos, z_negs=None, tau=0.5):
-        """Softmax over energies: true continuation must be the argmin.
-        In-batch negatives: every other item's z_pos is a negative for s_i."""
-        B = s.size(0)
-        # pairwise energies: E(s_i, z_j) for all i, j
-        s_rep = s.unsqueeze(1).expand(B, B, D_LAT).reshape(B * B, D_LAT)
-        z_rep = z_pos.unsqueeze(0).expand(B, B, D_LAT).reshape(B * B, D_LAT)
-        E = self.energy(s_rep, z_rep).view(B, B)         # [B, B]
-        if z_negs is not None:                           # extra hard negatives
-            En = self.energy(
-                s.repeat_interleave(z_negs.size(1), 0),
-                z_negs.reshape(-1, D_LAT)).view(B, -1)
-            E = torch.cat([E, En], dim=1)
-        labels = torch.arange(B, device=s.device)
-        return F.cross_entropy(-E / tau, labels)         # low E ⇒ high logit
-
-    @torch.no_grad()
-    def rank(self, s, candidates):
-        """candidates: [N, D_LAT]; returns energies [N] (lower = better)."""
-        return self.energy(s.expand(candidates.size(0), -1), candidates)
-
-
-# ============================================================================
-# 6. Bridge — LLM hidden states -> shared latent space (alignment)
-# ============================================================================
-class Bridge(nn.Module):
-    def __init__(self, d_hidden):
-        super().__init__()
-        self.proj = nn.Sequential(nn.Linear(d_hidden, D_LAT), nn.GELU(),
-                                  nn.Linear(D_LAT, D_LAT))
-
-    def forward(self, h):                                # [B,T,H] -> [B,D_LAT]
-        return self.proj(h.float().mean(1))
-
-    def info_nce(self, a, b, tau=0.07):
-        a, b = F.normalize(a, -1), F.normalize(b, -1)
-        logits = a @ b.t() / tau
-        y = torch.arange(a.size(0), device=a.device)
-        return 0.5 * (F.cross_entropy(logits, y) +
-                      F.cross_entropy(logits.t(), y))
-
-
-# ============================================================================
-# 7. THE ENSEMBLE — wiring + inference (plan -> generate -> energy-rank)
-# ============================================================================
-class WorldEnsemble(nn.Module):
-    def __init__(self, llm_spec="tiny"):
-        super().__init__()
-        self.llm = load_llm(llm_spec)
-        V, H = self.llm.vocab_size, self.llm.hidden_size
-        self.emb = Embedder(V)
-        self.jepa = JEPA(V)
-        self.world = WorldModel()
-        self.energy = EnergyModel()
-        self.bridge = Bridge(H)
-        self.store = VectorStore()
-
-    # ------------------------- inference ---------------------------------
-    @torch.no_grad()
-    def world_state(self, segments):
-        """Fold a list of [1,T] segment tensors into a latent state."""
-        s = self.world.init_state(1, "cpu")
-        for seg in segments:
-            z = self.jepa.encode(seg.cpu())
-            s, _ = self.world.step(s, z)
-        return s
-
-    @torch.no_grad()
-    def answer(self, query_ids, n_new=24, n_drafts=6, horizon=3):
-        """retrieve -> build world state -> imagine -> generate N -> argmin E."""
-        q_emb = self.emb(query_ids.cpu())
-        mems = self.store.search(q_emb, k=1)
-        segments = (mems + [query_ids.cpu()]) if mems else [query_ids.cpu()]
-        ctx = torch.cat(segments, dim=1)
-
-        s = self.world_state(segments)                   # latent context state
-        plan, _ = self.world.rollout(s, horizon)         # imagined future
-        # (plan is available for planning losses / steering; logged here)
-
-        drafts, lat = [], []
-        for _ in range(n_drafts):
-            out = self.llm.generate(ctx.to(self.llm.device), n_new=n_new,
-                                    temperature=0.9)
-            new = out[:, ctx.size(1):].cpu()
-            drafts.append(new)
-            lat.append(self.jepa.encode(new))
-        Z = torch.cat(lat, 0)                            # [N, D_LAT]
-        E = self.energy.rank(s, Z)                       # lower = better
-        best = E.argmin().item()
-        return {"output": drafts[best], "energy": E[best].item(),
-                "all_energies": E.tolist(),
-                "plan_alignment": F.cosine_similarity(
-                    plan[:, 0], Z[best:best + 1]).item()}
-
-    def memorize(self, ids):
-        self.store.add(self.emb(ids.cpu()), ids.cpu())
-
-    # ------------------------- training ----------------------------------
-    def train_step(self, seg_seq, opt, w=dict(lm=1.0, jepa=1.0, world=1.0,
-                                              ebm=1.0, bridge=0.1),
-                   hard_negs=True):
-        """seg_seq: [B, T_seg, L] — B documents, each split into T_seg
-        consecutive segments of length L (same tokenizer as the LLM)."""
-        B, T, L = seg_seq.shape
-        dev = self.llm.device
-
-        # (1) LM loss on the first segment (or all, batched, if budget allows)
-        flat = seg_seq[:, 0].to(dev)
-        logits, hidden = self.llm(flat)
-        lm = F.cross_entropy(
-            logits[:, :-1].reshape(-1, self.llm.vocab_size).float(),
-            flat[:, 1:].reshape(-1))
-
-        # (2) JEPA: adjacent segment pairs
-        jepa = self.jepa.loss(seg_seq[:, 0], seg_seq[:, 1])
-
-        # (3) World model: sequence of latents (online encoder, grads flow)
-        z_seq = torch.stack([self.jepa.enc(seg_seq[:, t])
-                             for t in range(T)], 1)      # [B, T, D_LAT]
-        world = self.world.sequence_loss(z_seq)
-
-        # (4) Energy: state after t=0 must give low E to true z_1,
-        #     high E to in-batch + (optionally) LLM-generated negatives
-        s = self.world.init_state(B, z_seq.device)
-        s, _ = self.world.step(s, z_seq[:, 0].detach())
-        z_pos = z_seq[:, 1].detach()
-        z_negs = None
-        if hard_negs:
-            with torch.no_grad():                        # model drafts as negs
-                gen = self.llm.generate(seg_seq[:, 0].to(dev), n_new=L)
-                gen_new = gen[:, seg_seq.size(2):].cpu()
-                z_negs = self.jepa.encode(gen_new).unsqueeze(1)  # [B,1,D]
-        ebm = self.energy.contrastive_loss(s, z_pos, z_negs)
-
-        # (5) Bridge: align LLM hidden(seg0) with JEPA latent(seg0)
-        bridge = self.bridge.info_nce(
-            self.bridge(hidden.cpu() if hidden.device.type != "cpu" else hidden),
-            self.jepa.enc(seg_seq[:, 0]).detach())
-
-        loss = (w["lm"] * lm.cpu() + w["jepa"] * jepa + w["world"] * world
-                + w["ebm"] * ebm + w["bridge"] * bridge)
-        opt.zero_grad(); loss.backward(); opt.step()
-        self.jepa.ema()
-        return dict(lm=lm.item(), jepa=jepa.item(), world=world.item(),
-                    ebm=ebm.item(), bridge=bridge.item())
-
-    def make_optimizer(self, lr_lora=2e-4, lr_aux=1e-3):
-        return torch.optim.AdamW([
-            {"params": list(self.llm.trainable_parameters()), "lr": lr_lora},
-            {"params": list(self.jepa.enc.parameters())
-                     + list(self.jepa.pred.parameters()), "lr": lr_aux},
-            {"params": list(self.world.parameters()), "lr": lr_aux},
-            {"params": list(self.energy.parameters()), "lr": lr_aux},
-            {"params": list(self.bridge.parameters())
-                     + list(self.emb.parameters()), "lr": lr_aux}])
-
-
-# ============================================================================
-# 8. Data helpers + demo
-# ============================================================================
-def toy_segment_sequences(B=8, T=4, L=24, vocab=1000):
-    """Random docs split into T consecutive segments. Replace with real
-    corpus: tokenize each document, reshape into [T, L] windows."""
-    return torch.randint(0, vocab, (B, T, L))
-
-
-def hf_segment_sequences(llm: HFLLM, texts, T=4, L=64):
-    seqs = []
-    for t in texts:
-        ids = llm.tokenizer(t, return_tensors="pt").input_ids[0]
-        n = (len(ids) // (T * L)) * T * L
-        if n:
-            seqs.append(ids[:n].view(-1, T, L))
-    if not seqs:
-        raise ValueError("corpus too short for T*L window")
-    return torch.cat(seqs, 0)
-
-
-def demo(spec="tiny", steps=60):
-    ens = WorldEnsemble(spec)
-    opt = ens.make_optimizer()
-
-    if spec == "tiny":
-        get_batch = lambda: toy_segment_sequences(vocab=ens.llm.vocab_size)
-    else:
-        corpus = ["Replace with your real documents. " * 200]
-        data = hf_segment_sequences(ens.llm, corpus, T=4, L=32)
-        get_batch = lambda: data[torch.randperm(len(data))[:4]]
-        steps = min(steps, 10)
-
-    t0 = time.time()
-    for s in range(steps):
-        logs = ens.train_step(get_batch(), opt,
-                              hard_negs=(s > steps // 2))  # warmup w/o negs
-        if s % 10 == 0:
-            print(f"step {s:3d} | " +
-                  " | ".join(f"{k} {v:.3f}" for k, v in logs.items()))
-    print(f"trained {steps} steps in {time.time()-t0:.1f}s")
-
-    # memory + inference
-    for _ in range(4):
-        if spec == "tiny":
-            ens.memorize(torch.randint(0, ens.llm.vocab_size, (1, 24)))
-    q = (torch.randint(0, ens.llm.vocab_size, (1, 12)) if spec == "tiny"
-         else ens.llm.tokenizer("What is this document about?",
-                                return_tensors="pt").input_ids)
-    res = ens.answer(q, n_drafts=6, horizon=3)
-    print(f"\nselected draft energy={res['energy']:.3f} "
-          f"(all: {[f'{e:.2f}' for e in res['all_energies']]})")
-    print(f"plan↔output alignment: {res['plan_alignment']:.3f}")
-
-
-if __name__ == "__main__":
-    demo(sys.argv[1] if len(sys.argv) > 1 else "tiny")

research/eval_harness.py

@@ -1,6 +0,0 @@
-"""Deprecated shim — use `ensemble.eval.jepa_harness` instead."""
-
-from ensemble.eval.jepa_harness import run, parse_args
-
-if __name__ == "__main__":
-    run(parse_args())

research/evals/USAGE.md

@@ -197,12 +197,6 @@ uv run --package slm-evals slm-lm-eval \
   --model openbmb/MiniCPM5-1B \
   --adapter ./models/finetuned/minicpm5-1b-lora \
   --experiment-name minicpm5-1b-lora__manual
-
-# Ensemble checkpoint (manifest.json auto-detected)
-uv run --package slm-evals slm-lm-eval \
-  --config research/evals/configs/lm_eval_smoke.yaml \
-  --model ./models/ensemble/jepa-lesson-pretrain \
-  --experiment-name ensemble-jepa__lm-eval
 ```
 
 ### Compare baseline vs candidate
@@ -259,8 +253,8 @@ slm-lm-eval [OPTIONS]
 --list-tasks-all        Full lm-eval task list
 --profile NAME          Shorthand for --config (reasoning, code, smoke, …)
 --config PATH           YAML config (tasks, seed, limit, …)
---preset KEY            models.yaml preset (base, LoRA, merged, ensemble)
---model PATH            HF Hub id, merged dir, or ensemble checkpoint
+--preset KEY            models.yaml preset (base, LoRA, merged)
+--model PATH            HF Hub id or merged checkpoint dir
 --adapter PATH          LoRA adapter (alternative to preset adapter_path)
 --tasks NAMES           Override task list
 --num-fewshot N
@@ -285,12 +279,6 @@ Each run writes to `<output_dir>/<experiment_name>/`:
 | `run_meta.json` | Preset, base model, adapter, tasks, seed |
 | `comparison.md` | Delta table (when `--compare-to` set) |
 
-### Ensemble backend notes
-
-- **`ensemble-lm`** loads JEPA checkpoints via `manifest.json`.
-- **`generate_until`** tasks (e.g. `gsm8k`) use the full ensemble stack (`generate_text`).
-- **`loglikelihood`** tasks (e.g. `arc_easy`, `hellaswag`) score the underlying HF LLM head (adapter 0), not the JEPA selector. Use [`jepa_harness`](../ensemble/README.md) to measure selector value on domain QA.
-
 ### PEFT / LoRA
 
 lm-eval expects `pretrained=<base>,peft=<adapter>`. The preset resolver handles this for keys like `minicpm5-1b-lesson-lora`. Merged checkpoints use `--preset minicpm5-1b-lesson-merged` or `--model ./models/finetuned/...-merged`.

research/evals/configs/ensemble_jepa_lesson.yaml

@@ -1,24 +0,0 @@
-# JEPA ensemble checkpoint (models/ensemble/jepa-lesson-pretrain)
-# Pretrain: uv run --package ensemble ensemble-pretrain --llm Qwen/Qwen2.5-0.5B-Instruct
-# Compare baseline: copy this file, set model_path to the base Hub id and model_type: hf
-
-model_path: "./models/ensemble/jepa-lesson-pretrain"
-model_type: "ensemble"
-device: "auto"
-dtype: "bfloat16"
-
-max_new_tokens: 512
-temperature: 0.0
-
-experiment_name: "jepa-ensemble-lesson__bfcl-tau__v1"
-output_dir: "results"
-
-benchmarks:
-  - bfcl
-  - tau_bench
-
-max_samples: 20
-
-benchmark_overrides:
-  tau_bench:
-    use_llm_user: false