diff --git a/src/soup_cli/trainer/prm.py b/src/soup_cli/trainer/prm.py index 9fc32fa..b7ad98f 100644 --- a/src/soup_cli/trainer/prm.py +++ b/src/soup_cli/trainer/prm.py @@ -16,6 +16,7 @@ per project policy — ``python -m soup_cli.cli --help`` must not pull torch. from __future__ import annotations import logging +import time from functools import lru_cache from pathlib import Path from typing import Any, Optional @@ -85,6 +86,45 @@ def make_prm_trainer_class(base_cls: type) -> type: return _PRMTrainer +def build_prm_train_result( + *, + log_history: list, + metrics: Any, + global_step: int, + duration_secs: float, + output_dir: str, +) -> dict: + """Build the standard trainer-result dict for the PRM path (v0.71.30). + + Mirrors the shape every other trainer wrapper returns so ``commands/train.py`` + can render its summary — previously the PRM wrapper returned a bespoke dict + missing ``initial_loss`` / ``final_loss`` / ``duration`` / ``total_steps``, + crashing the CLI with a ``KeyError`` right after ``save_model``. + """ + train_losses = [e["loss"] for e in log_history if isinstance(e, dict) and "loss" in e] + fallback = 0.0 + if isinstance(metrics, dict): + try: + fallback = float(metrics.get("train_loss", 0.0)) + except (TypeError, ValueError): + fallback = 0.0 + initial = train_losses[0] if train_losses else fallback + final = train_losses[-1] if train_losses else fallback + hours = int(duration_secs // 3600) + minutes = int((duration_secs % 3600) // 60) + duration = f"{hours}h {minutes}m" if hours > 0 else f"{minutes}m" + return { + "status": "ok", + "initial_loss": initial, + "final_loss": final, + "duration": duration, + "duration_secs": duration_secs, + "total_steps": global_step, + "output_dir": output_dir, + "metrics": metrics, + } + + def _prepare_prm_dataset(raw_rows: list[dict], tokenizer: Any, max_length: int) -> list[dict]: """Tokenise PRM rows into (input_ids, attention_mask, step_positions, labels). @@ -294,6 +334,18 @@ class PRMTrainerWrapper: data_collator=collator, ) console.print("[green]Starting PRM training...[/]") + start = time.time() result = self.trainer.train() self.trainer.save_model(str(output_dir)) - return {"status": "ok", "output_dir": str(output_dir), "metrics": result.metrics} + # v0.71.30 — save the tokenizer alongside the model so the PRM + # checkpoint is loadable standalone (soup shrink / PRMScorer / + # `soup train prm_reward=` all call AutoTokenizer.from_pretrained + # on the dir). Previously the tokenizer was never persisted. + self.tokenizer.save_pretrained(str(output_dir)) + return build_prm_train_result( + log_history=self.trainer.state.log_history, + metrics=result.metrics, + global_step=self.trainer.state.global_step, + duration_secs=time.time() - start, + output_dir=str(output_dir), + ) diff --git a/tests/test_v07130.py b/tests/test_v07130.py index d19f8d9..ec3acd1 100644 --- a/tests/test_v07130.py +++ b/tests/test_v07130.py @@ -727,6 +727,60 @@ class TestRecipes: assert len(RECIPES) == 137 +# --------------------------------------------------------------------------- +# PRM producer fixes surfaced by the live smoke (train-result shape + tokenizer) +# --------------------------------------------------------------------------- +class TestBuildPrmTrainResult: + def test_has_all_summary_keys(self): + from soup_cli.trainer.prm import build_prm_train_result + + out = build_prm_train_result( + log_history=[{"loss": 3.2}, {"loss": 1.1}], + metrics={"train_loss": 2.0}, + global_step=6, + duration_secs=125.0, + output_dir="./out", + ) + # These are exactly the keys commands/train.py's summary indexes. + for key in ( + "initial_loss", + "final_loss", + "duration", + "duration_secs", + "total_steps", + "output_dir", + ): + assert key in out, key + assert out["initial_loss"] == pytest.approx(3.2) + assert out["final_loss"] == pytest.approx(1.1) + assert out["total_steps"] == 6 + assert out["duration"] == "2m" + + def test_empty_log_history_falls_back_to_metrics(self): + from soup_cli.trainer.prm import build_prm_train_result + + out = build_prm_train_result( + log_history=[], + metrics={"train_loss": 2.5}, + global_step=0, + duration_secs=3700.0, + output_dir="./out", + ) + assert out["initial_loss"] == pytest.approx(2.5) + assert out["final_loss"] == pytest.approx(2.5) + assert out["duration"] == "1h 1m" + + def test_train_saves_tokenizer(self): + # Guard: PRMTrainerWrapper.train() must persist the tokenizer so the + # PRM checkpoint is loadable standalone by PRMScorer. + import inspect + + from soup_cli.trainer.prm import PRMTrainerWrapper + + src = inspect.getsource(PRMTrainerWrapper.train) + assert "self.tokenizer.save_pretrained" in src + + class TestNoTopLevelTorch: def test_prm_reward_has_no_top_level_torch(self): import soup_cli.utils.prm_reward as mod