mlxsmith 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

mlxsmith/train/self_verify.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+from typing import Iterable, Optional
+
+from rich.console import Console
+
+from ..accel import get_backend
+from ..config import ProjectConfig
+from ..models import resolve_model_spec
+from ..runs import RunPaths, new_run, snapshot_config
+from ..util import write_jsonl, now_ts, tree_add, tree_scale, clip_grad_norm
+from ..llm.registry import get_llm_backend
+from ..llm.backend import BackendNotAvailable
+from ..verifiers.llm_judge import verify as judge_verify
+from .lora import LoRAConfig
+
+console = Console()
+
+
+def _iter_prompts(path: Path) -> Iterable[str]:
+    for line in path.read_text(encoding="utf-8").splitlines():
+        if not line.strip():
+            continue
+        row = json.loads(line)
+        prompt = row.get("prompt") or row.get("instruction") or row.get("input") or row.get("question") or ""
+        if not prompt and "messages" in row:
+            msgs = row.get("messages") or []
+            if msgs:
+                prompt = "\n".join([m.get("content", "") for m in msgs])
+        if prompt:
+            yield str(prompt)
+
+
+def run_self_verify(
+    project_root: Path,
+    cfg: ProjectConfig,
+    data_path: Path,
+    model_id_or_path: str,
+    accel: str,
+    *,
+    verifier_model: Optional[str] = None,
+    verifier_backend: str = "mlx-lm",
+    rubric: Optional[str] = None,
+    max_new_tokens: Optional[int] = None,
+    temperature: Optional[float] = None,
+    judge_mock_response: Optional[str | list[str]] = None,
+) -> RunPaths:
+    run = new_run(project_root, "self_verify")
+    snapshot_config(cfg.model_dump(), run.config_snapshot_path)
+
+    prompts = list(_iter_prompts(data_path))
+    if not prompts:
+        raise RuntimeError("No prompts found in self-verify dataset")
+
+    backend = get_backend(accel)
+    backend.patch()
+    console.print(f"[bold]SELF-VERIFY[/bold] run: {run.run_dir.name} accel={backend.name}")
+
+    policy = get_llm_backend(cfg.model.backend)
+    base_model, adapter_path, _meta = resolve_model_spec(project_root, model_id_or_path, cfg)
+
+    try:
+        policy.load(
+            base_model,
+            max_seq_len=cfg.model.max_seq_len,
+            dtype=cfg.model.dtype,
+            trust_remote_code=cfg.model.trust_remote_code,
+        )
+        if adapter_path:
+            policy.apply_adapter(str(adapter_path))
+        else:
+            lora_cfg = LoRAConfig(
+                r=cfg.lora.r,
+                alpha=cfg.lora.alpha,
+                dropout=cfg.lora.dropout,
+                target_modules=list(cfg.lora.target_modules or []),
+                num_layers=cfg.lora.num_layers,
+                scale=cfg.lora.scale,
+                fine_tune_type=cfg.lora.fine_tune_type,
+            )
+            policy.apply_lora_from_config(lora_cfg)
+    except BackendNotAvailable as e:
+        console.print(f"[yellow]MLX backend unavailable[/yellow]: {e}")
+        (run.adapter_dir / "ADAPTER.txt").write_text(
+            f"Backend unavailable in this environment.\nmodel={model_id_or_path}\naccel={backend.name}\n",
+            encoding="utf-8",
+        )
+        return run
+
+    opt, _params = policy.optimizer_and_params(
+        lr=cfg.train.lr,
+        weight_decay=cfg.train.weight_decay,
+        optimizer=cfg.train.optimizer,
+        optimizer_kwargs=cfg.train.optimizer_kwargs,
+    )
+
+    total = int(cfg.train.iters)
+    grad_accum = max(1, int(cfg.train.grad_accum))
+    max_grad_norm = float(getattr(cfg.train, "max_grad_norm", 0.0))
+    max_new = int(max_new_tokens or cfg.rft.max_new_tokens)
+    temp = float(temperature if temperature is not None else cfg.rft.temperature)
+
+    rng = random.Random(cfg.train.seed)
+    accum_grads = None
+    accum_loss = 0.0
+    accum_count = 0
+    reward_ema = 0.0
+    ema_alpha = 0.1
+
+    def _next_mock(idx: int) -> Optional[str]:
+        if judge_mock_response is None:
+            return None
+        if isinstance(judge_mock_response, list):
+            if not judge_mock_response:
+                return None
+            return judge_mock_response[min(idx, len(judge_mock_response) - 1)]
+        return judge_mock_response
+
+    for step in range(1, total + 1):
+        prompt = rng.choice(prompts)
+        gen = policy.generate_with_logprobs(
+            prompt,
+            max_new_tokens=max_new,
+            temperature=temp,
+            seed=rng.randint(0, 2**31 - 1),
+            logprobs=0,
+        )
+        completion = gen.text[len(prompt) :] if gen.text.startswith(prompt) else gen.text
+
+        res = judge_verify(
+            prompt,
+            completion,
+            str(run.artifacts_dir),
+            model=verifier_model or model_id_or_path,
+            backend=verifier_backend,
+            rubric=rubric,
+            reward_mode="score",
+            mock_response=_next_mock(0),
+        )
+        reward = float(getattr(res, "reward", 0.0))
+        reward_ema = (1.0 - ema_alpha) * reward_ema + ema_alpha * reward
+        advantage = reward - reward_ema
+
+        token_ids = list(gen.token_ids)
+        prompt_len = int(gen.prompt_len)
+
+        def loss_fn(_model):
+            logp = policy.sequence_logprob(token_ids, prompt_len=prompt_len)
+            return -policy.mx.array(float(advantage)) * logp  # type: ignore
+
+        lval, grads = policy.value_and_grad(loss_fn)
+        accum_loss += float(lval.item()) if hasattr(lval, "item") else float(lval)
+        accum_count += 1
+        if grads is not None:
+            accum_grads = tree_add(accum_grads, grads)
+
+        if step % grad_accum == 0:
+            if accum_grads is not None:
+                scaled = tree_scale(accum_grads, 1.0 / grad_accum)
+                if max_grad_norm > 0:
+                    scaled = clip_grad_norm(scaled, max_grad_norm)
+                policy.apply_grads(opt, scaled)
+            accum_grads = None
+            accum_loss = 0.0
+            accum_count = 0
+
+        if step % cfg.train.log_every == 0 or step == 1 or step == total:
+            avg_loss = accum_loss / max(1, accum_count) if accum_count else float(lval)
+            write_jsonl(
+                run.metrics_path,
+                [
+                    {
+                        "ts": now_ts(),
+                        "step": step,
+                        "kind": "self_verify",
+                        "loss": avg_loss,
+                        "reward": reward,
+                        "advantage": advantage,
+                        "accel": backend.name,
+                    }
+                ],
+            )
+
+        if step % cfg.train.save_every == 0 or step == total:
+            policy.save_adapter(
+                str(run.adapter_dir),
+                metadata={
+                    "base_model": base_model,
+                    "source_adapter": str(adapter_path) if adapter_path else None,
+                    "run": run.run_dir.name,
+                    "kind": "self_verify",
+                },
+            )
+
+    console.print(f"[green]Saved adapter[/green] {run.adapter_dir}")
+    return run

