mirror of https://github.com/razor-ai/soup.git
274 lines
11 KiB
Python
274 lines
11 KiB
Python
"""Tests for v0.40.0 Part B — unified preference loss dispatcher.
|
|
|
|
Adds ``task: preference`` + ``training.preference_loss`` Literal
|
|
(dpo/simpo/orpo/ipo/bco). Existing per-task forms (``task: dpo`` etc) keep
|
|
working unchanged — the unified surface is *additive*, not a breaking
|
|
collapse. Backward-compat ``resolve_preference_loss`` maps legacy task
|
|
strings to their preference_loss equivalents for callers that want a
|
|
single dispatch entry-point.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock
|
|
from unittest.mock import patch as mock_patch
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from soup_cli.config.schema import SoupConfig
|
|
|
|
# ─── Schema Tests ───────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestPreferenceTaskField:
|
|
def test_preference_task_accepted(self):
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="preference",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
training={"preference_loss": "dpo"},
|
|
)
|
|
assert cfg.task == "preference"
|
|
assert cfg.training.preference_loss == "dpo"
|
|
|
|
@pytest.mark.parametrize("loss", ["dpo", "simpo", "orpo", "ipo", "bco"])
|
|
def test_preference_loss_each_value(self, loss):
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="preference",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
training={"preference_loss": loss},
|
|
)
|
|
assert cfg.training.preference_loss == loss
|
|
|
|
def test_preference_loss_unknown_rejected(self):
|
|
with pytest.raises(ValidationError, match="preference_loss"):
|
|
SoupConfig(
|
|
base="some-model",
|
|
task="preference",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
training={"preference_loss": "garbage"},
|
|
)
|
|
|
|
def test_preference_task_requires_preference_loss(self):
|
|
"""task=preference without preference_loss must error."""
|
|
with pytest.raises(ValidationError, match="preference_loss"):
|
|
SoupConfig(
|
|
base="some-model",
|
|
task="preference",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
)
|
|
|
|
def test_preference_loss_default_none_for_other_tasks(self):
|
|
"""Non-preference tasks default preference_loss=None."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="dpo",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
)
|
|
assert cfg.training.preference_loss is None
|
|
|
|
def test_preference_loss_set_outside_preference_task_rejected(self):
|
|
"""preference_loss is meaningful only when task=preference."""
|
|
with pytest.raises(ValidationError, match="preference_loss"):
|
|
SoupConfig(
|
|
base="some-model",
|
|
task="dpo",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
training={"preference_loss": "dpo"},
|
|
)
|
|
|
|
|
|
# ─── Resolver ──────────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestResolvePreferenceLoss:
|
|
def test_resolve_preference_task_returns_loss(self):
|
|
from soup_cli.trainer.preference import resolve_preference_loss
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="preference",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
training={"preference_loss": "simpo"},
|
|
)
|
|
assert resolve_preference_loss(cfg) == "simpo"
|
|
|
|
@pytest.mark.parametrize(
|
|
"task,expected",
|
|
[
|
|
("dpo", "dpo"),
|
|
("simpo", "simpo"),
|
|
("orpo", "orpo"),
|
|
("ipo", "ipo"),
|
|
("bco", "bco"),
|
|
],
|
|
)
|
|
def test_legacy_task_maps_to_loss(self, task, expected):
|
|
from soup_cli.trainer.preference import resolve_preference_loss
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task=task,
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
)
|
|
assert resolve_preference_loss(cfg) == expected
|
|
|
|
def test_resolve_non_preference_task_returns_none(self):
|
|
from soup_cli.trainer.preference import resolve_preference_loss
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="sft",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert resolve_preference_loss(cfg) is None
|
|
|
|
|
|
# ─── Wrapper Tests ─────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestPreferenceTrainerWrapper:
|
|
def test_import_exists(self):
|
|
from soup_cli.trainer.preference import PreferenceTrainerWrapper
|
|
|
|
assert PreferenceTrainerWrapper is not None
|
|
|
|
@pytest.mark.parametrize(
|
|
"loss,wrapper_path",
|
|
[
|
|
("dpo", "soup_cli.trainer.dpo.DPOTrainerWrapper"),
|
|
("simpo", "soup_cli.trainer.simpo.SimPOTrainerWrapper"),
|
|
("orpo", "soup_cli.trainer.orpo.ORPOTrainerWrapper"),
|
|
("ipo", "soup_cli.trainer.ipo.IPOTrainerWrapper"),
|
|
("bco", "soup_cli.trainer.bco.BCOTrainerWrapper"),
|
|
],
|
|
)
|
|
def test_dispatcher_routes_to_correct_wrapper(self, loss, wrapper_path):
|
|
"""PreferenceTrainerWrapper.setup() must delegate to the right wrapper.
|
|
|
|
Asserts the inner cfg sent to the wrapper has task=loss and that
|
|
preference_loss has been cleared (defends _make_inner_cfg's contract).
|
|
"""
|
|
from soup_cli.trainer.preference import PreferenceTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="preference",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
training={"preference_loss": loss},
|
|
)
|
|
with mock_patch(wrapper_path) as mock_inner:
|
|
mock_instance = MagicMock()
|
|
mock_instance.train.return_value = {
|
|
"initial_loss": 1.0, "final_loss": 0.5,
|
|
"duration": "1m", "duration_secs": 60.0,
|
|
"output_dir": "./out", "total_steps": 10,
|
|
}
|
|
mock_inner.return_value = mock_instance
|
|
|
|
wrapper = PreferenceTrainerWrapper(cfg, device="cpu")
|
|
wrapper.setup({"train": [{"prompt": "p", "chosen": "c", "rejected": "r"}]})
|
|
wrapper.train()
|
|
|
|
mock_inner.assert_called_once()
|
|
inner_cfg = mock_inner.call_args[0][0]
|
|
assert inner_cfg.task == loss
|
|
assert inner_cfg.training.preference_loss is None
|
|
mock_instance.setup.assert_called_once()
|
|
mock_instance.train.assert_called_once()
|
|
|
|
def test_setup_does_not_mutate_caller_cfg(self):
|
|
"""_make_inner_cfg must return a copy; cfg.task unchanged after setup."""
|
|
from soup_cli.trainer.preference import PreferenceTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="preference",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
training={"preference_loss": "dpo"},
|
|
)
|
|
with mock_patch("soup_cli.trainer.dpo.DPOTrainerWrapper") as mock_inner:
|
|
mock_inner.return_value = MagicMock()
|
|
wrapper = PreferenceTrainerWrapper(cfg, device="cpu")
|
|
wrapper.setup({"train": [{"prompt": "p", "chosen": "c", "rejected": "r"}]})
|
|
# Caller's cfg untouched.
|
|
assert cfg.task == "preference"
|
|
assert cfg.training.preference_loss == "dpo"
|
|
|
|
def test_train_before_setup_raises(self):
|
|
from soup_cli.trainer.preference import PreferenceTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="preference",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
training={"preference_loss": "dpo"},
|
|
)
|
|
wrapper = PreferenceTrainerWrapper(cfg, device="cpu")
|
|
with pytest.raises(RuntimeError, match="setup"):
|
|
wrapper.train()
|
|
|
|
def test_unknown_loss_raises_at_setup(self):
|
|
"""Defence-in-depth: schema gate prevents this, but if it slipped, raise."""
|
|
from soup_cli.trainer.preference import PreferenceTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="preference",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
training={"preference_loss": "dpo"},
|
|
)
|
|
# Mutate after construction to bypass schema validation.
|
|
cfg.training.preference_loss = "garbage"
|
|
wrapper = PreferenceTrainerWrapper(cfg, device="cpu")
|
|
with pytest.raises(ValueError, match="preference_loss"):
|
|
wrapper.setup({"train": [{"prompt": "p", "chosen": "c", "rejected": "r"}]})
|
|
|
|
|
|
# ─── Train Routing ─────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestPreferenceTrainRouting:
|
|
def test_sweep_routes_to_preference_wrapper(self):
|
|
from soup_cli.commands.sweep import _run_single
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="preference",
|
|
data={"train": "./data.jsonl", "format": "dpo"},
|
|
training={"preference_loss": "dpo"},
|
|
)
|
|
fake_dataset = {"train": [{"prompt": "Q", "chosen": "A", "rejected": "B"}]}
|
|
fake_result = {
|
|
"initial_loss": 1.0, "final_loss": 0.5,
|
|
"duration": "1m", "duration_secs": 60.0,
|
|
"output_dir": "./out", "total_steps": 10,
|
|
}
|
|
fake_gpu_info = {"memory_total": "0 MB", "memory_total_bytes": 0}
|
|
|
|
with mock_patch(
|
|
"soup_cli.data.loader.load_dataset", return_value=fake_dataset,
|
|
), mock_patch(
|
|
"soup_cli.utils.gpu.detect_device", return_value=("cpu", "CPU"),
|
|
), mock_patch(
|
|
"soup_cli.utils.gpu.get_gpu_info", return_value=fake_gpu_info,
|
|
), mock_patch(
|
|
"soup_cli.experiment.tracker.ExperimentTracker",
|
|
) as mock_tracker_cls, mock_patch(
|
|
"soup_cli.monitoring.display.TrainingDisplay",
|
|
), mock_patch(
|
|
"soup_cli.trainer.preference.PreferenceTrainerWrapper.setup",
|
|
), mock_patch(
|
|
"soup_cli.trainer.preference.PreferenceTrainerWrapper.train",
|
|
return_value=fake_result,
|
|
) as mock_train:
|
|
mock_tracker = MagicMock()
|
|
mock_tracker.start_run.return_value = "run-pref-1"
|
|
mock_tracker_cls.return_value = mock_tracker
|
|
result = _run_single(cfg, {}, "pref_run_1", None)
|
|
|
|
mock_train.assert_called_once()
|
|
assert result["run_id"] == "run-pref-1"
|