Merge pull request #4 from MSghais/experiment/small_model_building_testing

.cursor/plans/ensemble_refactor_benchmark_be904eec.plan.md

@@ -0,0 +1,267 @@
+---
+name: Ensemble refactor benchmark
+overview: Compare Option A vs Option B, then reorganize the three research scripts into a shared `research/ensemble/` package (Option B) and add a clear run + benchmark workflow from CPU smoke tests through real-model ablations.
+todos:
+  - id: scaffold-package
+    content: "Create research/ensemble/ package: pyproject.toml, workspace member, ensemble dependency group"
+    status: completed
+  - id: extract-shared
+    content: Extract shared modules (backends, memory, jepa, bridge, eval/metrics) from the three existing scripts
+    status: completed
+  - id: migrate-jepa
+    content: Move Ensemble + demo into jepa_ensemble.py; migrate eval_harness to eval/jepa_harness.py
+    status: completed
+  - id: migrate-world
+    content: Move WorldEnsemble into world_ensemble.py; add world_model.py and energy.py
+    status: completed
+  - id: world-harness
+    content: Add eval/world_harness.py with energy vs random vs oracle selector comparison
+    status: completed
+  - id: benchmark-data
+    content: Add research/data/benchmark-qa.jsonl and benchmark-kb.jsonl from lesson topics
+    status: completed
+  - id: run-docs
+    content: Add research/ensemble/README.md, scripts/smoke.sh, and .env.example ensemble vars with Tier 1–3 commands
+    status: completed
+  - id: cleanup-shims
+    content: Remove or thin re-export old top-level research/*.py files; verify smoke + toy harness pass
+    status: completed
+isProject: false
+---
+
+# Ensemble package (Option B) + run/benchmark plan
+
+## Option A vs Option B
+
+| | **Option A** (folder split only) | **Option B** (shared package) |
+|---|---|---|
+| **Layout** | `research/ensemble/jepa/` + `research/ensemble/world_model/` — move files, fix imports | `research/ensemble/` with shared `backends.py`, `jepa.py`, `memory.py`, etc. |
+| **Effort** | ~1 hour | ~3–4 hours |
+| **Duplication** | Keeps ~400 lines duplicated (two LLM backends, two JEPAs, two VectorStores) | Single implementation; bugfixes apply to both ensembles |
+| **Imports** | Still fragile (`from jepa.ensemble import ...`) | Clean: `from ensemble.jepa_ensemble import Ensemble` |
+| **Benchmarking** | Works for JEPA only ([`research/eval_harness.py`](research/eval_harness.py)); world model has no harness | Shared [`research/ensemble/eval/metrics.py`](research/ensemble/eval/metrics.py) + harnesses for both tracks |
+| **Future agent hook** | Harder — two divergent codepaths | One package to import from `libs/agent` later if needed |
+| **When to pick A** | You only need tidy folders for a demo and will not touch world-model eval | — |
+| **When to pick B** | You will run ablations, add world-model benchmarks, or train bridge checkpoints | **Your choice** |
+
+**Recommendation:** Go straight to **Option B under `research/ensemble/`** (experiments stay out of `libs/`). Option A is a reasonable **fallback** if time runs out: do the folder split and defer deduplication.
+
+```mermaid
+flowchart LR
+  subgraph shared [research/ensemble shared]
+    backends[backends.py]
+    memory[memory.py]
+    jepa_mod[jepa.py]
+    bridge[bridge.py]
+    metrics[eval/metrics.py]
+  end
+
+  subgraph jepa_track [JEPA track]
+    jepa_ens[jepa_ensemble.py]
+    jepa_eval[eval/jepa_harness.py]
+  end
+
+  subgraph world_track [World-model track]
+    world_mod[world_model.py]
+    energy[energy.py]
+    world_ens[world_ensemble.py]
+    world_eval[eval/world_harness.py]
+  end
+
+  backends --> jepa_ens
+  memory --> jepa_ens
+  jepa_mod --> jepa_ens
+  bridge --> jepa_ens
+  jepa_ens --> jepa_eval
+  metrics --> jepa_eval
+
+  backends --> world_ens
+  memory --> world_ens
+  jepa_mod --> world_ens
+  world_mod --> world_ens
+  energy --> world_ens
+  bridge --> world_ens
+  world_ens --> world_eval
+  metrics --> world_eval
+```
+
+---
+
+## Target layout
+
+```
+research/
+  ensemble/
+    pyproject.toml              # workspace member, optional deps
+    README.md                   # run + benchmark commands
+    ensemble/
+      __init__.py
+      backends.py               # TinyBackend + HFBackend (from jepa file)
+      memory.py                   # Embedder, VectorStore, Router
+      jepa.py                     # JEPA module
+      bridge.py
+      jepa_ensemble.py            # Ensemble class (from llm_emb_jepa_ensemble_pluggable.py)
+      world_model.py              # WorldModel
+      energy.py                   # EnergyModel
+      world_ensemble.py           # WorldEnsemble (from world_model_ensemble.py)
+      eval/
+        __init__.py
+        metrics.py                # EM, F1, paired_bootstrap (from eval_harness)
+        jepa_harness.py           # ablation ladder + best-of-N (from eval_harness)
+        world_harness.py          # NEW: energy vs random vs oracle on shared drafts
+    scripts/
+      smoke.sh                    # toy CPU checks for both ensembles
+  data/
+    education-lesson-chat.jsonl   # existing
+    benchmark-qa.jsonl            # NEW: short QA for eval (derived from lesson topics)
+    benchmark-kb.jsonl            # NEW: 1–2 sentence facts per topic for RAG
+  finetune.py                     # unchanged
+```
+
+Delete after migration (or leave thin re-export shims for one release):
+- [`research/llm_emb_jepa_ensemble_pluggable.py`](research/llm_emb_jepa_ensemble_pluggable.py)
+- [`research/eval_harness.py`](research/eval_harness.py)
+- [`research/world_model_ensemble.py`](research/world_model_ensemble.py)
+
+---
+
+## Package wiring
+
+1. Add [`research/ensemble/pyproject.toml`](research/ensemble/pyproject.toml) with `name = "ensemble"`, `torch` required, `transformers`/`peft`/`accelerate` optional (same pattern as root [`finetune`](pyproject.toml) group).
+2. Extend root [`pyproject.toml`](pyproject.toml):
+   - `[tool.uv.workspace] members` → add `"research/ensemble"`
+   - `[dependency-groups] ensemble = [...]` mirroring `finetune`
+   - Optional root dep: `"ensemble"` for `uv run` convenience
+3. Add ensemble env vars to [`.env.example`](.env.example) (model path, QA/KB paths, checkpoint path).
+
+---
+
+## How to try the models (3 tiers)
+
+### Tier 1 — Smoke (CPU, no HF download, ~30s)
+
+Validates imports and inference plumbing after refactor.
+
+```bash
+uv sync --group ensemble
+uv run --package ensemble python -m ensemble.jepa_ensemble tiny      # train 50 steps + answer
+uv run --package ensemble python -m ensemble.world_ensemble tiny     # train 60 steps + answer
+bash research/ensemble/scripts/smoke.sh
+```
+
+Uses `TinyBackend` / `TinyLLM` — random weights, synthetic segments. Confirms modules load and forward/generate paths work.
+
+### Tier 2 — Micro demo (real small model, ~2–5 min)
+
+Quick “does it run on GPU/CPU with a real tokenizer?”
+
+```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
+```
+
+Or a local/finetuned path from [`models.yaml`](models.yaml) / `FINETUNE_OUT`:
+
+```bash
+uv run --package ensemble python -m ensemble.jepa_ensemble ./models/finetuned/minicpm5-1b-lora-merged
+```
+
+### Tier 3 — Benchmark (ablation + significance)
+
+**JEPA ablation ladder** (existing logic, moved to `eval/jepa_harness.py`):
+
+```bash
+# Toy benchmark (no download)
+uv run --package ensemble python -m ensemble.eval.jepa_harness \
+  --llm tiny --toy --limit 20 --n_drafts 8
+
+# Real model + project-aligned 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
+
+# With bridge-trained checkpoint (C5)
+uv run --package ensemble python -m ensemble.eval.jepa_harness \
+  --llm ./models/finetuned/minicpm5-1b-lora-merged \
+  --qa research/data/benchmark-qa.jsonl \
+  --kb research/data/benchmark-kb.jsonl \
+  --ckpt ./checkpoints/ensemble_bridge.pt
+```
+
+**World-model benchmark** (new `eval/world_harness.py`, parallel API):
+
+- Same QA/KB inputs
+- Compare selectors on shared drafts: `first | random | energy | oracle`
+- Report mean energy gap and `P(energy > random)` via shared bootstrap
+
+**Optional continual-forgetting flag** (JEPA only): `--continual` as today.
+
+---
+
+## Benchmark data
+
+Add small files derived from lesson topics in [`research/data/education-lesson-chat.jsonl`](research/data/education-lesson-chat.jsonl):
+
+**`benchmark-qa.jsonl`** (~8–10 rows):
+
+```json
+{"question": "What is photosynthesis?", "answer": "Plants make food using sunlight, water, and carbon dioxide.", "domain": "science"}
+```
+
+**`benchmark-kb.jsonl`** (~8–10 rows):
+
+```json
+{"text": "Photosynthesis: plants use sunlight, water, and CO2 to make glucose and release oxygen."}
+```
+
+Keeps eval aligned with the hackathon domain (education) without needing a large external QA set.
+
+---
+
+## What “good” benchmark output looks like
+
+**JEPA harness** (from existing [`eval_harness.py`](research/eval_harness.py)):
+
+```
+config              EM      F1   lat(s)
+C1_base           0.120   0.180    0.05
+C2_rag            0.240   0.310    0.06
+C3_rag_router     0.260   0.330    0.06
+C4_full_jepa      0.320   0.400    0.45
+
+best-of-N selector comparison:
+  first    EM=0.28
+  random   EM=0.31
+  jepa     EM=0.38
+  oracle   EM=0.52
+  P(jepa > random) = 0.97   JEPA critic WORKS
+```
+
+On **toy + untrained weights**, deltas will be noise — smoke tier only checks the pipeline runs. Meaningful numbers need Tier 2/3 with a real LLM and/or a trained bridge checkpoint.
+
+**World harness** (new): same table shape but selectors `first | random | energy | oracle`.
+
+---
+
+## Implementation notes (minimal scope)
+
+- Preserve existing public classes: `Ensemble`, `WorldEnsemble`, same `answer_ids` / `answer` APIs.
+- `HFBackend` in JEPA track keeps LoRA adapter bank + router; world track keeps simpler `HFLLM` (no router) — unify only where behavior matches.
+- Add `if __name__ == "__main__"` entrypoints on `jepa_ensemble.py` and `world_ensemble.py` (move demos from current files).
+- No integration with Gradio/`libs/agent` in this pass — research-only per your choice.
+- No new pytest suite unless requested; `smoke.sh` + toy harness are the acceptance check.
+
+---
+
+## If time is short (Option A fallback)
+
+Skip deduplication; only:
+
+```
+research/ensemble/jepa/{ensemble.py, eval_harness.py}
+research/ensemble/world_model/{ensemble.py}
+```
+
+Fix imports, add `README.md` with Tier 1–3 commands. Revisit shared modules post-hackathon.

.cursor/plans/lm-eval_harness_integration_34af70c7.plan.md

@@ -0,0 +1,339 @@
+---
+name: lm-eval harness integration
+overview: Add EleutherAI lm-evaluation-harness as the primary academic benchmark track for finetuned HF/LoRA models and ensemble checkpoints, with preset-aware CLI, fair-comparison configs, and structured results — building on existing slm-evals and ensemble harnesses without replacing them.
+todos:
+  - id: deps-lm-eval
+    content: Add lm-eval[hf] optional dependency group to root pyproject.toml and slm-evals extras
+    status: completed
+  - id: preset-resolver
+    content: Implement preset_resolver.py mapping models.yaml presets to lm-eval model_args (base, LoRA, merged)
+    status: completed
+  - id: ensemble-lm-backend
+    content: Implement ensemble-lm custom LM class (generate_until + loglikelihood via underlying HF LLM)
+    status: completed
+  - id: slm-lm-eval-cli
+    content: Build run_lm_eval.py CLI with YAML config, simple_evaluate(), results + summary.md + --compare-to
+    status: completed
+  - id: lm-eval-configs
+    content: Add lm_eval_minicpm5.yaml, lm_eval_smoke.yaml, lm_eval_compare_study.yaml under research/evals/configs/
+    status: completed
+  - id: docs-workflows
+    content: Document finetune + ensemble lm-eval workflows in research/evals/USAGE.md and research/USAGE.md
+    status: completed
+  - id: finetune-hook
+    content: "Optional: add --lm-eval-after to finetune.py after CLI smoke tests pass"
+    status: completed
+  - id: phase2-verify
+    content: "Deferred: slm-compare, education_qa, harness JSON export from verification pipeline plan"
+    status: cancelled
+isProject: false
+---
+
+# lm-evaluation-harness Integration for research/ Models
+
+## Context: what you already have
+
+Your repo runs **three eval tracks** today; none use standard academic benchmarks (GSM8K, ARC, HellaSwag):
+
+| Track | Tool | Best for |
+|-------|------|----------|
+| Agentic | [`slm-benchmark`](research/evals/src/slm_evals/run_benchmark.py) | BFCL, τ-bench, GAIA, SWE |
+| Finetune training | [`finetune.py`](research/finetune.py) | eval_loss / perplexity only |
+| Ensemble components | [`jepa_harness`](research/ensemble/src/ensemble/eval/jepa_harness.py) | RAG/router/JEPA ablation on custom QA |
+
+The existing [Model Verification Pipeline plan](.cursor/plans/model_verification_pipeline_ed9d35ab.plan.md) defers lm-eval; per your choice, **lm-eval comes first**. `slm-compare` / `education_qa` stay Phase 2.
+
+```mermaid
+flowchart TB
+  subgraph models [Model artifacts]
+    basePreset[models.yaml preset]
+    loraDir[LoRA adapter dir]
+    mergedDir[Merged checkpoint]
+    ensCkpt[Ensemble manifest checkpoint]
+  end
+
+  subgraph newTrack [NEW: lm-eval track]
+    resolver[PresetResolver]
+    slmLmEval[slm-lm-eval CLI]
+    hfLm[lm_eval hf backend]
+    ensLm[ensemble-lm custom backend]
+    tasks[Standard tasks]
+    results[results/lm_eval/]
+  end
+
+  subgraph existing [Existing tracks unchanged]
+    agentic[slm-benchmark]
+    ablation[jepa_harness]
+  end
+
+  basePreset --> resolver
+  loraDir --> resolver
+  mergedDir --> resolver
+  resolver --> hfLm
+  ensCkpt --> ensLm
+  slmLmEval --> hfLm
+  slmLmEval --> ensLm
+  hfLm --> tasks
+  ensLm --> tasks
+  tasks --> results
+```
+
+---
+
+## Benchmark matrix (what to run for each claim)
+
+Match tasks to model size (~1B) and claim type:
+
+| Model | lm-eval tasks (primary) | Existing complement |
+|-------|-------------------------|---------------------|
+| **Base preset** (`minicpm5-1b`, `gemma4-e2b-mobile`) | `arc_easy`, `arc_challenge`, `hellaswag`, `piqa`, `boolq`, `gsm8k` | `slm-benchmark --benchmarks bfcl,tau_bench` |
+| **Finetuned LoRA** (`minicpm5-1b-lesson-lora`) | Same tasks; `pretrained=base,peft=adapter` | Same agentic suite |
+| **Merged finetune** (`minicpm5-1b-lesson-merged`, `gemma-merged-local`) | Same tasks; `pretrained=merged_path` | Same |
+| **Ensemble checkpoint** | Same generative + MC tasks via custom backend | `jepa_harness` for component ablation; `slm-benchmark --model-type ensemble` for agentic E2E |
+
+**Smoke profile** (CPU/GPU quick check, ~5–15 min): `--tasks arc_easy,hellaswag --num_fewshot 0 --limit 50`
+
+**Full profile** (reportable): `--tasks arc_easy,arc_challenge,hellaswag,piqa,boolq,gsm8k --num_fewshot 5` (gsm8k uses 8-shot per harness default)
+
+**Fair comparison rules** (from your guide, enforced in YAML):
+- Identical `tasks`, `num_fewshot`, `limit`, `seed`, `batch_size`
+- Same base tokenizer (preset resolution guarantees this for LoRA)
+- `temperature=0` / greedy decoding (lm-eval default for MC; gsm8k generative)
+- Never compare `training_results.json` `result_score` to lm-eval accuracy
+
+---
+
+## Phase 1 — Dependencies and install surface
+
+**Root [`pyproject.toml`](pyproject.toml)** — add optional group:
+
+```toml
+lm-eval = [
+    "lm-eval[hf]>=0.4.9",
+]
+```
+
+Install: `uv sync --group evals --group lm-eval --group finetune`
+
+Keep lm-eval **optional** so agentic-only workflows stay lightweight.
+
+---
+
+## Phase 2 — Preset → lm-eval model_args resolver
+
+New module: `research/evals/src/slm_evals/lm_eval/preset_resolver.py`
+
+Reuse [`inference.config.get_model_config`](libs/inference/src/inference/config.py) (same as [`finetune.py`](research/finetune.py)):
+
+| Preset shape | lm-eval invocation |
+|--------------|-------------------|
+| `model_id` only (base) | `pretrained={model_id},trust_remote_code=True` |
+| `model_id` + `adapter_path` (LoRA) | `pretrained={model_id},peft={adapter_path},trust_remote_code=True` |
+| Local merged dir | `pretrained={model_id},trust_remote_code=True` |
+| Ensemble (`jepa-ensemble-lesson` or path with `manifest.json`) | `--model ensemble-lm` (custom backend) |
+
+Reject multimodal / llama_cpp presets with a clear error (same rule as finetune).
+
+---
+
+## Phase 3 — CLI wrapper `slm-lm-eval`
+
+New entry: `research/evals/src/slm_evals/run_lm_eval.py`  
+Register in [`research/evals/pyproject.toml`](research/evals/pyproject.toml):
+
+```toml
+[project.scripts]
+slm-lm-eval = "slm_evals.run_lm_eval:main"
+```
+
+**Flags:**
+
+```bash
+uv run --package slm-evals slm-lm-eval \
+  --config research/evals/configs/lm_eval_minicpm5.yaml \
+  --preset minicpm5-1b \
+  --experiment-name minicpm5-1b__lm-eval-baseline
+```
+
+| Flag | Purpose |
+|------|---------|
+| `--preset` | Resolve from [`models.yaml`](models.yaml) |
+| `--model` | Override path/Hub id (merged dir or ensemble ckpt) |
+| `--adapter` | Override LoRA path (alternative to preset) |
+| `--config` | YAML: tasks, num_fewshot, limit, seed, device, batch_size |
+| `--tasks` | CLI override of task list |
+| `--compare-to` | Path to prior `results.json` → print delta table (lightweight pre-`slm-compare`) |
+| `--experiment-name` / `--output-dir` | Write under `results/lm_eval/{name}/` |
+
+Implementation: subprocess or programmatic call to `lm_eval.simple_evaluate()` (preferred over shelling to `lm_eval` CLI — easier to inject custom model registration).
+
+Output artifacts per run:
+- `results.json` (lm-eval native)
+- `summary.md` (human-readable table: task → acc/score)
+- `run_meta.json` (preset, base_model, adapter_path, tasks, seed, git hash optional)
+
+---
+
+## Phase 4 — Ensemble custom backend
+
+New file: `research/evals/src/slm_evals/lm_eval/ensemble_lm.py`
+
+Register with lm-eval:
+
+```python
+@register_model("ensemble-lm")
+class EnsembleLM(LM):
+    ...
+```
+
+Load via existing [`load_ensemble_model`](research/evals/src/slm_evals/utils/model_loader.py) / `ensemble.checkpoint.load_checkpoint`.
+
+| lm-eval method | Implementation |
+|----------------|----------------|
+| `generate_until` | `ens.generate_text(prompt, max_new_tokens=..., temperature=0)` — full JEPA+RAG+router stack |
+| `loglikelihood` | Delegate to underlying `ens.llm` HF model (default adapter index 0) for MC tasks; document that this evaluates **base LLM head**, not selector — pair with `generate_until` scores for full-stack generative tasks |
+
+This split is intentional: MC benchmarks need token logprobs; the ensemble’s value on generative QA shows up in `generate_until` (gsm8k) and in [`jepa_harness`](research/ensemble/src/ensemble/eval/jepa_harness.py).
+
+Import side-effect: ensure `ensemble_lm` is imported before `simple_evaluate()` so registration runs.
+
+---
+
+## Phase 5 — Experiment configs
+
+Add under `research/evals/configs/`:
+
+**`lm_eval_minicpm5.yaml`** — baseline template:
+
+```yaml
+tasks:
+  - arc_easy
+  - arc_challenge
+  - hellaswag
+  - piqa
+  - boolq
+  - gsm8k
+num_fewshot: 5        # gsm8k harness may override to 8 internally
+limit: null           # null = full; 100 for dev
+seed: 42
+batch_size: auto
+device: auto
+dtype: bfloat16
+trust_remote_code: true
+```
+
+**`lm_eval_smoke.yaml`** — `limit: 25`, tasks `[arc_easy, hellaswag]`
+
+**`lm_eval_compare_study.yaml`** — documents baseline + candidate preset names and shared settings
+
+---
+
+## Phase 6 — End-to-end workflows (finetune + ensemble)
+
+### Finetuned model verification
+
+```bash
+# 1. Baseline (same config, pinned seed)
+uv run --package slm-evals slm-lm-eval \
+  --config research/evals/configs/lm_eval_minicpm5.yaml \
+  --preset minicpm5-1b \
+  --experiment-name minicpm5-1b__baseline
+
+# 2. Train
+uv run python research/finetune.py --preset minicpm5-1b --mode lora --epochs 3
+
+# 3. Candidate — LoRA via preset (no merge required)
+uv run --package slm-evals slm-lm-eval \
+  --config research/evals/configs/lm_eval_minicpm5.yaml \
+  --preset minicpm5-1b-lesson-lora \
+  --experiment-name minicpm5-1b-lora__v1 \
+  --compare-to results/lm_eval/minicpm5-1b__baseline/results.json
+```
+
+For Gemma: use `gemma-lora-local` / `gemma-merged-local` presets after notebook or `finetune.py` training.
+
+### Ensemble verification
+
+```bash
+# Component ablation (domain QA) — existing
+uv run --package ensemble python -m ensemble.eval.jepa_harness \
+  --llm openbmb/MiniCPM5-1B \
+  --qa research/data/benchmark-qa.jsonl \
+  --kb research/data/benchmark-kb.jsonl
+
+# Academic benchmarks on saved ensemble
+uv run --package slm-evals slm-lm-eval \
+  --config research/evals/configs/lm_eval_minicpm5.yaml \
+  --model ./models/ensemble/jepa-lesson-pretrain \
+  --experiment-name ensemble-jepa__lm-eval
+```
+
+### Optional post-finetune hook (minimal, Phase 1.5)
+
+Add to [`finetune.py`](research/finetune.py) only after CLI stabilizes:
+
+- `--lm-eval-after` + `--lm-eval-config` → subprocess `slm-lm-eval` on output checkpoint
+- Append `lm_eval_summary` path to `training_results.json`
+
+Defer `--eval-baseline` auto-compare until Phase 2 `slm-compare` exists; use `--compare-to` on lm-eval outputs in the meantime.
+
+---
+
+## Phase 7 — Documentation
+
+Update:
+- [`research/evals/USAGE.md`](research/evals/USAGE.md) — lm-eval section, PEFT notes, task profiles
+- [`research/USAGE.md`](research/USAGE.md) — unified “verify finetune / ensemble” checklist
+- [`.env.example`](.env.example) — optional `LM_EVAL_TASKS`, `LM_EVAL_SEED` (low priority)
+
+Include the verification checklist from your guide (seeds, fair comparison, no val-set leakage) mapped to concrete commands.
+
+---
+
+## File change summary
+
+| File | Change |
+|------|--------|
+| [`pyproject.toml`](pyproject.toml) | `lm-eval` dependency group |
+| [`research/evals/pyproject.toml`](research/evals/pyproject.toml) | Optional `lm-eval[hf]` extra; `slm-lm-eval` script |
+| `research/evals/src/slm_evals/lm_eval/__init__.py` | Package init |
+| `research/evals/src/slm_evals/lm_eval/preset_resolver.py` | Preset → model_args |
+| `research/evals/src/slm_evals/lm_eval/ensemble_lm.py` | Custom `ensemble-lm` backend |
+| `research/evals/src/slm_evals/run_lm_eval.py` | Main CLI |
+| `research/evals/configs/lm_eval_*.yaml` | Baseline, smoke, compare templates |
+| [`research/finetune.py`](research/finetune.py) | Optional `--lm-eval-after` (after CLI stable) |
+| [`research/evals/USAGE.md`](research/evals/USAGE.md), [`research/USAGE.md`](research/USAGE.md) | Workflows |
+
+---
+
+## Testing plan
+
+1. **Import smoke**: `uv run --package slm-evals python -c "import lm_eval; import slm_evals.lm_eval.ensemble_lm"`
+2. **HF base**: `slm-lm-eval --config lm_eval_smoke.yaml --preset minicpm5-1b --limit 10`
+3. **LoRA**: same with `--preset minicpm5-1b-lesson-lora` (or `--adapter ./models/finetuned/...`)
+4. **Ensemble**: `slm-lm-eval --model ./models/ensemble/... --config lm_eval_smoke.yaml`
+5. **Compare**: baseline + candidate runs produce `--compare-to` delta without task/seed mismatch warnings
+6. **Regression**: existing `slm-benchmark` and `jepa_harness --toy` still pass unchanged
+
+---
+
+## Phase 2 (deferred — existing verification plan)
+
+After lm-eval stabilizes, implement from [model_verification_pipeline plan](.cursor/plans/model_verification_pipeline_ed9d35ab.plan.md):
+
+- `slm-compare` with paired bootstrap across **both** lm-eval and slm-benchmark JSON
+- `education_qa` domain benchmark in slm-evals
+- PEFT loading in `slm-benchmark` model_loader (reuse `preset_resolver`)
+- Harness JSON export (`jepa_harness --output-dir`)
+
+This gives you: **academic generalization** (lm-eval) + **agentic capability** (slm-evals) + **ensemble component proof** (jepa harness) with shared statistical comparison.
+
+---
+
+## Expected limitations (document, don’t hide)
+
+- **1B models** will score low on gsm8k/mmlu — use for **relative** baseline vs finetune deltas, not SOTA claims
+- **Ensemble loglikelihood** uses underlying LLM, not full selector stack — report both modes in docs
+- **MiniCPM / Gemma** need `trust_remote_code=True`
+- **First lm-eval run** downloads datasets; pin `HF_HOME` / cache for reproducibility
+- **Multi-seed training** (3–5 seeds) is manual until Phase 2 compare aggregates runs

.cursor/plans/model_verification_pipeline_ed9d35ab.plan.md

@@ -0,0 +1,280 @@
+---
+name: Model Verification Pipeline
+overview: Operationalize model verification for both fine-tuned SLMs and ensemble research by adding comparison tooling, preset-aware eval, a lesson-domain benchmark, reproducible experiment configs, and optional post-finetune eval — all grounded in your existing `slm-evals` and ensemble harness infrastructure.
+todos:
+  - id: eval-infra
+    content: Add preset resolution, PEFT loading, eval_seed sample pinning, and education_qa benchmark to slm-evals
+    status: pending
+  - id: compare-cli
+    content: Implement slm-compare with paired bootstrap, delta tables, and subset mismatch warnings
+    status: pending
+  - id: finetune-hook
+    content: Add --eval-after / --eval-baseline flags to finetune.py with subprocess eval + compare
+    status: pending
+  - id: ensemble-json
+    content: Persist jepa/world harness results to JSON (--output-dir) using shared stats helpers
+    status: pending
+  - id: docs-configs
+    content: Add baseline_minicpm5.yaml template and verification checklist to research/ USAGE docs
+    status: pending
+isProject: false
+---
+
+# Model Verification Pipeline for research/
+
+## Current state
+
+Your repo already has **three parallel eval tracks** with no unified comparison:
+
+| Track | Tool | Metrics | Results | Stats |
+|-------|------|---------|---------|-------|
+| Fine-tuned SLM | [`slm-benchmark`](research/evals/src/slm_evals/run_benchmark.py) | BFCL, τ-bench, GAIA, SWE pass rate | `results/<experiment>/results.json` | None |
+| Finetune training | [`finetune.py`](research/finetune.py) | eval_loss, perplexity, `result_score` | `training_results.json` | None |
+| Ensemble | [`jepa_harness`](research/ensemble/src/ensemble/eval/jepa_harness.py) / [`world_harness`](research/ensemble/src/ensemble/eval/world_harness.py) | EM, F1, ablation ladder | stdout only | [`paired_bootstrap`](research/ensemble/src/ensemble/eval/metrics.py) |
+
+The manual before/after loop in [`research/evals/USAGE.md`](research/evals/USAGE.md) works but lacks significance testing, preset resolution, and a domain-aligned benchmark for lesson fine-tuning.
+
+```mermaid
+flowchart TB
+  subgraph today [Today]
+    baseline[slm-benchmark baseline]
+    finetune[finetune.py]
+    reeval[slm-benchmark new model]
+    manual[Manual diff results.csv]
+    baseline --> manual
+    reeval --> manual
+    finetune --> reeval
+  end
+
+  subgraph target [Target pipeline]
+    preset[models.yaml preset]
+    evalRun[slm-benchmark + eval_seed]
+    compare[slm-compare]
+    report[comparison_report.md]
+    preset --> evalRun
+    evalRun --> compare
+    compare --> report
+  end
+```
+
+---
+
+## Verification strategy (mapped to your repo)
+
+### 1. Fair comparison checklist (enforce via shared eval config)
+
+Create a **single YAML experiment config** per comparison study (copy from [`experiment_001.yaml`](research/evals/configs/experiment_001.yaml)) and reuse it for baseline + candidate runs:
+
+- Same `benchmarks`, `max_samples`, `benchmark_overrides`, `temperature: 0.0`, `max_new_tokens`
+- Same `eval_seed` (new field) so sample subsets are identical across models
+- Baseline = base preset from [`models.yaml`](models.yaml) (e.g. `minicpm5-1b`); candidate = LoRA/merged preset (e.g. `minicpm5-1b-lesson-lora`)
+
+**Do not** compare `training_results.json` `result_score` against BFCL pass rate — they measure different things.
+
+### 2. Benchmark selection for your use cases
+
+| Claim | Benchmarks in this repo |
+|-------|-------------------------|
+| Agentic tool use (general) | `bfcl`, `tau_bench` (already in slm-evals) |
+| End-to-end assistant | `gaia` (levels 1–2 for small models) |
+| Code | `swe_bench` (keep `full_eval: false` unless Docker installed) |
+| Lesson/education domain (finetune target) | **New** `education_qa` benchmark on [`research/data/benchmark-qa.jsonl`](research/data/benchmark-qa.jsonl) |
+| Ensemble component value | Existing C1→C4 ablation ladder in JEPA harness |
+| JEPA critic beats random | Existing selector comparison + `paired_bootstrap` |
+
+Run **in-distribution** (`education_qa`) + **out-of-distribution** (`bfcl`, `tau_bench`) to show generalization, not just lesson memorization.
+
+### 3. Statistical validation
+
+Reuse existing `paired_bootstrap` from [`metrics.py`](research/ensemble/src/ensemble/eval/metrics.py) in a new shared module. For each benchmark:
+
+- Align per-sample `passed` (or `score`) arrays by `samples[].id`
+- Report: Δscore, win rate, `P(candidate > baseline)` from bootstrap, 95% CI via bootstrap percentiles
+- Flag `p > 0.95` as "significant improvement" (same threshold as JEPA harness)
+
+For multi-seed training runs: run eval once per checkpoint seed, aggregate mean ± std in comparison report (optional `--runs` glob in compare CLI).
+
+---
+
+## Implementation plan
+
+### Phase A — Shared eval infrastructure (`research/evals/`)
+
+**A1. Preset resolution in config loader**
+
+Extend [`config_loader.py`](research/evals/src/slm_evals/utils/config_loader.py) and CLI in [`run_benchmark.py`](research/evals/src/slm_evals/run_benchmark.py):
+
+- Add `--preset` flag and YAML field `preset:` (alternative to `model_path`)
+- Resolve via existing [`inference.config.get_model_config`](libs/inference/src/inference/config.py): extract `model_id` + optional `adapter_path`
+- Reject non-`transformers` / multimodal presets with clear error (same rule as finetune)
+
+**A2. PEFT adapter loading**
+
+Extend [`model_loader.py`](research/evals/src/slm_evals/utils/model_loader.py):
+
+- If `adapter_path` is set: load base from `model_id`, attach LoRA via `peft.PeftModel.from_pretrained`
+- Support merged checkpoints (adapter_path absent) unchanged
+- Record `base_model`, `adapter_path`, `param_count` in results metadata
+
+**A3. Reproducible sample subsets**
+
+Extend [`BaseBenchmark`](research/evals/src/slm_evals/benchmarks/base.py):
+
+- Accept `eval_seed` + `max_samples` from config
+- After loading dataset: `rng = random.Random(eval_seed); indices = rng.sample(range(len(data)), min(max_samples, len(data)))`
+- Persist `eval_seed`, `sample_ids` list in `results.json` so compare can verify identical subsets
+
+**A4. Lesson-domain benchmark**
+
+New file `research/evals/src/slm_evals/benchmarks/education_qa.py`:
+
+- Load [`research/data/benchmark-qa.jsonl`](research/data/benchmark-qa.jsonl)
+- Prompt: `"Answer briefly.\nQ: {question}\nA:"`
+- Score: token-overlap F1 + normalized substring EM (reuse logic from ensemble `metrics.py` — extract to shared `slm_evals/utils/scoring.py` or import from ensemble if dependency is acceptable)
+- Register as `education_qa` in `BENCHMARK_REGISTRY`
+
+**A5. Comparison CLI — `slm-compare`**
+
+New module `research/evals/src/slm_evals/compare_runs.py` + console script in [`pyproject.toml`](research/evals/pyproject.toml):
+
+```bash
+uv run --package slm-evals slm-compare \
+  --baseline results/minicpm5-1b__baseline/results.json \
+  --candidate results/minicpm5-1b-lora__v1/results.json \
+  --output results/comparisons/minicpm5-lora-vs-base.md
+```
+
+Outputs:
+- Per-benchmark delta table (score, passed/total, latency)
+- Paired bootstrap p-value per benchmark
+- Per-sample win/loss/tie counts (joined on sample `id`)
+- Warnings if `eval_seed`, `max_samples`, or benchmark sets differ
+
+**A6. Experiment config templates**
+
+Add two configs under `research/evals/configs/`:
+
+- `baseline_minicpm5.yaml` — preset `minicpm5-1b`, benchmarks `[education_qa, bfcl, tau_bench]`, `max_samples: 100`, `eval_seed: 42`
+- `compare_study.yaml` — documents baseline + candidate preset keys, shared eval settings, output naming convention
+
+Update [`research/evals/USAGE.md`](research/evals/USAGE.md) and [`research/USAGE.md`](research/USAGE.md) with the verification checklist from your research guide.
+
+---
+
+### Phase B — Finetune integration (`research/finetune.py`)
+
+**B1. Post-finetune eval hook**
+
+Add optional flags to [`finetune.py`](research/finetune.py):
+
+- `--eval-after` — run slm-benchmark after training completes
+- `--eval-config PATH` — YAML with benchmark settings (defaults to `baseline_minicpm5.yaml` structure)
+- `--eval-baseline PRESET` — also eval base preset for side-by-side comparison
+
+On completion:
+1. Write `training_results.json` (existing)
+2. Run eval on output checkpoint
+3. If `--eval-baseline` set, run baseline eval + invoke `slm-compare`
+4. Append comparison summary path to `training_results.json` under `"post_eval": {...}`
+
+Implementation: subprocess call to `uv run --package slm-evals slm-benchmark` (avoids circular imports).
+
+---
+
+### Phase C — Ensemble track unification (`research/ensemble/`)
+
+**C1. Persist harness results to JSON**
+
+Extend [`jepa_harness.py`](research/ensemble/src/ensemble/eval/jepa_harness.py) and [`world_harness.py`](research/ensemble/src/ensemble/eval/world_harness.py):
+
+- Add `--output-dir` flag
+- Write `results.json` matching slm-evals schema where possible:
+  - `benchmarks.C4_full_jepa.samples[]` with per-question EM/F1
+  - `benchmarks.C1_base` … `C4_full_jepa` aggregate scores
+  - `significance[]` with paired bootstrap between ladder steps
+
+**C2. Compare ensemble configs**
+
+`slm-compare` accepts ensemble JSON (detect by presence of ablation config keys) and renders ablation table + bootstrap block.
+
+**C3. Shared metrics module**
+
+Move or re-export `em_score`, `f1_score`, `paired_bootstrap` to `research/evals/src/slm_evals/utils/stats.py` (ensemble imports from there, or duplicate minimally to avoid cross-package coupling — prefer a tiny shared module under `research/` if both packages need it).
+
+---
+
+### Phase D — End-to-end workflow documentation
+
+Add a **Verification Checklist** section to [`research/README.md`](research/README.md):
+
+```bash
+# 1. Baseline (preset-aware, pinned seed)
+uv run --package slm-evals slm-benchmark \
+  --config research/evals/configs/baseline_minicpm5.yaml \
+  --preset minicpm5-1b \
+  --experiment-name minicpm5-1b__baseline
+
+# 2. Fine-tune
+uv run python research/finetune.py --preset minicpm5-1b --mode lora --epochs 3
+
+# 3. Candidate eval (same config, different preset/path)
+uv run --package slm-evals slm-benchmark \
+  --config research/evals/configs/baseline_minicpm5.yaml \
+  --preset minicpm5-1b-lesson-lora \
+  --experiment-name minicpm5-1b-lora__v1
+
+# 4. Statistical comparison
+uv run --package slm-evals slm-compare \
+  --baseline results/minicpm5-1b__baseline/results.json \
+  --candidate results/minicpm5-1b-lora__v1/results.json
+
+# 5. Ensemble ablation (domain QA)
+uv run --package ensemble python -m ensemble.eval.jepa_harness \
+  --llm openbmb/MiniCPM5-1B \
+  --qa research/data/benchmark-qa.jsonl \
+  --kb research/data/benchmark-kb.jsonl \
+  --output-dir results/ensemble/jepa-ablation-v1
+```
+
+**Claiming "better than baseline" requires:**
+- Δscore > 0 on target benchmarks
+- `P(candidate > baseline) > 0.95` on paired per-sample scores (or document inconclusive)
+- No regression > 2pp on OOD benchmarks unless tradeoff is explicitly claimed
+- Ablation shows which components contribute (ensemble) or which eval split improved (finetune)
+
+---
+
+## File change summary
+
+| File | Change |
+|------|--------|
+| [`research/evals/src/slm_evals/utils/config_loader.py`](research/evals/src/slm_evals/utils/config_loader.py) | `preset`, `eval_seed`, adapter fields |
+| [`research/evals/src/slm_evals/utils/model_loader.py`](research/evals/src/slm_evals/utils/model_loader.py) | PEFT adapter loading |
+| [`research/evals/src/slm_evals/benchmarks/base.py`](research/evals/src/slm_evals/benchmarks/base.py) | Seeded sample selection |
+| `research/evals/src/slm_evals/benchmarks/education_qa.py` | New lesson-domain benchmark |
+| `research/evals/src/slm_evals/compare_runs.py` | New comparison CLI |
+| `research/evals/src/slm_evals/utils/stats.py` | Bootstrap + scoring helpers |
+| [`research/evals/pyproject.toml`](research/evals/pyproject.toml) | Register `slm-compare` script |
+| [`research/finetune.py`](research/finetune.py) | `--eval-after`, `--eval-config`, `--eval-baseline` |
+| [`research/ensemble/src/ensemble/eval/jepa_harness.py`](research/ensemble/src/ensemble/eval/jepa_harness.py) | `--output-dir`, JSON persistence |
+| [`research/ensemble/src/ensemble/eval/world_harness.py`](research/ensemble/src/ensemble/eval/world_harness.py) | Same |
+| [`research/evals/configs/baseline_minicpm5.yaml`](research/evals/configs/baseline_minicpm5.yaml) | New template |
+| [`research/USAGE.md`](research/USAGE.md), [`research/evals/USAGE.md`](research/evals/USAGE.md) | Verification workflow docs |
+
+---
+
+## Testing plan
+
+- **Unit**: `paired_bootstrap` on synthetic paired arrays; sample-id alignment in compare; preset resolution for `minicpm5-1b` and `minicpm5-1b-lesson-lora`
+- **Smoke**: `slm-benchmark --preset minicpm5-1b --benchmarks education_qa --max-samples 5 --device cpu`
+- **Integration**: baseline → finetune (1 epoch) → re-eval → `slm-compare` produces report with no subset mismatch warnings
+- **Ensemble**: `jepa_harness --toy --output-dir /tmp/jepa-test` writes valid JSON readable by `slm-compare`
+
+---
+
+## Out of scope (defer)
+
+- W&B / MLflow experiment tracking
+- lm-evaluation-harness integration (GSM8K/MMLU — different from agentic suite; add later as separate benchmark group if needed)
+- Running ensemble inference through BFCL (requires adapter wrapper, not just raw HF model)
+- CI regression gates (can add after compare CLI stabilizes)

.env.example

@@ -28,3 +28,28 @@ ALLOW_MODEL_SWITCH=false
 # Optional: local fine-tuned merged weights
 # ACTIVE_MODEL=gemma-merged-local
 # MODEL_ID=./gemma_merged_model
+
+# --- Fine-tuning (research/finetune.py) ---
+# FINETUNE_PRESET=minicpm5-1b
+# FINETUNE_MODEL=openbmb/MiniCPM5-1B
+# FINETUNE_DATASET=./research/data/education-lesson-chat.jsonl
+# FINETUNE_DATASET=tatsu-lab/alpaca
+# FINETUNE_DATASET_CONFIG=
+# FINETUNE_DATASET_SPLIT=train
+# FINETUNE_MAX_SAMPLES=500
+# FINETUNE_OUT=./models/finetuned/minicpm5-1b-lora
+# FINETUNE_FORMAT=chat
+# After training, point Gradio at the adapter preset:
+# ACTIVE_MODEL=minicpm5-1b-lesson-lora
+
+# --- 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

@@ -10,4 +10,6 @@ models/
 dist/
 build/
 
-outputs/traces
+outputs/traces
+
+/results

README.md

@@ -46,8 +46,11 @@ apps/gradio-space/   # Gradio tabs (Lesson slides + Chat debug)
 libs/agent/          # Skill agent runner, tools, trace recorder
 libs/inference/      # Transformers + llama.cpp backends
 skills/              # SKILL.md task definitions
+research/            # Fine-tune, ensemble experiments, agentic evals (optional)
 ```
 
+Optional research tooling (not required for the Space): see [research/USAGE.md](research/USAGE.md).
+
 ## Environment variables
 
 | Variable | Default | Description |

libs/inference/pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
     "pyyaml>=6.0.2",
     "torch>=2.5.0",
     "torchvision>=0.20.0",
+    "peft>=0.14.0",
     "transformers>=5.7.0",
 ]
 

libs/inference/src/inference/config.py

@@ -23,6 +23,7 @@ class ModelConfig:
     model_file: str | None = None
     model_path: str | None = None
     model_id: str | None = None
+    adapter_path: str | None = None
     trust_remote_code: bool = False
     multimodal: bool = False
     n_ctx: int = 4096
@@ -37,6 +38,7 @@ class ModelConfig:
             self.model_file,
             self.model_path,
             self.model_id,
+            self.adapter_path,
             self.trust_remote_code,
             self.multimodal,
             self.n_ctx,
@@ -60,6 +62,9 @@ class ModelConfig:
         if self.model_id and self.model_id.startswith(("./", "../")):
             updates["model_id"] = str((base_dir / self.model_id).resolve())
 
+        if self.adapter_path and self.adapter_path.startswith(("./", "../")):
+            updates["adapter_path"] = str((base_dir / self.adapter_path).resolve())
+
         return replace(self, **updates) if updates else self
 
 
@@ -149,6 +154,7 @@ def _parse_model_entry(key: str, raw: dict[str, Any]) -> ModelConfig:
         model_file=raw.get("model_file"),
         model_path=raw.get("model_path"),
         model_id=raw.get("model_id"),
+        adapter_path=raw.get("adapter_path"),
         trust_remote_code=bool(raw.get("trust_remote_code", False)),
         multimodal=bool(raw.get("multimodal", False)),
         n_ctx=int(raw.get("n_ctx", 4096)),

libs/inference/src/inference/transformers.py

@@ -76,6 +76,18 @@ class TransformersBackend:
                 **model_kwargs,
             )
 
+        if self._config.adapter_path:
+            from pathlib import Path
+
+            from peft import PeftModel
+
+            adapter = Path(self._config.adapter_path)
+            if not adapter.is_dir():
+                raise FileNotFoundError(
+                    f"LoRA adapter not found for preset {self._config.key!r}: {adapter}"
+                )
+            self._model = PeftModel.from_pretrained(self._model, str(adapter))
+
         if device == "cpu":
             self._model.to(device)
 

models.yaml

@@ -2,9 +2,8 @@
 # Select active preset with ACTIVE_MODEL; override any field via .env (see .env.example).
 
 defaults:
-  active_model: minicpm-v-4.6
-
-  # active_model: minicpm5-1b
+  # active_model: minicpm-v-4.6
+  active_model: minicpm5-1b
 
   # Dev: set ALLOW_MODEL_SWITCH=true in .env to expose a dropdown in Gradio.
   # Space: keep false so visitors use one pinned model.
@@ -55,3 +54,22 @@ models:
     label: Fine-tuned LoRA adapter (local path)
     backend: transformers
     model_id: ./gemma_finetuned_model
+
+  minicpm5-1b-lesson-lora:
+    label: MiniCPM5 1B lesson LoRA (local)
+    backend: transformers
+    model_id: openbmb/MiniCPM5-1B
+    adapter_path: ./models/finetuned/minicpm5-1b-lora
+    trust_remote_code: true
+
+  minicpm5-1b-lesson-merged:
+    label: MiniCPM5 1B lesson merged (local)
+    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

pyproject.toml

@@ -6,6 +6,7 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "agent",
+    "ensemble",
     "gradio-space",
     "inference",
 ]
@@ -15,14 +16,34 @@ dev = [
     "ruff>=0.9.0",
     "pytest>=8.0.0",
 ]
+finetune = [
+    "bitsandbytes>=0.46.1",
+    "datasets>=3.0.0",
+    "peft>=0.14.0",
+]
+ensemble = [
+    "accelerate>=1.2.0",
+    "peft>=0.14.0",
+    "transformers>=5.7.0",
+]
+evals = [
+    "slm-evals",
+]
+lm-eval = [
+    "slm-evals[lm-eval]",
+]
 
 [tool.uv.workspace]
 members = [
     "apps/*",
     "libs/*",
+    "research/ensemble",
+    "research/evals",
 ]
 
 [tool.uv.sources]
 agent = { workspace = true }
+ensemble = { workspace = true }
 gradio-space = { workspace = true }
 inference = { workspace = true }
+slm-evals = { workspace = true }

research/README.md

@@ -0,0 +1,48 @@
+# 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.
+
+| 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 |
+
+## 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
+```
+
+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`) |
+
+## Typical workflow
+
+```text
+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)
+```
+
+See [USAGE.md](USAGE.md) for copy-paste commands.

research/USAGE.md

@@ -0,0 +1,311 @@
+# 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.
+
+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)
+- 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
+
+# Or one at a time
+uv sync --group finetune
+uv sync --group ensemble
+uv sync --group evals
+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, …) |
+
+---
+
+## 1. Fine-tuning (`research/finetune.py`)
+
+Single script for **full**, **LoRA**, and **QLoRA** training. Defaults to the lesson-agent chat dataset at `research/data/education-lesson-chat.jsonl` and writes checkpoints under `models/finetuned/`.
+
+### Model resolution (first match wins)
+
+1. `--model <hf-id-or-path>`
+2. `--preset <key>` from root `models.yaml`
+3. Env: `FINETUNE_MODEL`, `MODEL_ID`, or `BASE`
+4. `ACTIVE_MODEL` preset from `.env`
+
+### Quick start
+
+```bash
+# LoRA on default lesson chat data, 1 epoch
+uv run python research/finetune.py --preset minicpm5-1b --mode lora --epochs 1
+
+# Smoke run (50 steps)
+uv run python research/finetune.py --mode lora --max_steps 50
+
+# QLoRA on a Hub instruction dataset
+uv run python research/finetune.py \
+  --model Qwen/Qwen2.5-0.5B-Instruct \
+  --dataset tatsu-lab/alpaca --format alpaca \
+  --mode qlora --epochs 1
+
+# Merge LoRA adapter into standalone weights
+uv run python research/finetune.py \
+  --merge ./models/finetuned/minicpm5-1b-lora \
+  --out ./models/finetuned/minicpm5-1b-merged
+```
+
+### Dataset formats (`--format`)
+
+| Format | Expected columns |
+| ------ | ---------------- |
+| `chat` | `messages`: `[{"role": "...", "content": "..."}]` |
+| `alpaca` | `instruction`, optional `input`, `output` |
+| `prompt` | `prompt` / `completion` (or `response`) |
+| `text` | `text`, or a plain `.txt` file |
+
+Local files: `.json`, `.jsonl`, `.csv`, `.txt`. Hub ids: any `datasets` repo id.
+
+### Outputs
+
+Training writes to `<out>/` (default `./models/finetuned/<preset>-<mode>/`):
+
+- Adapter or full weights
+- `training_results.json` — train/eval loss, perplexity, `result_score` (0–100)
+
+### Env vars
+
+| Variable | Description |
+| -------- | ----------- |
+| `FINETUNE_PRESET` | Preset key from `models.yaml` |
+| `FINETUNE_DATASET` | Override dataset path or Hub id |
+| `FINETUNE_DATASET_CONFIG` | Hub config name |
+| `FINETUNE_DATASET_SPLIT` | Hub split (e.g. `train[:500]`) |
+| `ACTIVE_MODEL` | Fallback preset when `--preset` omitted |
+
+---
+
+## 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/`)
+
+Evaluate local HuggingFace checkpoints on BFCL, τ-bench, GAIA, and SWE-bench Verified.
+
+Install: `uv sync --group evals`
+
+```bash
+# Smoke test (20 samples, two benchmarks)
+uv run --package slm-evals slm-benchmark \
+  --model openbmb/MiniCPM5-1B \
+  --benchmarks bfcl tau_bench \
+  --max-samples 20
+
+# Full config-driven run
+uv run --package slm-evals slm-benchmark \
+  --config research/evals/configs/experiment_001.yaml
+```
+
+Full reference: [evals/USAGE.md](evals/USAGE.md).
+
+---
+
+## 4. Academic benchmarks (`slm-lm-eval`)
+
+Standard lm-evaluation-harness tasks (ARC, HellaSwag, GSM8K, …) for base presets, LoRA adapters, merged checkpoints, and ensemble manifests.
+
+Install: `uv sync --group lm-eval`
+
+Profile guide: [evals/docs/eval_profiles.md](evals/docs/eval_profiles.md)
+
+```bash
+# List claim-matched profiles (reasoning, code, understanding, …)
+uv run --package slm-evals slm-lm-eval --list-profiles
+
+# Run by profile name
+uv run --package slm-evals slm-lm-eval \
+  --profile reasoning \
+  --preset minicpm5-1b \
+  --experiment-name minicpm5-1b__reasoning-baseline
+
+# Smoke (25 samples, arc_easy + hellaswag)
+uv run --package slm-evals slm-lm-eval \
+  --profile smoke \
+  --preset minicpm5-1b \
+  --experiment-name minicpm5-1b__smoke
+
+# Full profile
+uv run --package slm-evals slm-lm-eval \
+  --config research/evals/configs/lm_eval_minicpm5.yaml \
+  --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:
+
+```bash
+uv run python research/finetune.py \
+  --preset minicpm5-1b --mode lora --max_steps 50 \
+  --lm-eval-after \
+  --lm-eval-baseline minicpm5-1b
+```
+
+Full reference: [evals/USAGE.md](evals/USAGE.md#lm-evaluation-harness-slm-lm-eval).
+
+---
+
+## Shared data (`research/data/`)
+
+| 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 |
+
+---
+
+## Suggested end-to-end pipeline
+
+1. **Baseline lm-eval** — academic benchmarks on the base preset (pinned seed):
+   ```bash
+   uv run --package slm-evals slm-lm-eval \
+     --config research/evals/configs/lm_eval_compare_study.yaml \
+     --preset minicpm5-1b \
+     --experiment-name minicpm5-1b__baseline
+   ```
+
+2. **Baseline agentic eval** (optional):
+   ```bash
+   uv run --package slm-evals slm-benchmark \
+     --model openbmb/MiniCPM5-1B --benchmarks bfcl --max-samples 50
+   ```
+
+3. **Fine-tune** on lesson data:
+   ```bash
+   uv run python research/finetune.py --preset minicpm5-1b --mode lora --epochs 1
+   ```
+
+4. **Re-eval candidate** with the same lm-eval config:
+   ```bash
+   uv run --package slm-evals slm-lm-eval \
+     --config research/evals/configs/lm_eval_compare_study.yaml \
+     --preset minicpm5-1b-lesson-lora \
+     --experiment-name minicpm5-1b-lora__v1 \
+     --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.
+
+---
+
+## Troubleshooting
+
+| 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 |
+| BFCL / GAIA download slow | Set `max_samples` low first; cache HF datasets under `~/.cache/huggingface` |
+| SWE-bench Docker errors | Keep `full_eval: false` in YAML unless `swebench` + Docker are installed |
+| τ-bench API costs | Keep `use_llm_user: false` (rule-based user simulator) |

research/data/benchmark-kb.jsonl

@@ -0,0 +1,10 @@
+{"text": "Photosynthesis: plants use sunlight, water, and CO2 to make glucose and release oxygen."}
+{"text": "Fractions: the top number counts parts and the bottom number shows total equal parts."}
+{"text": "Water cycle: evaporation turns liquid water into vapor; condensation forms clouds."}
+{"text": "Precipitation returns water to land and bodies of water as rain, snow, sleet, or hail."}
+{"text": "AI agents observe state, decide actions, and update plans using tools and memory."}
+{"text": "American Revolution: colonists protested taxation without representation and trade restrictions."}
+{"text": "Boston Tea Party: colonists dumped tea to protest the Tea Act in 1773."}
+{"text": "Enlightenment ideas about rights and consent spread through pamphlets and speeches."}
+{"text": "One-half means one of two equal parts; one-fourth means one of four equal parts."}
+{"text": "Runoff carries water from land back toward rivers, lakes, and oceans."}

research/data/benchmark-qa.jsonl

@@ -0,0 +1,10 @@
+{"question": "What is photosynthesis?", "answer": "Plants make food using sunlight, water, and carbon dioxide.", "domain": "science"}
+{"question": "What do plants produce during photosynthesis?", "answer": "Glucose and oxygen.", "domain": "science"}
+{"question": "What is a fraction?", "answer": "A fraction shows equal parts of a whole.", "domain": "math"}
+{"question": "What does one-half mean?", "answer": "One of two equal parts.", "domain": "math"}
+{"question": "What is evaporation in the water cycle?", "answer": "Heat turns liquid water into water vapor.", "domain": "science"}
+{"question": "What is precipitation?", "answer": "Water falls as rain, snow, sleet, or hail.", "domain": "science"}
+{"question": "What is an AI agent?", "answer": "Software that takes goals and acts to reach them.", "domain": "technology"}
+{"question": "What is a colonial grievance before the American Revolution?", "answer": "Taxation without representation.", "domain": "history"}
+{"question": "What happened at the Boston Tea Party?", "answer": "Colonists protested the Tea Act.", "domain": "history"}
+{"question": "Why does photosynthesis matter?", "answer": "It feeds nearly all life on Earth and keeps the atmosphere balanced.", "domain": "science"}

research/data/education-lesson-chat.jsonl

@@ -0,0 +1,8 @@
+{"messages": [{"role": "system", "content": "You are a lesson-planning assistant for teachers.\nFollow the skill workflow below and output ONLY valid JSON (no markdown fences).\n\nSkill workflow:\n1. Ask for topic, audience grade, and slide count (3–8 content slides).\n2. Produce a JSON outline with `title` and `slides` (each slide has `title`, `bullets`, `speaker_note`).\n3. Call `create_pptx` with the validated outline.\n4. Return a download link and markdown preview for the teacher.\n\nJSON schema:\n{\n  \"title\": \"string — presentation title\",\n  \"slides\": [\n    {\n      \"title\": \"string — slide heading\",\n      \"bullets\": [\"string\", \"...\"],\n      \"speaker_note\": \"string — one sentence for the teacher\"\n    }\n  ]\n}\n\nRules:\n- Use exactly the requested number of content slides (title slide is added separately by the tool).\n- Bullets should be short, age-appropriate, and factual.\n- speaker_note is optional but helpful for each slide."}, {"role": "user", "content": "Topic: Photosynthesis\nGrade level: 6\nNumber of content slides: 4\nReturn JSON only."}, {"role": "assistant", "content": "{\"title\": \"Photosynthesis for 6th Grade\", \"slides\": [{\"title\": \"What Is Photosynthesis?\", \"bullets\": [\"Plants make their own food using sunlight\", \"The process happens mainly in leaves\"], \"speaker_note\": \"Ask students what they already know about how plants eat.\"}, {\"title\": \"Ingredients Plants Need\", \"bullets\": [\"Sunlight provides energy\", \"Water comes from roots and carbon dioxide from air\"], \"speaker_note\": \"Draw a simple leaf diagram while you explain inputs.\"}, {\"title\": \"What Plants Produce\", \"bullets\": [\"Glucose is the sugar plants use for energy\", \"Oxygen is released into the air\"], \"speaker_note\": \"Connect oxygen release to why forests matter for breathing.\"}, {\"title\": \"Why It Matters\", \"bullets\": [\"Photosynthesis feeds nearly all life on Earth\", \"It helps keep our atmosphere balanced\"], \"speaker_note\": \"End with a quick pair-share: one reason photosynthesis matters.\"}]}"}]}
+{"messages": [{"role": "system", "content": "You are a lesson-planning assistant for teachers.\nFollow the skill workflow below and output ONLY valid JSON (no markdown fences).\n\nSkill workflow:\n1. Ask for topic, audience grade, and slide count (3–8 content slides).\n2. Produce a JSON outline with `title` and `slides` (each slide has `title`, `bullets`, `speaker_note`).\n3. Call `create_pptx` with the validated outline.\n4. Return a download link and markdown preview for the teacher.\n\nJSON schema:\n{\n  \"title\": \"string — presentation title\",\n  \"slides\": [\n    {\n      \"title\": \"string — slide heading\",\n      \"bullets\": [\"string\", \"...\"],\n      \"speaker_note\": \"string — one sentence for the teacher\"\n    }\n  ]\n}\n\nRules:\n- Use exactly the requested number of content slides (title slide is added separately by the tool).\n- Bullets should be short, age-appropriate, and factual.\n- speaker_note is optional but helpful for each slide."}, {"role": "user", "content": "Topic: Fractions\nGrade level: 4\nNumber of content slides: 3\nReturn JSON only."}, {"role": "assistant", "content": "{\"title\": \"Understanding Fractions\", \"slides\": [{\"title\": \"Parts of a Whole\", \"bullets\": [\"A fraction shows equal parts of a whole\", \"The top number counts parts; the bottom shows total equal parts\"], \"speaker_note\": \"Use a paper pizza cut into equal slices as a visual.\"}, {\"title\": \"Reading Common Fractions\", \"bullets\": [\"One-half means one of two equal parts\", \"One-fourth means one of four equal parts\"], \"speaker_note\": \"Have students name fractions they see in the room.\"}, {\"title\": \"Comparing Fractions\", \"bullets\": [\"Larger denominators can mean smaller pieces\", \"Use drawings to compare two simple fractions\"], \"speaker_note\": \"Give two fraction cards and ask which shows more pizza.\"}]}"}]}
+{"messages": [{"role": "system", "content": "You are a lesson-planning assistant for teachers.\nFollow the skill workflow below and output ONLY valid JSON (no markdown fences).\n\nSkill workflow:\n1. Ask for topic, audience grade, and slide count (3–8 content slides).\n2. Produce a JSON outline with `title` and `slides` (each slide has `title`, `bullets`, `speaker_note`).\n3. Call `create_pptx` with the validated outline.\n4. Return a download link and markdown preview for the teacher.\n\nJSON schema:\n{\n  \"title\": \"string — presentation title\",\n  \"slides\": [\n    {\n      \"title\": \"string — slide heading\",\n      \"bullets\": [\"string\", \"...\"],\n      \"speaker_note\": \"string — one sentence for the teacher\"\n    }\n  ]\n}\n\nRules:\n- Use exactly the requested number of content slides (title slide is added separately by the tool).\n- Bullets should be short, age-appropriate, and factual.\n- speaker_note is optional but helpful for each slide."}, {"role": "user", "content": "Topic: The water cycle\nGrade level: 5\nNumber of content slides: 5\nReturn JSON only."}, {"role": "assistant", "content": "{\"title\": \"The Water Cycle\", \"slides\": [{\"title\": \"Water Moves Constantly\", \"bullets\": [\"Water travels between Earth, air, and living things\", \"The same water has been recycled for millions of years\"], \"speaker_note\": \"Show a simple cycle diagram before diving into each stage.\"}, {\"title\": \"Evaporation\", \"bullets\": [\"Heat turns liquid water into water vapor\", \"Oceans, lakes, and puddles lose water to the air\"], \"speaker_note\": \"Ask what happens to a wet sidewalk on a sunny day.\"}, {\"title\": \"Condensation\", \"bullets\": [\"Water vapor cools and forms tiny droplets\", \"Clouds form when droplets gather in the sky\"], \"speaker_note\": \"Hold a cold cup and point out the outside moisture.\"}, {\"title\": \"Precipitation\", \"bullets\": [\"Droplets combine and fall as rain, snow, sleet, or hail\", \"Precipitation returns water to land and water bodies\"], \"speaker_note\": \"Let students name types of precipitation they have seen.\"}, {\"title\": \"Collection and Runoff\", \"bullets\": [\"Water collects in rivers, lakes, and groundwater\", \"Runoff carries water back toward oceans\"], \"speaker_note\": \"Trace a raindrop's path on your diagram to close the loop.\"}]}"}]}
+{"messages": [{"role": "system", "content": "You are a lesson-planning assistant for teachers.\nFollow the skill workflow below and output ONLY valid JSON (no markdown fences).\n\nSkill workflow:\n1. Ask for topic, audience grade, and slide count (3–8 content slides).\n2. Produce a JSON outline with `title` and `slides` (each slide has `title`, `bullets`, `speaker_note`).\n3. Call `create_pptx` with the validated outline.\n4. Return a download link and markdown preview for the teacher.\n\nJSON schema:\n{\n  \"title\": \"string — presentation title\",\n  \"slides\": [\n    {\n      \"title\": \"string — slide heading\",\n      \"bullets\": [\"string\", \"...\"],\n      \"speaker_note\": \"string — one sentence for the teacher\"\n    }\n  ]\n}\n\nRules:\n- Use exactly the requested number of content slides (title slide is added separately by the tool).\n- Bullets should be short, age-appropriate, and factual.\n- speaker_note is optional but helpful for each slide."}, {"role": "user", "content": "Topic: AI agents\nGrade level: 8\nNumber of content slides: 4\nReturn JSON only."}, {"role": "assistant", "content": "{\"title\": \"Introduction to AI Agents\", \"slides\": [{\"title\": \"What Is an AI Agent?\", \"bullets\": [\"An AI agent is software that takes goals and acts to reach them\", \"Agents can use tools, memory, and feedback loops\"], \"speaker_note\": \"Contrast a chatbot that only replies with an agent that plans steps.\"}, {\"title\": \"How Agents Work\", \"bullets\": [\"They observe input or environment state\", \"They decide actions and update their plan\"], \"speaker_note\": \"Walk through a simple example like booking a meeting.\"}, {\"title\": \"Everyday Examples\", \"bullets\": [\"Navigation apps reroute based on traffic\", \"Email filters can sort and draft responses\"], \"speaker_note\": \"Invite students to suggest other agent-like tools they use.\"}, {\"title\": \"Benefits and Risks\", \"bullets\": [\"Agents can save time on repetitive tasks\", \"They need clear rules and human oversight\"], \"speaker_note\": \"Close with a short discussion on when automation helps or hurts.\"}]}"}]}
+{"messages": [{"role": "system", "content": "You are a lesson-planning assistant for teachers.\nFollow the skill workflow below and output ONLY valid JSON (no markdown fences).\n\nSkill workflow:\n1. Ask for topic, audience grade, and slide count (3–8 content slides).\n2. Produce a JSON outline with `title` and `slides` (each slide has `title`, `bullets`, `speaker_note`).\n3. Call `create_pptx` with the validated outline.\n4. Return a download link and markdown preview for the teacher.\n\nJSON schema:\n{\n  \"title\": \"string — presentation title\",\n  \"slides\": [\n    {\n      \"title\": \"string — slide heading\",\n      \"bullets\": [\"string\", \"...\"],\n      \"speaker_note\": \"string — one sentence for the teacher\"\n    }\n  ]\n}\n\nRules:\n- Use exactly the requested number of content slides (title slide is added separately by the tool).\n- Bullets should be short, age-appropriate, and factual.\n- speaker_note is optional but helpful for each slide."}, {"role": "user", "content": "Topic: The American Revolution\nGrade level: 8\nNumber of content slides: 5\nReturn JSON only."}, {"role": "assistant", "content": "{\"title\": \"Causes of the American Revolution\", \"slides\": [{\"title\": \"Colonial Grievances\", \"bullets\": [\"Colonists protested taxation without representation\", \"Trade restrictions limited colonial economies\"], \"speaker_note\": \"Connect taxes on everyday goods to student purchases today.\"}, {\"title\": \"Key Events\", \"bullets\": [\"The Boston Tea Party protested the Tea Act\", \"The Intolerable Acts tightened British control\"], \"speaker_note\": \"Place events on a timeline as you introduce each one.\"}, {\"title\": \"Ideas That Spread\", \"bullets\": [\"Enlightenment ideas emphasized rights and consent\", \"Pamphlets and speeches spread revolutionary arguments\"], \"speaker_note\": \"Read a short excerpt from a primary source if time allows.\"}, {\"title\": \"Leaders and Groups\", \"bullets\": [\"Committees and congresses organized resistance\", \"Leaders like Adams and Washington gained prominence\"], \"speaker_note\": \"Highlight diverse roles beyond famous military figures.\"}, {\"title\": \"Road to Independence\", \"bullets\": [\"The Declaration explained the break from Britain\", \"War followed when compromise failed\"], \"speaker_note\": \"End by asking what principle students would defend most strongly.\"}]}"}]}
+{"messages": [{"role": "system", "content": "You are a lesson-planning assistant for teachers.\nFollow the skill workflow below and output ONLY valid JSON (no markdown fences).\n\nSkill workflow:\n1. Ask for topic, audience grade, and slide count (3–8 content slides).\n2. Produce a JSON outline with `title` and `slides` (each slide has `title`, `bullets`, `speaker_note`).\n3. Call `create_pptx` with the validated outline.\n4. Return a download link and markdown preview for the teacher.\n\nJSON schema:\n{\n  \"title\": \"string — presentation title\",\n  \"slides\": [\n    {\n      \"title\": \"string — slide heading\",\n      \"bullets\": [\"string\", \"...\"],\n      \"speaker_note\": \"string — one sentence for the teacher\"\n    }\n  ]\n}\n\nRules:\n- Use exactly the requested number of content slides (title slide is added separately by the tool).\n- Bullets should be short, age-appropriate, and factual.\n- speaker_note is optional but helpful for each slide."}, {"role": "user", "content": "Topic: Kindness and friendship\nGrade level: 2\nNumber of content slides: 3\nReturn JSON only."}, {"role": "assistant", "content": "{\"title\": \"Being a Good Friend\", \"slides\": [{\"title\": \"What Friends Do\", \"bullets\": [\"Friends listen and share\", \"Friends help when someone feels sad\"], \"speaker_note\": \"Ask students to share one kind thing a friend did for them.\"}, {\"title\": \"Using Kind Words\", \"bullets\": [\"Please and thank you show respect\", \"Compliments can make others feel happy\"], \"speaker_note\": \"Practice turning an unkind sentence into a kind one together.\"}, {\"title\": \"Solving Small Problems\", \"bullets\": [\"Take turns when you disagree\", \"Ask a grown-up if you need help\"], \"speaker_note\": \"Role-play a simple sharing conflict and a peaceful solution.\"}]}"}]}
+{"messages": [{"role": "system", "content": "You are a lesson-planning assistant for teachers.\nFollow the skill workflow below and output ONLY valid JSON (no markdown fences).\n\nSkill workflow:\n1. Ask for topic, audience grade, and slide count (3–8 content slides).\n2. Produce a JSON outline with `title` and `slides` (each slide has `title`, `bullets`, `speaker_note`).\n3. Call `create_pptx` with the validated outline.\n4. Return a download link and markdown preview for the teacher.\n\nJSON schema:\n{\n  \"title\": \"string — presentation title\",\n  \"slides\": [\n    {\n      \"title\": \"string — slide heading\",\n      \"bullets\": [\"string\", \"...\"],\n      \"speaker_note\": \"string — one sentence for the teacher\"\n    }\n  ]\n}\n\nRules:\n- Use exactly the requested number of content slides (title slide is added separately by the tool).\n- Bullets should be short, age-appropriate, and factual.\n- speaker_note is optional but helpful for each slide."}, {"role": "user", "content": "Topic: Linear equations\nGrade level: 9\nNumber of content slides: 4\nReturn JSON only."}, {"role": "assistant", "content": "{\"title\": \"Introduction to Linear Equations\", \"slides\": [{\"title\": \"What Makes an Equation Linear?\", \"bullets\": [\"A linear equation graphs as a straight line\", \"The highest power of the variable is one\"], \"speaker_note\": \"Show y = 2x + 1 on a coordinate plane.\"}, {\"title\": \"Slope-Intercept Form\", \"bullets\": [\"y = mx + b uses slope m and y-intercept b\", \"m tells how steep the line is\"], \"speaker_note\": \"Have students identify m and b from one example.\"}, {\"title\": \"Graphing from an Equation\", \"bullets\": [\"Plot the y-intercept first\", \"Use the slope to find another point\"], \"speaker_note\": \"Work one example slowly on the board.\"}, {\"title\": \"Real-World Connections\", \"bullets\": [\"Phone plans with a flat fee plus per-minute cost are linear\", \"Distance at constant speed changes linearly over time\"], \"speaker_note\": \"Ask students to invent their own linear situation.\"}]}"}]}
+{"messages": [{"role": "system", "content": "You are a lesson-planning assistant for teachers.\nFollow the skill workflow below and output ONLY valid JSON (no markdown fences).\n\nSkill workflow:\n1. Ask for topic, audience grade, and slide count (3–8 content slides).\n2. Produce a JSON outline with `title` and `slides` (each slide has `title`, `bullets`, `speaker_note`).\n3. Call `create_pptx` with the validated outline.\n4. Return a download link and markdown preview for the teacher.\n\nJSON schema:\n{\n  \"title\": \"string — presentation title\",\n  \"slides\": [\n    {\n      \"title\": \"string — slide heading\",\n      \"bullets\": [\"string\", \"...\"],\n      \"speaker_note\": \"string — one sentence for the teacher\"\n    }\n  ]\n}\n\nRules:\n- Use exactly the requested number of content slides (title slide is added separately by the tool).\n- Bullets should be short, age-appropriate, and factual.\n- speaker_note is optional but helpful for each slide."}, {"role": "user", "content": "Topic: The solar system\nGrade level: 3\nNumber of content slides: 4\nReturn JSON only."}, {"role": "assistant", "content": "{\"title\": \"Our Solar System\", \"slides\": [{\"title\": \"The Sun at the Center\", \"bullets\": [\"The Sun is a star that gives Earth light and heat\", \"Planets orbit, or travel around, the Sun\"], \"speaker_note\": \"Use a lamp and balls to model orbits if you have them.\"}, {\"title\": \"Inner Planets\", \"bullets\": [\"Mercury, Venus, Earth, and Mars are rocky planets\", \"Earth is the only one known to have life\"], \"speaker_note\": \"Compare sizes with classroom objects.\"}, {\"title\": \"Outer Planets\", \"bullets\": [\"Jupiter and Saturn are gas giants\", \"Uranus and Neptune are icy giants farther from the Sun\"], \"speaker_note\": \"Mention that Jupiter is the largest planet.\"}, {\"title\": \"Other Objects\", \"bullets\": [\"Asteroids and comets also orbit the Sun\", \"The Moon orbits Earth, not the Sun directly\"], \"speaker_note\": \"Clarify that moons are not planets.\"}]}"}]}

research/docs/overview.md

@@ -0,0 +1,114 @@
+# Research overview
+
+How `research/` relates to the main hackathon repo and what each component does.
+
+## Position in the repo
+
+```text
+small-model-hackathon/
+├── apps/gradio-space/     ← shipped Lesson Agent UI
+├── libs/agent/            ← skill loop, tools, traces
+├── libs/inference/        ← transformers + llama.cpp backends
+├── models.yaml            ← model presets (shared with finetune)
+└── 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.
+
+## Three tracks
+
+### Fine-tuning
+
+`research/finetune.py` adapts a small HF causal LM on instruction or chat data. It reuses root `models.yaml` presets and the shared inference config loader, so the same `minicpm5-1b` preset used in the Gradio app can be fine-tuned without duplicating model metadata.
+
+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)
+
+`research/ensemble/` explores a modular stack inspired by LeCun-style architectures:
+
+```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.
+
+## Data flow
+
+```mermaid
+flowchart LR
+  subgraph data [research/data]
+    lesson[education-lesson-chat.jsonl]
+    qa[benchmark-qa.jsonl]
+    kb[benchmark-kb.jsonl]
+  end
+
+  subgraph train [finetune.py]
+    ckpt[models/finetuned/]
+  end
+
+  subgraph evals [slm-evals]
+    bfcl[BFCL]
+    tau[tau-bench]
+    gaia[GAIA]
+    swe[SWE-bench]
+  end
+
+  subgraph ens [ensemble]
+    jepa[JEPA harness]
+    world[World harness]
+  end
+
+  lesson --> train
+  train --> ckpt
+  ckpt --> evals
+  qa --> jepa
+  kb --> jepa
+  qa --> world
+  kb --> world
+```
+
+## When to use which tool
+
+| Goal | Tool |
+| ---- | ---- |
+| 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 |
+| 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`:
+
+- `research/ensemble` → import name `ensemble`
+- `research/evals` → import name `slm_evals`, CLI `slm-benchmark`
+
+Run with `uv run --package <name>` from the repo root so uv resolves workspace paths and shared lockfile versions.

research/ensemble/README.md

@@ -0,0 +1,113 @@
+# 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

@@ -0,0 +1,16 @@
+[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

@@ -0,0 +1,35 @@
+#!/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

@@ -0,0 +1,15 @@
+"""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

@@ -0,0 +1,418 @@
+"""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

@@ -0,0 +1,28 @@
+"""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

@@ -0,0 +1,149 @@
+"""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

@@ -0,0 +1,163 @@
+"""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

@@ -0,0 +1,45 @@
+"""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

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

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

@@ -0,0 +1,266 @@
+"""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

@@ -0,0 +1,42 @@
+"""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

@@ -0,0 +1,174 @@
+"""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

@@ -0,0 +1,309 @@
+"""
+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

@@ -0,0 +1,75 @@
+"""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

@@ -0,0 +1,232 @@
+"""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

@@ -0,0 +1,507 @@
+"""
+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

@@ -0,0 +1,46 @@
+"""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

@@ -0,0 +1,198 @@
+"""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

@@ -0,0 +1,228 @@
+"""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

@@ -0,0 +1,40 @@
+"""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

@@ -0,0 +1,499 @@
+"""
+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

@@ -0,0 +1,6 @@
+"""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/README.md

@@ -0,0 +1,98 @@
+# SLM Agentic Benchmark Suite
+
+A uv workspace package to evaluate **local HuggingFace models** against agentic and academic benchmarks.
+
+**Docs:** [USAGE.md](USAGE.md) (commands and workflows) · [docs/benchmarks.md](docs/benchmarks.md) (per-benchmark reference) · [../USAGE.md](../USAGE.md) (full research tree)
+
+| Suite | CLI | What it measures |
+|---|---|---|
+| **Agentic** | `slm-benchmark` | BFCL, τ-bench, GAIA, SWE-bench |
+| **Academic** | `slm-lm-eval` | ARC, HellaSwag, GSM8K, … (lm-evaluation-harness) |
+
+## Install
+
+From the repo root:
+
+```bash
+uv sync --group evals
+uv sync --group lm-eval   # optional: slm-lm-eval academic benchmarks
+```
+
+## Quickstart
+
+```bash
+# From repo root (recommended)
+uv run --package slm-evals slm-benchmark \
+  --model openbmb/MiniCPM5-1B \
+  --benchmarks bfcl tau_bench \
+  --max-samples 20
+
+# Or as a module
+uv run --package slm-evals python -m slm_evals.run_benchmark \
+  --model openbmb/MiniCPM5-1B \
+  --benchmarks bfcl tau_bench \
+  --max-samples 20
+
+# YAML config
+uv run --package slm-evals slm-benchmark \
+  --config research/evals/configs/experiment_001.yaml
+```
+
+## Project structure
+
+```
+research/evals/
+├── pyproject.toml
+├── configs/
+│   └── experiment_001.yaml
+├── src/slm_evals/
+│   ├── run_benchmark.py
+│   ├── benchmarks/
+│   │   ├── base.py
+│   │   ├── bfcl.py
+│   │   ├── tau_bench.py
+│   │   ├── gaia.py
+│   │   └── swe_bench.py
+│   └── utils/
+│       ├── model_loader.py
+│       ├── reporter.py
+│       └── config_loader.py
+└── results/              # created at runtime (relative to cwd)
+```
+
+## CLI reference
+
+```
+--model          Path to local HF model dir (or Hub ID)
+--benchmarks     Space-separated: bfcl tau_bench gaia swe_bench all
+--config         YAML config file (overrides CLI flags)
+--max-samples    Cap samples per benchmark
+--output-dir     Results directory (default: ./results)
+--experiment-name  Tag for this run
+--device         auto | cpu | cuda | cuda:0
+--dtype          float32 | float16 | bfloat16 | int8 | int4
+--max-new-tokens Max tokens per generation (default: 512)
+--temperature    Sampling temp (default: 0.0 = greedy)
+```
+
+## Adding a custom benchmark
+
+1. Create `src/slm_evals/benchmarks/my_bench.py` and subclass `BaseBenchmark`.
+2. Register it in `src/slm_evals/run_benchmark.py` → `BENCHMARK_REGISTRY`.
+3. Run: `uv run --package slm-evals slm-benchmark --model ./my-model --benchmarks my_bench`
+
+## Output formats
+
+Results are written under `<output-dir>/<experiment_name>/`:
+
+- `results.json` — full structured dump
+- `results.csv` — one row per benchmark
+- `report.md` — human-readable summary
+
+## Notes
+
+**τ-bench user simulator**: Default is a lightweight rule-based simulator. Set `use_llm_user: true` in config for the GPT-4o user agent (API cost).
+
+**SWE-bench full eval**: Set `full_eval: true` to run the official Docker harness (`pip install swebench docker`).
+
+**GAIA tools**: Offline by default (`tool_mode: describe`). Wire real tools in `gaia.py` for live eval.

research/evals/USAGE.md

@@ -0,0 +1,376 @@
+# Evals usage
+
+Run the **SLM Agentic Benchmark Suite** (`slm-evals`) against a local HuggingFace model directory or Hub id.
+
+Benchmark details: [docs/benchmarks.md](docs/benchmarks.md). Package overview: [README.md](README.md).
+
+## Install
+
+From the repo root:
+
+```bash
+uv sync --group evals
+```
+
+For academic benchmarks (lm-evaluation-harness):
+
+```bash
+uv sync --group lm-eval
+```
+
+This installs the `slm-evals` workspace package and registers the `slm-benchmark` and `slm-lm-eval` console scripts.
+
+## Quick start
+
+```bash
+# Two benchmarks, capped samples (good first run)
+uv run --package slm-evals slm-benchmark \
+  --model openbmb/MiniCPM5-1B \
+  --benchmarks bfcl tau_bench \
+  --max-samples 20
+
+# All four benchmarks
+uv run --package slm-evals slm-benchmark \
+  --model ./models/finetuned/minicpm5-1b-lora \
+  --benchmarks all \
+  --max-samples 50
+
+# Equivalent module invocation
+uv run --package slm-evals python -m slm_evals.run_benchmark \
+  --model openbmb/MiniCPM5-1B \
+  --benchmarks bfcl \
+  --max-samples 10
+```
+
+## Config-driven runs
+
+Copy and edit the template, then pass `--config`:
+
+```bash
+cp research/evals/configs/experiment_001.yaml research/evals/configs/my_run.yaml
+# edit model_path, benchmarks, max_samples, overrides
+
+uv run --package slm-evals slm-benchmark \
+  --config research/evals/configs/my_run.yaml
+```
+
+When `--config` is set, **YAML values override CLI flags**. Use configs for reproducible experiment names and per-benchmark settings.
+
+### Template fields
+
+| Key | Description |
+| --- | ----------- |
+| `model_path` | Local directory or HF Hub id |
+| `device` | `auto`, `cpu`, `cuda`, `cuda:0`, … |
+| `dtype` | `float32`, `float16`, `bfloat16`, `int8`, `int4` |
+| `max_new_tokens` | Cap per generation (default 512) |
+| `temperature` | `0.0` = greedy (recommended for evals) |
+| `experiment_name` | Folder name under `output_dir` |
+| `output_dir` | Root for results (default `results`) |
+| `benchmarks` | List: `bfcl`, `tau_bench`, `gaia`, `swe_bench` |
+| `max_samples` | Cap per benchmark; omit or `null` for full split |
+| `benchmark_overrides` | Per-benchmark dict (see [docs/benchmarks.md](docs/benchmarks.md)) |
+
+---
+
+## CLI reference
+
+```
+slm-benchmark [OPTIONS]
+
+--list-benchmarks       Show agentic benchmark keys and preset suites
+--model PATH            Local HF dir or Hub id (required unless --config)
+--benchmarks NAMES      bfcl tau_bench gaia swe_bench all  (default: all)
+--config PATH           YAML config (overrides other flags)
+--max-samples N         Cap samples per benchmark
+--output-dir DIR        Results root (default: ./results)
+--experiment-name TAG   Run folder name (auto timestamp if omitted)
+--device MAP            auto | cpu | cuda | cuda:0
+--dtype TYPE            float32 | float16 | bfloat16 | int8 | int4
+--max-new-tokens N      Default 512
+--temperature T         Default 0.0
+```
+
+---
+
+## Results
+
+Each run writes to `<output_dir>/<experiment_name>/`:
+
+| File | Contents |
+| ---- | -------- |
+| `results.json` | Full structured payload (per-sample + aggregates) |
+| `results.csv` | One row per benchmark |
+| `report.md` | Human-readable summary |
+
+Example layout:
+
+```text
+results/
+└── minicpm5-1b__bfcl-tau__v1/
+    ├── results.json
+    ├── results.csv
+    └── report.md
+```
+
+`output_dir` is relative to **current working directory**. Run from repo root so paths stay predictable, or set an absolute `output_dir` in YAML.
+
+---
+
+## Per-benchmark tips
+
+### BFCL (function calling)
+
+- Default: downloads from `gorilla-llm/Berkeley-Function-Calling-Leaderboard`
+- `strict: false` in YAML — fuzzy argument matching (better for small models)
+- Local JSONL: set `benchmark_overrides.bfcl.data_path`
+
+### τ-bench (multi-turn tools)
+
+- Domains: `retail`, `airline`, or `both`
+- `use_llm_user: false` — free rule-based user simulator (default)
+- `use_llm_user: true` — GPT-4o user agent (**API cost**)
+
+### GAIA
+
+- Default split: `validation` (public)
+- `tool_mode: describe` — offline tool descriptions (no live web)
+- Level filter: `levels: [1, 2]` or `[1, 2, 3]`
+
+### SWE-bench Verified
+
+- Default: lightweight patch-generation scoring (no Docker)
+- `full_eval: true` — official harness (`pip install swebench docker`)
+
+See [docs/benchmarks.md](docs/benchmarks.md) for scoring semantics.
+
+---
+
+## lm-evaluation-harness (`slm-lm-eval`)
+
+Run standard academic benchmarks (ARC, HellaSwag, PIQA, BoolQ, GSM8K) via [EleutherAI lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness).
+
+Install: `uv sync --group lm-eval`
+
+Full profile guide: [docs/eval_profiles.md](docs/eval_profiles.md)
+
+### Discover profiles and tasks
+
+```bash
+# Claim-matched lm-eval profiles (reasoning, code, smoke, …)
+uv run --package slm-evals slm-lm-eval --list-profiles
+
+# Also show agentic suites + external benchmark notes
+uv run --package slm-evals slm-lm-eval --list-profiles-all
+
+# lm-eval task names
+uv run --package slm-evals slm-lm-eval --list-tasks
+
+# Agentic benchmarks (BFCL, τ-bench, GAIA, SWE)
+uv run --package slm-evals slm-benchmark --list-benchmarks
+```
+
+### Quick start
+
+```bash
+# By profile name (recommended)
+uv run --package slm-evals slm-lm-eval \
+  --profile reasoning \
+  --preset minicpm5-1b \
+  --experiment-name minicpm5-1b__reasoning-baseline
+
+# Smoke profile (25 samples)
+uv run --package slm-evals slm-lm-eval \
+  --profile smoke \
+  --preset minicpm5-1b \
+  --experiment-name minicpm5-1b__smoke
+
+# LoRA adapter via preset (base + peft resolved automatically)
+uv run --package slm-evals slm-lm-eval \
+  --config research/evals/configs/lm_eval_minicpm5.yaml \
+  --preset minicpm5-1b-lesson-lora \
+  --experiment-name minicpm5-1b-lora__v1
+
+# Explicit base + adapter
+uv run --package slm-evals slm-lm-eval \
+  --config research/evals/configs/lm_eval_smoke.yaml \
+  --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
+
+Use the **same config** for both runs; only change `--preset` / `--experiment-name`:
+
+```bash
+uv run --package slm-evals slm-lm-eval \
+  --config research/evals/configs/lm_eval_compare_study.yaml \
+  --preset minicpm5-1b \
+  --experiment-name minicpm5-1b__baseline
+
+uv run --package slm-evals slm-lm-eval \
+  --config research/evals/configs/lm_eval_compare_study.yaml \
+  --preset minicpm5-1b-lesson-lora \
+  --experiment-name minicpm5-1b-lora__v1 \
+  --compare-to results/lm_eval/minicpm5-1b__baseline/results.json
+```
+
+### Config templates
+
+Catalog: `configs/eval_profiles.yaml` — maps **claim → profile → tasks**.
+
+| Profile (`--profile`) | Config file | Purpose |
+| --------------------- | ----------- | ------- |
+| `smoke` | `lm_eval_smoke.yaml` | Fast validation (`limit: 25`, 2 tasks) |
+| `reasoning` | `lm_eval_reasoning.yaml` | Math + commonsense (GSM8K, ARC, HellaSwag) |
+| `understanding` | `lm_eval_understanding.yaml` | NLU (BoolQ, PIQA, COPA, RTE) |
+| `code` | `lm_eval_code.yaml` | HumanEval + MBPP |
+| `instructions` | `lm_eval_instructions.yaml` | IFEval instruction following |
+| `general_slm` | `lm_eval_minicpm5.yaml` | Full ~1B SLM profile (6 tasks) |
+| `compare_study` | `lm_eval_compare_study.yaml` | Baseline vs finetune comparison defaults |
+
+| Key | Description |
+| --- | ----------- |
+| `tasks` | lm-eval task names (e.g. `arc_easy`, `gsm8k`) |
+| `num_fewshot` | Few-shot count (gsm8k may use task default 8) |
+| `limit` | Max samples per task; `null` = full split |
+| `seed` | Random seed (applied to all lm-eval RNGs) |
+| `batch_size` | `auto` or integer |
+| `device` | `auto`, `cpu`, `cuda`, … |
+| `dtype` | `bfloat16`, `float16`, `int4`, … |
+| `trust_remote_code` | Required for MiniCPM / Gemma presets |
+| `output_dir` | Root for runs (default `results/lm_eval`) |
+
+### CLI reference
+
+```
+slm-lm-eval [OPTIONS]
+
+--list-profiles         Show claim-matched profiles and example commands
+--list-profiles-all     Include agentic suites and external benchmark notes
+--list-tasks            List lm-eval task names (catalog fallback if not installed)
+--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
+--adapter PATH          LoRA adapter (alternative to preset adapter_path)
+--tasks NAMES           Override task list
+--num-fewshot N
+--limit N               Cap samples per task
+--seed N
+--batch-size VALUE
+--device MAP
+--dtype TYPE
+--output-dir DIR        Default: results/lm_eval
+--experiment-name TAG   Run folder name
+--compare-to PATH       Baseline results.json for delta table
+```
+
+### Results
+
+Each run writes to `<output_dir>/<experiment_name>/`:
+
+| File | Contents |
+| ---- | -------- |
+| `results.json` | lm-eval native payload + `run_meta` |
+| `summary.md` | Task → metric table |
+| `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`.
+
+---
+
+## Adding a custom benchmark
+
+1. Create `src/slm_evals/benchmarks/my_bench.py` subclassing `BaseBenchmark`:
+   - `load_dataset()` → list of sample dicts
+   - `build_prompt(sample)` → prompt string
+   - `evaluate_sample(sample, prediction)` → `{passed, score, note}`
+
+2. Register in `src/slm_evals/run_benchmark.py` → `BENCHMARK_REGISTRY`.
+
+3. Run:
+   ```bash
+   uv run --package slm-evals slm-benchmark \
+     --model ./my-model --benchmarks my_bench --max-samples 10
+   ```
+
+---
+
+## Suggested workflows
+
+### Smoke (CPU/GPU, ~5 min)
+
+```bash
+uv run --package slm-evals slm-benchmark \
+  --model openbmb/MiniCPM5-1B \
+  --benchmarks bfcl \
+  --max-samples 5 \
+  --device cpu
+```
+
+### Before / after fine-tune
+
+```bash
+BASE=openbmb/MiniCPM5-1B
+ADAPTER=./models/finetuned/minicpm5-1b-lora
+
+for M in "$BASE" "$ADAPTER"; do
+  uv run --package slm-evals slm-benchmark \
+    --model "$M" \
+    --benchmarks bfcl tau_bench \
+    --max-samples 100 \
+    --experiment-name "$(basename "$M")__bfcl-tau"
+done
+```
+
+### Full experiment (YAML)
+
+Edit `configs/experiment_001.yaml` with your `model_path` and `experiment_name`, then:
+
+```bash
+uv run --package slm-evals slm-benchmark \
+  --config research/evals/configs/experiment_001.yaml
+```
+
+---
+
+## Troubleshooting
+
+| Symptom | Likely cause | Fix |
+| ------- | ------------ | --- |
+| `error: --model is required` | No `--config` and no `--model` | Pass one of them |
+| CUDA OOM | Model too large for VRAM | `--dtype int4` or `--device cpu` |
+| HF dataset 401 on GAIA test | Gated split | Use `split: validation` |
+| τ-bench hangs / costs | LLM user enabled | Set `use_llm_user: false` |
+| Empty `results/` | Wrong cwd | Run from repo root or use absolute `output_dir` |
+| Import errors | Evals group not synced | `uv sync --group evals` |
+
+---
+
+## Entry points
+
+| Path | Role |
+| ---- | ---- |
+| `slm-benchmark` | Agentic benchmarks (BFCL, τ-bench, GAIA, SWE) |
+| `slm-lm-eval` | Academic benchmarks via lm-evaluation-harness |
+| `python -m slm_evals.run_benchmark` | Same as `slm-benchmark` |
+| `python -m slm_evals.run_lm_eval` | Same as `slm-lm-eval` |
+| `research/evals/run_benchmark.py` | Thin shim for backward compatibility |

research/evals/configs/ensemble_jepa_lesson.yaml

@@ -0,0 +1,24 @@
+# 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

research/evals/configs/eval_profiles.yaml

@@ -0,0 +1,160 @@
+# Eval profile catalog — match your model claim to a profile, then run with:
+#   slm-lm-eval --list-profiles
+#   slm-lm-eval --profile reasoning --preset minicpm5-1b --experiment-name my-run
+#
+# See research/evals/docs/eval_profiles.md for full guide.
+
+profiles:
+  smoke:
+    tool: slm-lm-eval
+    claim: Quick sanity check
+    description: Fast validation before a long run (2 tasks, capped samples).
+    config: lm_eval_smoke.yaml
+    tasks:
+      - arc_easy
+      - hellaswag
+
+  reasoning:
+    tool: slm-lm-eval
+    claim: Better reasoning
+    description: Math + commonsense + science QA (generation + multiple-choice).
+    config: lm_eval_reasoning.yaml
+    tasks:
+      - gsm8k
+      - arc_easy
+      - arc_challenge
+      - hellaswag
+
+  understanding:
+    tool: slm-lm-eval
+    claim: Better language understanding
+    description: NLU / reading comprehension (SuperGLUE-style multiple-choice).
+    config: lm_eval_understanding.yaml
+    tasks:
+      - boolq
+      - piqa
+      - copa
+      - rte
+
+  code:
+    tool: slm-lm-eval
+    claim: Better code generation
+    description: Python function synthesis (HumanEval + MBPP via lm-eval).
+    config: lm_eval_code.yaml
+    tasks:
+      - humaneval
+      - mbpp
+
+  instructions:
+    tool: slm-lm-eval
+    claim: Better instruction following
+    description: Verifiable instruction constraints (IFEval).
+    config: lm_eval_instructions.yaml
+    tasks:
+      - ifeval
+
+  general_slm:
+    tool: slm-lm-eval
+    claim: General ~1B SLM baseline
+    description: Balanced academic mix for before/after fine-tune on chat data.
+    config: lm_eval_minicpm5.yaml
+    tasks:
+      - arc_easy
+      - arc_challenge
+      - hellaswag
+      - piqa
+      - boolq
+      - gsm8k
+
+  compare_study:
+    tool: slm-lm-eval
+    claim: Baseline vs finetune comparison
+    description: Same tasks as general_slm with limit 100 for paired studies.
+    config: lm_eval_compare_study.yaml
+    tasks:
+      - arc_easy
+      - arc_challenge
+      - hellaswag
+      - piqa
+      - boolq
+      - gsm8k
+
+suites:
+  agentic_tool_use:
+    tool: slm-benchmark
+    claim: Tool use and function calling
+    description: BFCL + tau-bench for agents that call tools.
+    command: >-
+      uv run --package slm-evals slm-benchmark
+      --model <model> --benchmarks bfcl tau_bench --max-samples 50
+    benchmarks:
+      - bfcl
+      - tau_bench
+
+  agentic_gaia:
+    tool: slm-benchmark
+    claim: End-to-end assistant tasks
+    description: GAIA — multi-step reasoning with optional tools.
+    command: >-
+      uv run --package slm-evals slm-benchmark
+      --model <model> --benchmarks gaia --max-samples 20
+    benchmarks:
+      - gaia
+
+  agentic_code:
+    tool: slm-benchmark
+    claim: Real-world code repair
+    description: SWE-bench Verified patch generation (lightweight mode by default).
+    command: >-
+      uv run --package slm-evals slm-benchmark
+      --model <model> --benchmarks swe_bench --max-samples 10
+    benchmarks:
+      - swe_bench
+
+  agentic_all:
+    tool: slm-benchmark
+    claim: Full agentic suite
+    description: All four slm-benchmark benchmarks.
+    command: >-
+      uv run --package slm-evals slm-benchmark
+      --model <model> --benchmarks all --max-samples 50
+    benchmarks:
+      - bfcl
+      - tau_bench
+      - gaia
+      - swe_bench
+
+  jepa_selector:
+    tool: jepa_harness
+    claim: JEPA / draft selector quality
+    description: >-
+      Domain QA on benchmark-qa.jsonl; compares first vs random vs JEPA vs oracle
+      draft selection. Use for ensemble checkpoints, not base LLM alone.
+    command: >-
+      uv run --package ensemble python -m ensemble.eval.jepa_harness
+      --checkpoint ./models/ensemble/<name> --data research/data/benchmark-qa.jsonl
+
+  world_model:
+    tool: world_harness
+    claim: World model / energy-based planning
+    description: Energy-ranked draft selection on domain QA (world-model track).
+    command: >-
+      uv run --package ensemble python -m ensemble.eval.world_harness
+      --checkpoint ./models/ensemble/<name> --data research/data/benchmark-qa.jsonl
+
+external:
+  embeddings_mteb:
+    tool: mteb
+    claim: Better embeddings
+    description: MTEB, BEIR, STS — not wired in this repo; use embeddings-benchmark/mteb.
+    tasks_note: 56+ embedding tasks; requires a dedicated encoder model.
+
+  chat_judge:
+    tool: mt_bench_alpacaeval
+    claim: Chat / instruction quality (judge-based)
+    description: MT-Bench, AlpacaEval — require LLM-as-judge APIs; not in slm-lm-eval.
+
+  rl_envs:
+    tool: babyai_minigrid
+    claim: Embodied planning
+    description: BabyAI, MiniGrid — separate RL env evals; not integrated here.

research/evals/configs/experiment_001.yaml

@@ -0,0 +1,53 @@
+# configs/experiment_001.yaml
+# ─────────────────────────────────────────────────────────────────────────────
+# Template config for a local SLM agentic benchmark run.
+# Copy and edit this file for each new experiment.
+# Run with: uv run --package slm-evals slm-benchmark --config research/evals/configs/experiment_001.yaml
+# ─────────────────────────────────────────────────────────────────────────────
+
+# ── Model ─────────────────────────────────────────────────────────────────────
+model_path: "./models/my-small-model-7b"   # local directory or HF Hub ID
+device: "auto"                              # auto | cpu | cuda | cuda:0
+dtype: "bfloat16"                           # float32 | float16 | bfloat16 | int8 | int4
+
+# ── Inference ─────────────────────────────────────────────────────────────────
+max_new_tokens: 512
+temperature: 0.0           # 0 = greedy (recommended for deterministic evals)
+
+# ── Experiment ────────────────────────────────────────────────────────────────
+experiment_name: "my-model-7b__bfcl-tau-gaia-swe__v1"
+output_dir: "results"
+
+# ── Benchmarks ────────────────────────────────────────────────────────────────
+benchmarks:
+  - bfcl
+  - tau_bench
+  - gaia
+  - swe_bench
+
+max_samples: 100    # null = run all; set a number for quick smoke tests
+
+# ── Per-benchmark overrides ───────────────────────────────────────────────────
+benchmark_overrides:
+
+  bfcl:
+    # data_path: "./data/bfcl_v4_local.jsonl"   # uncomment to use local copy
+    strict: false                # false = fuzzy arg matching (recommended for SLMs)
+    categories: []               # [] = all categories
+
+  tau_bench:
+    # data_path: "./data/tau_bench_retail.jsonl"
+    domain: "retail"             # retail | airline | both
+    max_turns: 15
+    use_llm_user: false          # true = use GPT-4o as user simulator (costs $)
+
+  gaia:
+    # data_path: "./data/gaia_validation.jsonl"
+    split: "validation"          # validation (public) | test (needs HF auth)
+    levels: [1, 2]               # [1] easy only, [1,2,3] all levels
+    tool_mode: "describe"        # describe (offline) | none
+
+  swe_bench:
+    # data_path: "./data/swe_bench_verified.jsonl"
+    full_eval: false             # true = run Docker harness (requires swebench pkg)
+    context_lines: 80

research/evals/configs/lm_eval_code.yaml

@@ -0,0 +1,19 @@
+# Code generation profile — HumanEval + MBPP
+# Run: slm-lm-eval --profile code --preset minicpm5-1b-lesson-lora
+# Note: small models often score low; use --limit 25 for smoke tests.
+
+profile: code
+claim: Better code generation
+
+tasks:
+  - humaneval
+  - mbpp
+
+num_fewshot: 0
+limit: 50
+seed: 42
+batch_size: auto
+device: auto
+dtype: bfloat16
+trust_remote_code: true
+output_dir: results/lm_eval

research/evals/configs/lm_eval_compare_study.yaml

@@ -0,0 +1,42 @@
+# configs/lm_eval_compare_study.yaml
+# ─────────────────────────────────────────────────────────────────────────────
+# Shared settings for baseline vs finetuned comparisons.
+# Use the SAME config for both runs; only change --preset / --experiment-name.
+#
+# Baseline:
+#   uv run --package slm-evals slm-lm-eval \
+#     --config research/evals/configs/lm_eval_compare_study.yaml \
+#     --preset minicpm5-1b \
+#     --experiment-name minicpm5-1b__baseline
+#
+# Candidate (after finetune):
+#   uv run --package slm-evals slm-lm-eval \
+#     --config research/evals/configs/lm_eval_compare_study.yaml \
+#     --preset minicpm5-1b-lesson-lora \
+#     --experiment-name minicpm5-1b-lora__v1 \
+#     --compare-to results/lm_eval/minicpm5-1b__baseline/results.json
+# ─────────────────────────────────────────────────────────────────────────────
+
+study:
+  baseline_preset: minicpm5-1b
+  candidate_preset: minicpm5-1b-lesson-lora
+  notes: >
+    Keep tasks, num_fewshot, limit, and seed identical across runs.
+    Do not compare training_results.json result_score to lm-eval accuracy.
+
+tasks:
+  - arc_easy
+  - arc_challenge
+  - hellaswag
+  - piqa
+  - boolq
+  - gsm8k
+
+num_fewshot: 5
+limit: 100
+seed: 42
+batch_size: auto
+device: auto
+dtype: bfloat16
+trust_remote_code: true
+output_dir: results/lm_eval

research/evals/configs/lm_eval_instructions.yaml

@@ -0,0 +1,18 @@
+# Instruction following profile — IFEval (verifiable constraints)
+# Run: slm-lm-eval --profile instructions --preset minicpm5-1b
+# Requires lm-eval extras; install with: uv sync --group lm-eval
+
+profile: instructions
+claim: Better instruction following
+
+tasks:
+  - ifeval
+
+num_fewshot: 0
+limit: 100
+seed: 42
+batch_size: auto
+device: auto
+dtype: bfloat16
+trust_remote_code: true
+output_dir: results/lm_eval

research/evals/configs/lm_eval_minicpm5.yaml

@@ -0,0 +1,26 @@
+# configs/lm_eval_minicpm5.yaml
+# ─────────────────────────────────────────────────────────────────────────────
+# Full academic benchmark profile for ~1B SLMs (MiniCPM5, Gemma E2B, etc.).
+# Run with:
+#   uv run --package slm-evals slm-lm-eval \
+#     --config research/evals/configs/lm_eval_minicpm5.yaml \
+#     --preset minicpm5-1b \
+#     --experiment-name minicpm5-1b__baseline
+# ─────────────────────────────────────────────────────────────────────────────
+
+tasks:
+  - arc_easy
+  - arc_challenge
+  - hellaswag
+  - piqa
+  - boolq
+  - gsm8k
+
+num_fewshot: 5
+limit: null
+seed: 42
+batch_size: auto
+device: auto
+dtype: bfloat16
+trust_remote_code: true
+output_dir: results/lm_eval

research/evals/configs/lm_eval_reasoning.yaml

@@ -0,0 +1,20 @@
+# Reasoning profile — math + commonsense + science QA
+# Run: slm-lm-eval --profile reasoning --preset minicpm5-1b --experiment-name reasoning-baseline
+
+profile: reasoning
+claim: Better reasoning
+
+tasks:
+  - gsm8k
+  - arc_easy
+  - arc_challenge
+  - hellaswag
+
+num_fewshot: 5
+limit: 100
+seed: 42
+batch_size: auto
+device: auto
+dtype: bfloat16
+trust_remote_code: true
+output_dir: results/lm_eval