mlxsmith/train/sft.py

@@ -79,7 +79,12 @@ def run_sft(project_root: Path, cfg: ProjectConfig, data_dir: Path, model_id_or_
         )
         return run
 
-    opt, _params = llm.optimizer_and_params(lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
+    opt, _params = llm.optimizer_and_params(
+        lr=cfg.train.lr,
+        weight_decay=cfg.train.weight_decay,
+        optimizer=cfg.train.optimizer,
+        optimizer_kwargs=cfg.train.optimizer_kwargs,
+    )
 
     total = int(cfg.train.iters)
     grad_accum = max(1, int(cfg.train.grad_accum))
@@ -89,6 +94,7 @@ def run_sft(project_root: Path, cfg: ProjectConfig, data_dir: Path, model_id_or_
     rng = random.Random(cfg.train.seed)
     accum_grads = None
     accum_loss = 0.0
+    accum_count = 0
 
     for step in range(1, total + 1):
         row = rng.choice(rows)
@@ -110,6 +116,7 @@ def run_sft(project_root: Path, cfg: ProjectConfig, data_dir: Path, model_id_or_
 
         lval, grads = llm.value_and_grad(loss_fn)
         accum_loss += float(lval.item()) if hasattr(lval, "item") else float(lval)
+        accum_count += 1
         if grads is not None:
             accum_grads = tree_add(accum_grads, grads)
 
@@ -121,8 +128,12 @@ def run_sft(project_root: Path, cfg: ProjectConfig, data_dir: Path, model_id_or_
                 llm.apply_grads(opt, scaled)
             accum_grads = None
             accum_loss = 0.0
+            accum_count = 0
 
         if step % cfg.train.log_every == 0 or step == 1 or step == total:
+            avg_loss = (accum_loss / max(1, accum_count)) if accum_count else (
+                float(lval.item()) if hasattr(lval, "item") else float(lval)
+            )
             write_jsonl(
                 run.metrics_path,
                 [
@@ -130,7 +141,7 @@ def run_sft(project_root: Path, cfg: ProjectConfig, data_dir: Path, model_id_or_
                         "ts": now_ts(),
                         "step": step,
                         "kind": "sft",
-                        "loss": float(lval.item()) if hasattr(lval, "item") else float(lval),
+                        "loss": avg_loss,
                         "accel": backend.name,
                     }
                 ],

mlxsmith/verifiers/llm_judge.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+import time
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from .types import VerifyResult
+from ..llm.registry import get_llm_backend
+
+_STATE: Dict[str, Any] = {
+    "backend": None,
+    "backend_name": None,
+    "model_id": None,
+}
+
+
+def _read_text(value: Optional[str]) -> Optional[str]:
+    if not value:
+        return None
+    if value.startswith("@"):
+        path = Path(value[1:])
+        if path.exists():
+            return path.read_text(encoding="utf-8")
+    path = Path(value)
+    if path.exists():
+        return path.read_text(encoding="utf-8")
+    return value
+
+
+def _load_backend(model_id: str, backend_name: str, *, max_seq_len: Optional[int], dtype: Optional[str], trust_remote_code: bool) -> Any:
+    if (
+        _STATE["backend"] is None
+        or _STATE["backend_name"] != backend_name
+        or _STATE["model_id"] != model_id
+    ):
+        backend = get_llm_backend(backend_name)
+        backend.load(
+            model_id,
+            max_seq_len=max_seq_len,
+            dtype=dtype,
+            trust_remote_code=trust_remote_code,
+        )
+        _STATE["backend"] = backend
+        _STATE["backend_name"] = backend_name
+        _STATE["model_id"] = model_id
+    return _STATE["backend"]
+
+
+def _extract_json(text: str) -> Optional[dict]:
+    if not text:
+        return None
+    start = text.find("{")
+    end = text.rfind("}")
+    if start == -1 or end == -1 or end <= start:
+        return None
+    snippet = text[start : end + 1].strip()
+    try:
+        return json.loads(snippet)
+    except json.JSONDecodeError:
+        cleaned = re.sub(r",\s*}", "}", snippet)
+        cleaned = re.sub(r",\s*]", "]", cleaned)
+        cleaned = cleaned.replace("'", "\"")
+        try:
+            return json.loads(cleaned)
+        except json.JSONDecodeError:
+            return None
+
+
+def _coerce_float(val: Any) -> Optional[float]:
+    if val is None:
+        return None
+    try:
+        return float(val)
+    except (TypeError, ValueError):
+        return None
+
+
+def _aggregate_scores(scores: list[float], mode: str) -> Optional[float]:
+    if not scores:
+        return None
+    mode = (mode or "product").lower()
+    if mode == "min":
+        return min(scores)
+    if mode == "mean":
+        return sum(scores) / float(len(scores))
+    prod = 1.0
+    for s in scores:
+        prod *= s
+    return prod
+
+
+def _score_step(
+    judge,
+    *,
+    system_prompt: str,
+    step_text: str,
+    prompt: str,
+    completion: str,
+    rubric_text: str,
+    temperature: float,
+    max_new_tokens: int,
+) -> Optional[float]:
+    step_prompt = (
+        f"{system_prompt}\n\n"
+        "Score this single step from a solution.\n\n"
+        f"## Task\n{prompt}\n\n"
+        f"## Model Answer\n{completion}\n\n"
+        f"## Step\n{step_text}\n\n"
+        f"## Rubric\n{rubric_text}\n\n"
+        "Return JSON only."
+    )
+    gen = judge.generate(
+        step_prompt,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=1.0,
+        top_k=None,
+    )
+    raw = gen.text[len(step_prompt) :] if gen.text.startswith(step_prompt) else gen.text
+    parsed = _extract_json(raw) or {}
+    return _coerce_float(parsed.get("score"))
+
+
+def verify(
+    prompt: str,
+    completion: str,
+    workdir: str,
+    *,
+    model: Optional[str] = None,
+    backend: str = "mlx-lm",
+    system_prompt: Optional[str] = None,
+    rubric: Optional[str] = None,
+    mode: str = "judge",
+    temperature: float = 0.0,
+    max_new_tokens: int = 256,
+    min_score: float = 0.5,
+    reward_pass: float = 1.0,
+    reward_fail: float = 0.0,
+    reward_mode: str = "score",
+    max_seq_len: Optional[int] = None,
+    dtype: Optional[str] = None,
+    trust_remote_code: bool = False,
+    mock_response: Optional[str] = None,
+    process_agg: str = "product",
+    max_steps: int = 8,
+    **kwargs,
+) -> VerifyResult:
+    """LLM-based verifier with JSON output.
+
+    The judge should return JSON: {"passed": bool, "score": 0..1, "reason": "..."}.
+    Set mock_response to bypass backend loading (useful for tests).
+    """
+    model_id = model or os.environ.get("MLXSMITH_JUDGE_MODEL")
+    if not model_id and not mock_response:
+        raise RuntimeError("llm_judge requires `model` or MLXSMITH_JUDGE_MODEL")
+
+    rubric_text = _read_text(rubric) or "Assess correctness and completeness."
+    mode = (mode or "judge").strip().lower()
+    sys_prompt = system_prompt or (
+        "You are a strict verifier. Return ONLY JSON with keys: "
+        "passed (bool), score (0-1), reason (string)."
+    )
+
+    if mode == "thinkprm":
+        sys_prompt = system_prompt or (
+            "You are a process reward model. Evaluate the reasoning quality. "
+            "Return ONLY JSON with keys: passed (bool), score (0-1), reason (string), steps (array)."
+        )
+
+    user_prompt = (
+        "## Task\n"
+        f"{prompt}\n\n"
+        "## Model Answer\n"
+        f"{completion}\n\n"
+        "## Rubric\n"
+        f"{rubric_text}\n\n"
+        "Return JSON only."
+    )
+
+    t0 = time.time()
+    if mock_response is not None:
+        raw = str(mock_response)
+    else:
+        judge = _load_backend(
+            model_id,
+            backend,
+            max_seq_len=max_seq_len,
+            dtype=dtype,
+            trust_remote_code=trust_remote_code,
+        )
+        full_prompt = f"{sys_prompt}\n\n{user_prompt}"
+        gen = judge.generate(
+            full_prompt,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=1.0,
+            top_k=None,
+        )
+        raw = gen.text[len(full_prompt) :] if gen.text.startswith(full_prompt) else gen.text
+
+    parsed = _extract_json(raw) or {}
+    score = _coerce_float(parsed.get("score"))
+    passed_val = parsed.get("passed")
+    steps_raw = parsed.get("steps") if isinstance(parsed, dict) else None
+    step_texts: list[str] = []
+    step_scores: list[float] = []
+    if mode == "thinkprm" and isinstance(steps_raw, list):
+        for step in steps_raw[: max(1, int(max_steps))]:
+            if isinstance(step, dict):
+                text = step.get("text") or step.get("step") or step.get("content") or ""
+                if text:
+                    step_texts.append(str(text))
+                s_val = _coerce_float(step.get("score"))
+                if s_val is not None:
+                    step_scores.append(float(s_val))
+            elif isinstance(step, str):
+                step_texts.append(step)
+
+        if step_texts and (len(step_scores) < len(step_texts)) and mock_response is None:
+            judge = _load_backend(
+                model_id,
+                backend,
+                max_seq_len=max_seq_len,
+                dtype=dtype,
+                trust_remote_code=trust_remote_code,
+            )
+            for idx, step_text in enumerate(step_texts):
+                if idx < len(step_scores):
+                    continue
+                s_val = _score_step(
+                    judge,
+                    system_prompt=sys_prompt,
+                    step_text=step_text,
+                    prompt=prompt,
+                    completion=completion,
+                    rubric_text=rubric_text,
+                    temperature=temperature,
+                    max_new_tokens=max_new_tokens,
+                )
+                if s_val is not None:
+                    step_scores.append(float(s_val))
+
+    process_score = _aggregate_scores(step_scores, process_agg) if step_scores else None
+    if mode == "thinkprm" and process_score is not None:
+        score = process_score
+    if passed_val is None and score is not None:
+        passed_val = score >= min_score
+    passed = bool(passed_val) if passed_val is not None else False
+    reason = parsed.get("reason") or parsed.get("explanation") or ""
+
+    if reward_mode == "score" and score is not None:
+        reward = max(0.0, min(1.0, float(score)))
+    else:
+        reward = reward_pass if passed else reward_fail
+
+    latency_ms = (time.time() - t0) * 1000.0
+
+    return VerifyResult(
+        reward=reward,
+        passed=passed,
+        info={
+            "mode": mode,
+            "model": model_id,
+            "score": score,
+            "process_score": process_score,
+            "process_agg": process_agg,
+            "steps": step_texts,
+            "step_scores": step_scores,
+            "passed": passed,
+            "reason": reason,
+            "raw": raw,
+            "verifier_latency_ms": latency_ms,
+        },
+        artifacts_dir=workdir,
+    )

mlxsmith/verifiers/prime.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import importlib.util
+import re
+from typing import Any, Dict, List, Optional
+
+from .types import VerifyResult
+
+_STATE: Dict[str, Dict[str, float]] = {
+    "values": {},
+    "counts": {},
+}
+
+
+def _load_verifier(path: str):
+    import sys
+    from pathlib import Path as _Path
+
+    verifier_path = _Path(path).resolve()
+
+    # If the file lives inside a Python package, set __package__ so that
+    # relative imports (e.g. ``from .types import ...``) work correctly.
+    pkg_name: Optional[str] = None
+    if (verifier_path.parent / "__init__.py").exists():
+        parts: list[str] = []
+        p = verifier_path.parent
+        while (p / "__init__.py").exists():
+            parts.insert(0, p.name)
+            p = p.parent
+        pkg_name = ".".join(parts)
+        root = str(p)
+        if root not in sys.path:
+            sys.path.insert(0, root)
+
+    mod_name = f"{pkg_name}._prime_loaded" if pkg_name else "prime_verifier"
+    spec = importlib.util.spec_from_file_location(mod_name, str(verifier_path))
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Could not load verifier: {verifier_path}")
+    module = importlib.util.module_from_spec(spec)
+    if pkg_name is not None:
+        module.__package__ = pkg_name
+    spec.loader.exec_module(module)  # type: ignore
+    verify_fn = getattr(module, "verify", None)
+    if not callable(verify_fn):
+        raise RuntimeError(f"Verifier must define verify(...): {verifier_path}")
+    return verify_fn
+
+
+def _extract_steps(text: str, *, max_steps: int = 12) -> List[str]:
+    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
+    steps = []
+    for ln in lines:
+        if re.match(r"^(\d+\.|\-|\*|\+)\s+", ln):
+            steps.append(re.sub(r"^(\d+\.|\-|\*|\+)\s+", "", ln).strip())
+    if not steps:
+        steps = lines[:max_steps]
+    return steps[:max_steps]
+
+
+def _aggregate(values: List[float], mode: str) -> float:
+    if not values:
+        return 0.0
+    mode = (mode or "mean").lower()
+    if mode == "min":
+        return min(values)
+    if mode == "product":
+        out = 1.0
+        for v in values:
+            out *= v
+        return out
+    return sum(values) / float(len(values))
+
+
+def verify(
+    prompt: str,
+    completion: str,
+    workdir: str,
+    *,
+    verifier: str,
+    verifier_kwargs: Optional[Dict[str, Any]] = None,
+    ema_alpha: float = 0.2,
+    max_steps: int = 12,
+    agg: str = "mean",
+    reward_mode: str = "process",
+    min_score: float = 0.0,
+    **kwargs,
+) -> VerifyResult:
+    """PRIME-style implicit process rewards.
+
+    Uses outcome reward from a base verifier to update per-step values.
+    """
+    verify_fn = _load_verifier(verifier)
+    base = verify_fn(prompt, completion, workdir, **(verifier_kwargs or {}), **kwargs)
+    outcome_reward = float(getattr(base, "reward", 0.0))
+    steps = _extract_steps(completion, max_steps=max_steps)
+
+    step_values: List[float] = []
+    for step in steps:
+        prev = _STATE["values"].get(step, outcome_reward)
+        new_val = (1.0 - ema_alpha) * prev + ema_alpha * outcome_reward
+        _STATE["values"][step] = new_val
+        _STATE["counts"][step] = _STATE["counts"].get(step, 0.0) + 1.0
+        step_values.append(new_val)
+
+    process_reward = _aggregate(step_values, agg)
+    if reward_mode == "combined":
+        reward = (process_reward + outcome_reward) / 2.0
+    else:
+        reward = process_reward
+
+    passed = bool(getattr(base, "passed", False)) and reward >= min_score
+
+    return VerifyResult(
+        reward=reward,
+        passed=passed,
+        info={
+            "mode": "prime",
+            "base_reward": outcome_reward,
+            "process_reward": process_reward,
+            "steps": steps,
+            "step_values": step_values,
+            "agg": agg,
+            "ema_alpha": ema_alpha,
+            "base_info": getattr(base, "info", {}),
+        },
+        artifacts_dir=workdir,
+    )