diff --git a/README.md b/README.md index a3b93c2..e9212dc 100644 --- a/README.md +++ b/README.md @@ -300,12 +300,14 @@ Nothing here is needed for a draft. These commands stay in the CLI for local exp ```bash personality-protect build-writer-sft +personality-protect select-writer-holdouts --apply +personality-protect index-voice --from-carve personality-protect train --writer --backend mlx -personality-protect eval-write-holdout --out receipt.json +personality-protect eval-writer-adapter --archive-on-fail personality-protect write --adapter --topic "…" --points "…" ``` -Keep an adapter only if `eval-write-holdout` shows it beating RAG-alone on held-out pieces. Otherwise delete it and stay on the default. Training is not a prerequisite for `write`, and an untested adapter is not an upgrade. +`build-writer-sft` builds de-voiced brief→post pairs. `train --writer` uses a short writer recipe (3 epochs) and keeps per-chunk checkpoints under `adapters/latest/checkpoints/`. Keep an adapter only when `eval-writer-adapter` decides `keep`; otherwise archive it and stay on `adapter=none`. Training is not a prerequisite for `write`. ### Other experiment commands diff --git a/src/personality_protect/cli.py b/src/personality_protect/cli.py index d21c28f..ad35876 100644 --- a/src/personality_protect/cli.py +++ b/src/personality_protect/cli.py @@ -941,6 +941,12 @@ def eval_writer_adapter_cmd( "--archive-on-fail", help="Move the adapter aside when the gate fails (write returns to adapter=none).", ), + sweep_checkpoints: bool = typer.Option( + False, + "--sweep-checkpoints", + help="Gate every adapters/latest/checkpoints/step_* dir, earliest first; " + "install the first that ships.", + ), out: Optional[Path] = typer.Option(None, "--out", help="Write the receipt JSON here."), profile: str = typer.Option(DEFAULT_PROFILE, "--profile"), home: Optional[Path] = typer.Option(None, "--home"), @@ -952,6 +958,7 @@ def eval_writer_adapter_cmd( device — importing MLX without one aborts the interpreter. """ from personality_protect.eval_writer_adapter import ( + run_checkpoint_gate_sweep, run_writer_adapter_gate, write_gate_receipt, ) @@ -978,6 +985,63 @@ def eval_writer_adapter_cmd( ) raise typer.Exit(1) + if sweep_checkpoints: + try: + generate_rag = make_mlx_generator(base_model=config.base_model) + except RuntimeError as exc: + console.print(f"[red]{exc}[/red]") + raise typer.Exit(1) from exc + + def _make_adapter(adapter_path: str): + return make_mlx_generator( + base_model=config.base_model, adapter_path=adapter_path + ) + + def _on_ckpt(row: dict, _receipt: dict) -> None: + if as_json: + return + wins = row["wins"] + console.print( + f"[dim]{row['checkpoint']}: adapter {wins['adapter']} — " + f"rag {wins['rag']} — tie {wins['tie']} → {row['decision']}[/dim]" + ) + + try: + sweep = run_checkpoint_gate_sweep( + paths, + holdout_ids, + make_adapter_generate=_make_adapter, + generate_fn_rag=generate_rag, + k=k, + max_tokens=max_tokens, + alpha=alpha, + on_checkpoint=None if as_json else _on_ckpt, + ) + except (ValueError, FileNotFoundError) as exc: + console.print(f"[red]{exc}[/red]") + raise typer.Exit(1) from exc + + if sweep["decision"] == "archive" and archive_on_fail: + sweep["archived_to"] = archive_writer_adapter( + paths, reason="gate-fail-sweep" + ) + + target = out or ( + paths.root / "dogfood" / "writer_adapter_checkpoint_sweep_receipt.json" + ) + write_gate_receipt(sweep, target) + if as_json: + typer.echo(json.dumps(sweep, indent=2, ensure_ascii=False)) + else: + console.print( + f"sweep: evaluated {sweep['evaluated']}/{sweep['n_checkpoints']} " + f"→ decision [bold]{sweep['decision']}[/bold] " + f"(kept={sweep['kept_checkpoint']}) → {target}" + ) + if sweep["decision"] != "keep": + raise typer.Exit(1) + return + adapter_path = resolve_writer_adapter(paths) if adapter_path is None: console.print( diff --git a/src/personality_protect/eval_writer_adapter.py b/src/personality_protect/eval_writer_adapter.py index 0a6b4d1..593d8d7 100644 --- a/src/personality_protect/eval_writer_adapter.py +++ b/src/personality_protect/eval_writer_adapter.py @@ -253,3 +253,74 @@ def write_gate_receipt(receipt: dict[str, Any], path: Any) -> Any: json.dumps(receipt, indent=2, ensure_ascii=False) + "\n", encoding="utf-8" ) return path + + +def run_checkpoint_gate_sweep( + paths: ProfilePaths, + holdout_ids: Sequence[str], + *, + make_adapter_generate: Any, + generate_fn_rag: GenerateFn, + k: int = DEFAULT_WRITE_K, + max_tokens: int = DEFAULT_WRITE_MAX_TOKENS, + alpha: float = SHIP_ALPHA, + on_checkpoint: Any = None, +) -> dict[str, Any]: + """Gate every durable step checkpoint, earliest first; keep the first that ships. + + ``make_adapter_generate(adapter_path)`` must build a fresh generator for the + installed weights — reusing a prior MLX load would score the wrong adapter. + """ + from personality_protect.mlx_train import list_step_checkpoints + from personality_protect.write import install_writer_checkpoint + + latest = paths.adapters_dir / "latest" + checkpoints = list_step_checkpoints(latest) + if not checkpoints: + raise FileNotFoundError( + f"No step checkpoints under {latest / 'checkpoints'}. " + "Retrain so each chunk persists checkpoints/step_NNNNNN/." + ) + + results: list[dict[str, Any]] = [] + kept: str | None = None + for ckpt in checkpoints: + install_writer_checkpoint(paths, ckpt) + generate_adapter = make_adapter_generate(str(ckpt)) + receipt = run_writer_adapter_gate( + paths, + holdout_ids, + generate_fn_adapter=generate_adapter, + generate_fn_rag=generate_fn_rag, + k=k, + max_tokens=max_tokens, + alpha=alpha, + ) + row = { + "checkpoint": ckpt.name, + "decision": receipt["decision"], + "wins": receipt["wins"], + "disqualified": receipt["disqualified"], + "p_value": receipt["p_value"], + "blocking_reasons": receipt["blocking_reasons"], + "n_holdouts": receipt["n_holdouts"], + } + results.append(row) + if on_checkpoint is not None: + on_checkpoint(row, receipt) + if receipt["decision"] == "keep": + kept = ckpt.name + install_writer_checkpoint(paths, ckpt) + break + + sweep: dict[str, Any] = { + "kind": "eval_writer_adapter_checkpoint_sweep", + "created_at": datetime.now(timezone.utc).isoformat(), + "n_checkpoints": len(checkpoints), + "evaluated": len(results), + "kept_checkpoint": kept, + "decision": "keep" if kept else "archive", + "results": results, + } + assert_receipt_contoso_safe(sweep) + return sweep diff --git a/src/personality_protect/mlx_train.py b/src/personality_protect/mlx_train.py index f690ce0..d2190a9 100644 --- a/src/personality_protect/mlx_train.py +++ b/src/personality_protect/mlx_train.py @@ -42,7 +42,10 @@ WRITER_NUM_LAYERS = 16 WRITER_LORA_RANK = 16 WRITER_LEARNING_RATE = 3e-5 -WRITER_EPOCHS = 10 +# 10 epochs on ~60 de-voiced pairs drove train loss ~0.08 and raised invention / +# parroting on the n=20 gate. Three passes is enough to move the adapter without +# memorizing the tiny set; mid-train step snapshots let a later gate pick earlier. +WRITER_EPOCHS = 3 # Cap wired Metal memory: leave OS/apps breathing room. DEFAULT_WIRED_FRACTION = 0.40 DEFAULT_WIRED_CAP_BYTES = 16 * 10**9 # 16 GB hard cap (leave Studio headroom) @@ -51,6 +54,10 @@ ProgressCallback = Callable[[dict[str, Any]], None] CHECKPOINT_META_NAME = "train_chunks.json" +# Durable per-chunk copies under adapter_dir/checkpoints/step_NNNNNN/. +# Distinct from mlx-lm's ephemeral ``0000050_adapters.safetensors`` which each +# chunk overwrites with the same name. +STEP_CHECKPOINTS_DIRNAME = "checkpoints" # Snapshot before each chunk so nan / crash never leaves a wiped or poisoned adapter. LAST_GOOD_ADAPTER_NAME = "adapters.safetensors.last_good" @@ -79,6 +86,47 @@ def restore_last_good_adapter(adapter_dir: Path) -> bool: return True +def persist_step_checkpoint(adapter_dir: Path, completed_steps: int) -> Path | None: + """Copy live weights into ``checkpoints/step_NNNNNN/`` after a good chunk. + + mlx-lm's own numbered files reuse the same basename every chunk, so earlier + steps disappear. These directories keep every completed step count so a gate + can evaluate under-trained adapters without a full retrain. + """ + src = adapter_dir / "adapters.safetensors" + if not src.is_file(): + return None + steps = max(0, int(completed_steps)) + dest_dir = adapter_dir / STEP_CHECKPOINTS_DIRNAME / f"step_{steps:06d}" + dest_dir.mkdir(parents=True, exist_ok=True) + shutil.copy2(src, dest_dir / "adapters.safetensors") + for name in ("adapter_config.json", "adapter_config.yaml"): + cfg = adapter_dir / name + if cfg.is_file(): + shutil.copy2(cfg, dest_dir / name) + return dest_dir + + +def list_step_checkpoints(adapter_dir: Path) -> list[Path]: + """Return step checkpoint dirs oldest-first (under-trained → final).""" + root = adapter_dir / STEP_CHECKPOINTS_DIRNAME + if not root.is_dir(): + return [] + dirs = [ + path + for path in root.iterdir() + if path.is_dir() and (path / "adapters.safetensors").is_file() + ] + return sorted(dirs, key=lambda path: path.name) + + +def clear_step_checkpoints(adapter_dir: Path) -> None: + """Drop durable step snapshots (used on ``--force-retrain``).""" + root = adapter_dir / STEP_CHECKPOINTS_DIRNAME + if root.is_dir(): + shutil.rmtree(root) + + def plan_train_chunks(total_steps: int, chunk_size: int) -> list[int]: """Split ``total_steps`` into positive chunk sizes (last chunk may be shorter).""" total = max(0, int(total_steps)) @@ -164,9 +212,13 @@ def _clear_adapter_weights(adapter_dir: Path) -> None: adapter_file.unlink() for stale in adapter_dir.glob("*_adapters.safetensors"): stale.unlink() + last_good = adapter_dir / LAST_GOOD_ADAPTER_NAME + if last_good.is_file(): + last_good.unlink() meta_path = adapter_dir / CHECKPOINT_META_NAME if meta_path.is_file(): meta_path.unlink() + clear_step_checkpoints(adapter_dir) def resolve_train_plan( @@ -697,6 +749,7 @@ def _on_line(line: str, *, _completed=completed, _n=n_iters) -> None: completed += n_iters status = "complete" if completed >= plan.total_steps else "in_progress" + step_ckpt = persist_step_checkpoint(adapter_dir, completed) chunk_meta = { "status": status, "completed_steps": completed, @@ -715,6 +768,7 @@ def _on_line(line: str, *, _completed=completed, _n=n_iters) -> None: "already_completed": plan.already_completed, "steps_this_run": completed - plan.already_completed, "chunks": total_chunk_count, + "step_checkpoint": str(step_ckpt) if step_ckpt is not None else None, } write_train_checkpoint_meta(adapter_dir, chunk_meta) @@ -727,6 +781,7 @@ def _on_line(line: str, *, _completed=completed, _n=n_iters) -> None: "completed_steps": completed, "total_steps": plan.total_steps, "peak_mem_gb": result.peak_mem_gb, + "step_checkpoint": str(step_ckpt) if step_ckpt is not None else None, } ) diff --git a/src/personality_protect/write.py b/src/personality_protect/write.py index b4509d6..a21d181 100644 --- a/src/personality_protect/write.py +++ b/src/personality_protect/write.py @@ -8,6 +8,7 @@ import re from collections.abc import Callable, MutableSequence, Sequence +from pathlib import Path from typing import Any from personality_protect.chat_prompt import ( @@ -176,7 +177,31 @@ def archive_writer_adapter(paths: ProfilePaths, *, reason: str) -> str | None: stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") target = paths.adapters_dir / f"writer-{reason}-{stamp}" latest.rename(target) - return str(target) + # Basename only — absolute paths under ~/.personality-protect embed the + # local username and poison Contoso-safe receipts. + return target.name + + +def install_writer_checkpoint(paths: ProfilePaths, checkpoint_dir: str | Path) -> str: + """Copy a step checkpoint into ``adapters/latest`` for gating or shipping. + + Leaves durable ``checkpoints/`` snapshots in place. Overwrites only the live + ``adapters.safetensors`` (and config) that :func:`resolve_writer_adapter` reads. + """ + import shutil + + src = Path(checkpoint_dir) + weights = src / "adapters.safetensors" + if not weights.is_file(): + raise FileNotFoundError(f"No adapters.safetensors in checkpoint {src}") + latest = paths.adapters_dir / "latest" + latest.mkdir(parents=True, exist_ok=True) + shutil.copy2(weights, latest / "adapters.safetensors") + for name in ("adapter_config.json", "adapter_config.yaml"): + cfg = src / name + if cfg.is_file(): + shutil.copy2(cfg, latest / name) + return str(latest) _SENTENCE_START_WORD_RE = re.compile( diff --git a/tests/test_eval_writer_adapter.py b/tests/test_eval_writer_adapter.py index d47b28b..fe5f4fd 100644 --- a/tests/test_eval_writer_adapter.py +++ b/tests/test_eval_writer_adapter.py @@ -162,3 +162,48 @@ def test_gate_refuses_a_holdout_that_leaked_into_retrieval(tmp_path: Path, monke generate_fn_rag=_fixed(RAG_DRAFT), k=0, ) + + +def test_checkpoint_sweep_keeps_the_first_shipping_step(tmp_path: Path): + from personality_protect.eval_writer_adapter import run_checkpoint_gate_sweep + from personality_protect.mlx_train import persist_step_checkpoint + + paths = _profile(tmp_path) + latest = paths.adapters_dir / "latest" + (latest / "adapters.safetensors").write_bytes(b"early") + persist_step_checkpoint(latest, 50) + (latest / "adapters.safetensors").write_bytes(b"late") + persist_step_checkpoint(latest, 100) + + calls: list[str] = [] + + def _make(adapter_path: str): + calls.append(Path(adapter_path).name) + # Early checkpoint wins the distance game; late one mirrors the RAG loser. + if Path(adapter_path).name == "step_000050": + return _fixed(ADAPTER_DRAFT) + return _fixed(RAG_DRAFT) + + sweep = run_checkpoint_gate_sweep( + paths, + ["hold1", "hold2"], + make_adapter_generate=_make, + generate_fn_rag=_fixed(RAG_DRAFT), + k=0, + alpha=1.0, # majority alone is enough for Contoso stub n=2 + ) + assert sweep["decision"] == "keep" + assert sweep["kept_checkpoint"] == "step_000050" + assert sweep["evaluated"] == 1 # stop at first keep + assert calls == ["step_000050"] + assert (latest / "adapters.safetensors").read_bytes() == b"early" + assert_receipt_contoso_safe(sweep) + + +def test_writer_epochs_default_is_short_enough_to_avoid_overfit(): + from personality_protect.mlx_train import WRITER_EPOCHS + from personality_protect.train import auto_max_steps, writer_train_settings + + assert WRITER_EPOCHS == 3 + assert writer_train_settings()["epochs"] == 3 + assert auto_max_steps(60, epochs=WRITER_EPOCHS) == 180 diff --git a/tests/test_mlx_train_safety.py b/tests/test_mlx_train_safety.py index 63f3d9f..2f6b326 100644 --- a/tests/test_mlx_train_safety.py +++ b/tests/test_mlx_train_safety.py @@ -570,6 +570,37 @@ def _side_effect(**kwargs): assert metas_during[1] is not None assert metas_during[1]["completed_steps"] == 50 assert metas_during[1]["status"] == "in_progress" + from personality_protect.mlx_train import list_step_checkpoints + + steps = [p.name for p in list_step_checkpoints(adapter_dir)] + assert steps == ["step_000050", "step_000100", "step_000120"] + assert (adapter_dir / "checkpoints" / "step_000050" / "adapters.safetensors").is_file() + + +def test_persist_step_checkpoint_roundtrip(tmp_path: Path): + from personality_protect.mlx_train import ( + clear_step_checkpoints, + list_step_checkpoints, + persist_step_checkpoint, + ) + + adapter_dir = tmp_path / "adapters" + adapter_dir.mkdir() + (adapter_dir / "adapters.safetensors").write_bytes(b"w50") + (adapter_dir / "adapter_config.json").write_text("{}", encoding="utf-8") + dest = persist_step_checkpoint(adapter_dir, 50) + assert dest is not None + assert dest.name == "step_000050" + assert (dest / "adapters.safetensors").read_bytes() == b"w50" + assert (dest / "adapter_config.json").read_text(encoding="utf-8") == "{}" + (adapter_dir / "adapters.safetensors").write_bytes(b"w100") + persist_step_checkpoint(adapter_dir, 100) + assert [p.name for p in list_step_checkpoints(adapter_dir)] == [ + "step_000050", + "step_000100", + ] + clear_step_checkpoints(adapter_dir) + assert list_step_checkpoints(adapter_dir) == [] def test_chunked_train_resume_skips_completed_steps(tmp_path: Path):