mindxtrain/governance/proof_loop.py

"""dcoach proof loop — prove a CPU-trained model recalls its training.

Chains the whole thing end-to-end: compose a persona/skills script → derive training
params (nudged by past feedback) → imprint-train a tiny actor on CPU → probe before
(base) vs after (adapter) recall → the **classroom** tests recall/persona → the
**boardroom** decides success/failure → record **autotune feedback** that improves the
next run. This is mindXtrain's first-run proof + the autotune feedback loop.

Heavy (real training + generation) — runs on the `trl_local` CPU lane. `on_event(phase,
msg)` reports progress so a UI can stream it.
"""

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path
from typing import Any

import yaml
from pydantic import BaseModel, ConfigDict

from mindxtrain.governance.classroom import ClassroomReport

_DEFAULT_BASE_MODEL = "HuggingFaceTB/SmolLM2-135M"
_IMPRINT_RECIPE = "mindx_persona_imprint_local"


class ProofResult(BaseModel):
    """The outcome of one proof-loop run."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    run_id: str
    dataset_path: str
    rows: int
    train_params: dict[str, int]
    classroom: ClassroomReport
    boardroom_outcome: str
    boardroom_rationale: str
    passed: bool
    next_params: dict[str, int]


def _build_cfg(base_model: str, script_path: Path, params: dict[str, int], run_id: str) -> Any:
    """Render the imprint recipe, override base/data/params, and load an XTrainConfig."""
    from mindxtrain.config.loader import load_config, render_recipe

    raw = yaml.safe_load(render_recipe(_IMPRINT_RECIPE))
    raw["meta"]["run_name"] = run_id
    raw["model"]["name"] = base_model
    raw["data"]["path"] = str(script_path)
    raw["data"]["max_samples"] = 64
    raw["data"]["seq_len"] = 128
    raw["train"]["schedule"]["epochs"] = int(params.get("epochs", 12))
    raw["train"]["batch"]["grad_accum"] = int(params.get("grad_accum", 1))
    raw["train"]["batch"]["per_device"] = int(params.get("per_device", 1))
    import tempfile

    with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as fh:
        fh.write(yaml.safe_dump(raw))
        tmp = fh.name
    return load_config(tmp)


def run_proof_loop(
    *,
    run_id: str,
    persona: str = "codephreak",
    skills: list[str] | None = None,
    exchanges: list[Any] | None = None,
    base_model: str = _DEFAULT_BASE_MODEL,
    out_dir: str | Path,
    inquiries: list[str] | None = None,
    board_preset: str = "classic_triad",
    board_model: str | None = None,
    force_cpu: bool = True,
    max_new_tokens: int = 32,
    on_event: Callable[[str, str], None] | None = None,
    feedback_path: Path | None = None,
) -> ProofResult:
    """Run the full proof loop and return a structured `ProofResult`."""
    emit = on_event or (lambda _phase, _msg: None)
    out = Path(out_dir)

    # 1) Compose persona + skills → script.
    from mindxtrain.data import personas as _pz
    from mindxtrain.data.scripts import (
        build_script_rows,
        derive_training_params,
        write_script_jsonl,
    )

    base_persona, skill_exchanges = _pz.compose(persona, skills or [])
    all_exchanges = list(exchanges or []) + skill_exchanges
    rows_list = build_script_rows(base_persona, all_exchanges, seed_voice=True)
    script_path = write_script_jsonl(rows_list, out / "script.jsonl")
    rows = len(rows_list)
    emit("dataset", f"authored {rows} rows for persona '{base_persona.name}'")

    # 2) Derive params, nudged by past feedback.
    from mindxtrain.autotune import feedback as _fb

    params = _fb.suggest_from_history(derive_training_params(rows), path=feedback_path)
    emit("params", f"epochs={params['epochs']} grad_accum={params['grad_accum']}")

    # 3) Imprint-train the tiny actor.
    from mindxtrain.autotune.benchmark import run_autotune
    from mindxtrain.train.backend_trl_cpu import run_trl_local

    cfg = _build_cfg(base_model, script_path, params, run_id)
    run_dir = out / "run"
    run_dir.mkdir(parents=True, exist_ok=True)
    emit("train", "imprinting the persona…")
    run_trl_local(cfg, run_autotune(dry_run=True), run_dir, force_cpu=force_cpu)
    adapter = run_dir / "checkpoint"

    # 4) Probe before (base) vs after (adapter) recall.
    from mindxtrain.data.scripts import persona_system_prompt
    from mindxtrain.eval.imprint import default_inquiries, probe_recall

    inq = inquiries or [e.user for e in all_exchanges][:4] or default_inquiries(base_persona.name)
    baseline = [e.assistant for e in all_exchanges] or list(base_persona.voice_examples)
    system = persona_system_prompt(base_persona)  # match the conditioning the adapter trained under
    emit("probe", "recall before training…")
    before = probe_recall(
        base_model, inq, system=system, force_cpu=force_cpu, max_new_tokens=max_new_tokens,
    )
    emit("probe", "recall after training…")
    after = probe_recall(
        base_model, inq, adapter_dir=adapter, system=system,
        force_cpu=force_cpu, max_new_tokens=max_new_tokens,
    )

    # 5) Classroom test.
    from mindxtrain.governance.classroom import evaluate_classroom, graduate

    classroom = evaluate_classroom(inq, before, after, baseline)
    emit("classroom", f"passed={classroom.passed} recall {classroom.before_recall}→{classroom.recall}")

    # 6) Boardroom decision.
    from mindxtrain.eval.imprint import score_imprint
    from mindxtrain.governance import Boardroom
    from mindxtrain.governance.boardroom import board_from_preset

    grad = graduate(score_imprint(inq, before, after, baseline), run_id=run_id)
    board = Boardroom(members=board_from_preset(board_preset, model=board_model or ""))
    if board_model:
        from mindxtrain.governance import panel as _panel

        ballot: Any = _panel.model_ballot(default_model=board_model)
    else:
        vote = "approve" if classroom.passed else "reject"
        ballot = {m.id: vote for m in board.members}
    decision = board.convene(grad.motion, ballot)
    emit("boardroom", f"{decision.outcome}: {decision.rationale}")

    # 7) Record feedback + suggest the next run's params.
    outcome = decision.outcome
    _fb.record(
        run_id=run_id, params=params, classroom_score=classroom.imprint_delta,
        passed=classroom.passed, boardroom_outcome=outcome, path=feedback_path,  # type: ignore[arg-type]
    )
    next_params = _fb.suggest_next_params(
        params, passed=classroom.passed, classroom_score=classroom.imprint_delta,
    )
    emit("feedback", f"next: epochs={next_params['epochs']} grad_accum={next_params['grad_accum']}")

    return ProofResult(
        run_id=run_id, dataset_path=str(script_path), rows=rows, train_params=params,
        classroom=classroom, boardroom_outcome=outcome, boardroom_rationale=decision.rationale,
        passed=classroom.passed and decision.outcome == "approved", next_params=next_params,
    )


__all__ = ["ProofResult", "run_proof_loop"]