| """Tests for embedded SFT cold-start gate and degenerate probe splits.""" |
| import re |
|
|
| from opsd_utils.diagnostics import summarize_generate_probe_stats |
| from opsd_utils.gate_policy import ( |
| in_sft_cold_start, |
| resolve_max_training_steps, |
| resolve_skip_degenerate_opsd, |
| sft_cold_start_steps, |
| sft_slots_for_step, |
| ) |
| from reward_utils.format_checks import evaluate_format_reward |
|
|
|
|
| def _gate_cfg(): |
| return { |
| "gate": { |
| "skip_degenerate_for_opsd": True, |
| "degen_skip_warmup_steps": 200, |
| "sft_warmup_steps": 500, |
| "sft_warmup_slots_per_group": 4, |
| "sft_cold_start_frac": 0.08, |
| } |
| } |
|
|
|
|
| def test_sft_cold_start_steps_from_frac(): |
| cfg = _gate_cfg() |
| assert sft_cold_start_steps(cfg, 1000) == 80 |
| assert in_sft_cold_start(cfg, 79, 1000) is True |
| assert in_sft_cold_start(cfg, 80, 1000) is False |
|
|
|
|
| def test_skip_degenerate_after_cold_start_and_warmup(): |
| cfg = _gate_cfg() |
| max_steps = 1000 |
| assert resolve_skip_degenerate_opsd(cfg, 50, max_steps) is False |
| assert resolve_skip_degenerate_opsd(cfg, 280, max_steps) is True |
|
|
|
|
| def test_sft_slots_zero_during_cold_start(): |
| cfg = _gate_cfg() |
| assert sft_slots_for_step(cfg, 10, 1000) == 0 |
| assert sft_slots_for_step(cfg, 200, 1000) == 4 |
|
|
|
|
| def test_resolve_max_training_steps_from_state(): |
| class _State: |
| max_steps = 5890 |
|
|
| class _Args: |
| max_steps = -1 |
| num_train_epochs = 1 |
|
|
| class _Trainer: |
| args = _Args() |
| state = _State() |
|
|
| assert resolve_max_training_steps(_Trainer()) == 5890 |
| assert sft_cold_start_steps(_gate_cfg(), 5890) == 471 |
|
|
|
|
| def test_resolve_max_training_steps_from_args(): |
| class _Args: |
| max_steps = 200 |
| num_train_epochs = 10 |
|
|
| class _Trainer: |
| args = _Args() |
| state = None |
|
|
| assert resolve_max_training_steps(_Trainer()) == 200 |
|
|
|
|
| def test_graded_chart_format_partial_credit(): |
| partial = "Goal: x\nAnswer: 42" |
| assert ( |
| evaluate_format_reward( |
| partial, |
| "answer:", |
| re.compile(r"(?i)answer:"), |
| task="chart", |
| ) |
| == 0.5 |
| ) |
|
|
|
|
| def test_degenerate_probe_split_format_vs_repeat(): |
| import torch |
|
|
| |
| ids = torch.tensor([[17, 15, 18, 15, 151645]], dtype=torch.long) |
| mask = torch.tensor([[1, 1, 1, 1, 1]], dtype=torch.long) |
| is_eos = ids == 151645 |
| stats = summarize_generate_probe_stats( |
| ids, |
| mask, |
| is_eos, |
| eos_id=151645, |
| completions=["2030"], |
| answer_flag="Answer:", |
| max_completion_length=128, |
| ) |
| assert stats["degenerate_rate_format"] == 1.0 |
| assert stats["degenerate_rate_repeat"] == 0.0 |
|
|
