Blame for code/tests/test_cispo_config.py

"""configs/p4_cispo_lean.toml: the draft loads under the stage loader's struct
mode, selects CISPO + FP32 head, passes the P4 gate, and its denominator is the
DERIVED constant (a hand-typed one is refused by GrpoStageConfig.validate)."""

from __future__ import annotations

import pathlib

import pytest

from align.rl.cispo import assert_p4_cispo_conformant
from align.rl.grpo_loss import GRPOLossType
from align.rl.stage_config import load_stage_config

TOML = pathlib.Path(__file__).resolve().parents[1] / "configs" / "p4_cispo_lean.toml"


def _load(tmp_path, edit=lambda s: s):
    band = tmp_path / "band.jsonl"
    band.write_text("{}\n")
    text = TOML.read_text().replace("/apps/development/rlhf/align/data/p4/lean_band0.jsonl", str(band))
    p = tmp_path / "p4.toml"
    p.write_text(edit(text))
    return load_stage_config(str(p))


def test_draft_loads_and_selects_cispo_fp32(tmp_path):
    cfg = _load(tmp_path)
    g = cfg.grpo
    assert g.loss_fn == GRPOLossType.cispo and g.fp32_head is True
    assert g.cispo_clip_lower is None  # paper: no lower bound
    assert g.loss_denominator == 16 * 16 * 1536.0
    assert_p4_cispo_conformant(g)


def test_hand_typed_denominator_still_refused(tmp_path):
    with pytest.raises(ValueError, match="DERIVED"):
        _load(tmp_path, lambda s: s.replace("loss_denominator = 393216.0", "loss_denominator = 100000.0"))


def test_unknown_grpo_key_refused(tmp_path):
    with pytest.raises(TypeError):
        _load(tmp_path, lambda s: s.replace("fp32_head = true", "fp32_hed = true"))