mirror of https://github.com/razor-ai/soup.git
427 lines
14 KiB
Python
427 lines
14 KiB
Python
"""Tests for GRPO training — config, rewards, data preparation, template."""
|
|
|
|
import textwrap
|
|
|
|
import pytest
|
|
|
|
from soup_cli.config.schema import TEMPLATES, SoupConfig
|
|
|
|
# ─── Config Tests ───────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestGRPOConfig:
|
|
"""Test GRPO task config validation."""
|
|
|
|
def test_grpo_task_accepted(self):
|
|
"""GRPO task should be a valid task type."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.task == "grpo"
|
|
|
|
def test_grpo_beta_default(self):
|
|
"""grpo_beta should default to 0.1."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.training.grpo_beta == 0.1
|
|
|
|
def test_grpo_beta_custom(self):
|
|
"""Custom grpo_beta should be accepted."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"grpo_beta": 0.04},
|
|
)
|
|
assert cfg.training.grpo_beta == pytest.approx(0.04)
|
|
|
|
def test_grpo_beta_must_be_positive(self):
|
|
"""grpo_beta must be > 0."""
|
|
with pytest.raises(Exception):
|
|
SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"grpo_beta": 0},
|
|
)
|
|
|
|
def test_num_generations_default(self):
|
|
"""num_generations should default to 4."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.training.num_generations == 4
|
|
|
|
def test_num_generations_custom(self):
|
|
"""Custom num_generations should be accepted."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"num_generations": 8},
|
|
)
|
|
assert cfg.training.num_generations == 8
|
|
|
|
def test_num_generations_minimum(self):
|
|
"""num_generations must be >= 2."""
|
|
with pytest.raises(Exception):
|
|
SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"num_generations": 1},
|
|
)
|
|
|
|
def test_reward_fn_default(self):
|
|
"""reward_fn should default to 'accuracy'."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.training.reward_fn == "accuracy"
|
|
|
|
def test_reward_fn_custom_path(self):
|
|
"""reward_fn should accept a custom file path."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"reward_fn": "./my_reward.py"},
|
|
)
|
|
assert cfg.training.reward_fn == "./my_reward.py"
|
|
|
|
def test_grpo_full_config(self):
|
|
"""Full GRPO config should validate correctly."""
|
|
cfg = SoupConfig(
|
|
base="meta-llama/Llama-3.1-8B-Instruct",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl", "format": "sharegpt", "max_length": 4096},
|
|
training={
|
|
"epochs": 3,
|
|
"lr": 1e-5,
|
|
"grpo_beta": 0.1,
|
|
"num_generations": 4,
|
|
"reward_fn": "format",
|
|
"lora": {"r": 64, "alpha": 16},
|
|
"quantization": "4bit",
|
|
},
|
|
)
|
|
assert cfg.task == "grpo"
|
|
assert cfg.training.reward_fn == "format"
|
|
assert cfg.training.num_generations == 4
|
|
assert cfg.data.max_length == 4096
|
|
|
|
|
|
# ─── Reward Function Tests ──────────────────────────────────────────────────
|
|
|
|
|
|
class TestAccuracyReward:
|
|
"""Test the accuracy reward function."""
|
|
|
|
def test_exact_match(self):
|
|
from soup_cli.trainer.rewards import accuracy_reward
|
|
|
|
completions = [[{"role": "assistant", "content": "The answer is #### 42"}]]
|
|
rewards = accuracy_reward(completions, answer=["42"])
|
|
assert rewards == [1.0]
|
|
|
|
def test_boxed_match(self):
|
|
from soup_cli.trainer.rewards import accuracy_reward
|
|
|
|
completions = [[{"role": "assistant", "content": "So \\boxed{42} is the result"}]]
|
|
rewards = accuracy_reward(completions, answer=["42"])
|
|
assert rewards == [1.0]
|
|
|
|
def test_partial_match(self):
|
|
from soup_cli.trainer.rewards import accuracy_reward
|
|
|
|
completions = [[{"role": "assistant", "content": "The answer is 42 degrees"}]]
|
|
rewards = accuracy_reward(completions, answer=["42"])
|
|
assert rewards == [0.5]
|
|
|
|
def test_no_match(self):
|
|
from soup_cli.trainer.rewards import accuracy_reward
|
|
|
|
completions = [[{"role": "assistant", "content": "I don't know"}]]
|
|
rewards = accuracy_reward(completions, answer=["42"])
|
|
assert rewards == [0.0]
|
|
|
|
def test_multiple_completions(self):
|
|
from soup_cli.trainer.rewards import accuracy_reward
|
|
|
|
completions = [
|
|
[{"role": "assistant", "content": "#### 42"}],
|
|
[{"role": "assistant", "content": "Wrong answer"}],
|
|
[{"role": "assistant", "content": "The answer is 42"}],
|
|
]
|
|
rewards = accuracy_reward(completions, answer=["42", "42", "42"])
|
|
assert rewards == [1.0, 0.0, 0.5]
|
|
|
|
def test_empty_completion(self):
|
|
from soup_cli.trainer.rewards import accuracy_reward
|
|
|
|
completions = [[]]
|
|
rewards = accuracy_reward(completions, answer=["42"])
|
|
assert rewards == [0.0]
|
|
|
|
|
|
class TestFormatReward:
|
|
"""Test the format reward function."""
|
|
|
|
def test_perfect_format(self):
|
|
from soup_cli.trainer.rewards import format_reward
|
|
|
|
content = "<think>Let me think step by step...</think>\nThe answer is 42."
|
|
completions = [[{"role": "assistant", "content": content}]]
|
|
rewards = format_reward(completions)
|
|
assert rewards == [1.0]
|
|
|
|
def test_think_only(self):
|
|
from soup_cli.trainer.rewards import format_reward
|
|
|
|
content = "<think>Thinking...</think>"
|
|
completions = [[{"role": "assistant", "content": content}]]
|
|
rewards = format_reward(completions)
|
|
assert rewards == [0.5]
|
|
|
|
def test_no_format(self):
|
|
from soup_cli.trainer.rewards import format_reward
|
|
|
|
completions = [[{"role": "assistant", "content": "Just a plain answer"}]]
|
|
rewards = format_reward(completions)
|
|
assert rewards == [0.0]
|
|
|
|
def test_multiple_completions(self):
|
|
from soup_cli.trainer.rewards import format_reward
|
|
|
|
completions = [
|
|
[{"role": "assistant", "content": "<think>A</think>\nB"}],
|
|
[{"role": "assistant", "content": "No format"}],
|
|
]
|
|
rewards = format_reward(completions)
|
|
assert rewards == [1.0, 0.0]
|
|
|
|
|
|
class TestExtractAnswer:
|
|
"""Test answer extraction from model output."""
|
|
|
|
def test_hash_format(self):
|
|
from soup_cli.trainer.rewards import _extract_answer
|
|
|
|
assert _extract_answer("Some work\n#### 42") == "42"
|
|
|
|
def test_boxed_format(self):
|
|
from soup_cli.trainer.rewards import _extract_answer
|
|
|
|
assert _extract_answer("So \\boxed{42} is the answer") == "42"
|
|
|
|
def test_no_answer(self):
|
|
from soup_cli.trainer.rewards import _extract_answer
|
|
|
|
assert _extract_answer("Just plain text") is None
|
|
|
|
def test_multiple_hashes(self):
|
|
from soup_cli.trainer.rewards import _extract_answer
|
|
|
|
assert _extract_answer("#### step\n#### 42") == "42"
|
|
|
|
|
|
class TestLoadRewardFn:
|
|
"""Test reward function loading."""
|
|
|
|
def test_load_builtin_accuracy(self):
|
|
from soup_cli.trainer.rewards import accuracy_reward, load_reward_fn
|
|
|
|
fn = load_reward_fn("accuracy")
|
|
assert fn is accuracy_reward
|
|
|
|
def test_load_builtin_format(self):
|
|
from soup_cli.trainer.rewards import format_reward, load_reward_fn
|
|
|
|
fn = load_reward_fn("format")
|
|
assert fn is format_reward
|
|
|
|
def test_load_custom_file(self, tmp_path):
|
|
from soup_cli.trainer.rewards import load_reward_fn
|
|
|
|
custom_file = tmp_path / "my_reward.py"
|
|
custom_file.write_text(textwrap.dedent("""\
|
|
def reward_fn(completions, **kwargs):
|
|
return [1.0] * len(completions)
|
|
"""))
|
|
fn = load_reward_fn(str(custom_file))
|
|
result = fn([[{"content": "test"}]])
|
|
assert result == [1.0]
|
|
|
|
def test_load_custom_file_missing_fn(self, tmp_path):
|
|
from soup_cli.trainer.rewards import load_reward_fn
|
|
|
|
custom_file = tmp_path / "bad_reward.py"
|
|
custom_file.write_text("x = 1\n")
|
|
with pytest.raises(ValueError, match="must define a 'reward_fn'"):
|
|
load_reward_fn(str(custom_file))
|
|
|
|
def test_load_unknown_name(self):
|
|
from soup_cli.trainer.rewards import load_reward_fn
|
|
|
|
with pytest.raises(ValueError, match="Unknown reward function"):
|
|
load_reward_fn("nonexistent")
|
|
|
|
|
|
# ─── Data Preparation Tests ─────────────────────────────────────────────────
|
|
|
|
|
|
class TestPrepareGRPODataset:
|
|
"""Test GRPO dataset preparation."""
|
|
|
|
def test_from_prompt_string(self):
|
|
from soup_cli.trainer.grpo import _prepare_grpo_dataset
|
|
|
|
data = [{"prompt": "What is 2+2?", "answer": "4"}]
|
|
result = _prepare_grpo_dataset(data)
|
|
assert len(result) == 1
|
|
assert result[0]["prompt"] == [{"role": "user", "content": "What is 2+2?"}]
|
|
assert result[0]["answer"] == "4"
|
|
|
|
def test_from_messages(self):
|
|
from soup_cli.trainer.grpo import _prepare_grpo_dataset
|
|
|
|
data = [
|
|
{
|
|
"messages": [
|
|
{"role": "system", "content": "You are helpful."},
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi!"},
|
|
]
|
|
}
|
|
]
|
|
result = _prepare_grpo_dataset(data)
|
|
assert len(result) == 1
|
|
# Should only include non-assistant messages as prompt
|
|
assert len(result[0]["prompt"]) == 2
|
|
assert result[0]["prompt"][0]["role"] == "system"
|
|
assert result[0]["prompt"][1]["role"] == "user"
|
|
|
|
def test_from_prompt_message_list(self):
|
|
from soup_cli.trainer.grpo import _prepare_grpo_dataset
|
|
|
|
data = [
|
|
{
|
|
"prompt": [{"role": "user", "content": "What is 2+2?"}],
|
|
"answer": "4",
|
|
}
|
|
]
|
|
result = _prepare_grpo_dataset(data)
|
|
assert result[0]["prompt"] == [{"role": "user", "content": "What is 2+2?"}]
|
|
assert result[0]["answer"] == "4"
|
|
|
|
def test_from_alpaca_format(self):
|
|
from soup_cli.trainer.grpo import _prepare_grpo_dataset
|
|
|
|
data = [{"instruction": "Translate hello", "input": "", "output": "hola"}]
|
|
result = _prepare_grpo_dataset(data)
|
|
assert result[0]["prompt"] == [{"role": "user", "content": "Translate hello"}]
|
|
assert result[0]["answer"] == "hola"
|
|
|
|
def test_multiple_rows(self):
|
|
from soup_cli.trainer.grpo import _prepare_grpo_dataset
|
|
|
|
data = [
|
|
{"prompt": "Q1", "answer": "A1"},
|
|
{"prompt": "Q2", "answer": "A2"},
|
|
{"prompt": "Q3", "answer": "A3"},
|
|
]
|
|
result = _prepare_grpo_dataset(data)
|
|
assert len(result) == 3
|
|
|
|
|
|
# ─── Template Tests ──────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestReasoningTemplate:
|
|
"""Test the reasoning/GRPO template."""
|
|
|
|
def test_reasoning_template_exists(self):
|
|
assert "reasoning" in TEMPLATES
|
|
|
|
def test_reasoning_template_valid_yaml(self):
|
|
import yaml
|
|
|
|
config = yaml.safe_load(TEMPLATES["reasoning"])
|
|
assert config["task"] == "grpo"
|
|
assert config["training"]["grpo_beta"] == 0.1
|
|
assert config["training"]["num_generations"] == 4
|
|
assert config["training"]["reward_fn"] == "accuracy"
|
|
|
|
def test_reasoning_template_valid_config(self):
|
|
import yaml
|
|
|
|
raw = yaml.safe_load(TEMPLATES["reasoning"])
|
|
cfg = SoupConfig(**raw)
|
|
assert cfg.task == "grpo"
|
|
assert cfg.training.grpo_beta == 0.1
|
|
|
|
|
|
# ─── Train Command Routing Tests ─────────────────────────────────────────────
|
|
|
|
|
|
class TestGRPOTrainRouting:
|
|
"""Test that train command routes to GRPO trainer."""
|
|
|
|
def test_grpo_import_exists(self):
|
|
"""GRPOTrainerWrapper should be importable."""
|
|
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
|
|
|
assert GRPOTrainerWrapper is not None
|
|
|
|
def test_grpo_wrapper_init(self):
|
|
"""GRPOTrainerWrapper should initialize without error."""
|
|
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
wrapper = GRPOTrainerWrapper(cfg, device="cpu")
|
|
assert wrapper.config.task == "grpo"
|
|
assert wrapper.device == "cpu"
|
|
assert wrapper.model is None
|
|
assert wrapper.trainer is None
|
|
|
|
|
|
# ─── Sweep Shortcut Tests ────────────────────────────────────────────────────
|
|
|
|
|
|
class TestGRPOSweepParams:
|
|
"""Test GRPO parameter shortcuts in sweep."""
|
|
|
|
def test_grpo_beta_shortcut(self):
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config = {"training": {"grpo_beta": 0.1}}
|
|
_set_nested_param(config, "grpo_beta", 0.04)
|
|
assert config["training"]["grpo_beta"] == 0.04
|
|
|
|
def test_num_generations_shortcut(self):
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config = {"training": {"num_generations": 4}}
|
|
_set_nested_param(config, "num_generations", 8)
|
|
assert config["training"]["num_generations"] == 8
|
|
|
|
def test_reward_fn_shortcut(self):
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config = {"training": {"reward_fn": "accuracy"}}
|
|
_set_nested_param(config, "reward_fn", "format")
|
|
assert config["training"]["reward_fn"] == "format"
|