mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
dced3bfc45
* init * fix tests * unify codeio * filtered for libraries not present in reasoning-gym * fix more bounds * puzzle24 * knight swap curriculum * fix number sorting * fix attributes * add validation of config in creation of dataset * dry run for instantiating and validating the datasets * remove unused imports * fix curriculum tests to reference newly updated attribute names
186 lines
6.3 KiB
Python
186 lines
6.3 KiB
Python
from random import Random
|
|
|
|
import pytest
|
|
|
|
from reasoning_gym.arc import Arc1DConfig, Arc1DCurriculum, Arc1DDataset
|
|
|
|
|
|
def test_arc_1d_config_validation():
|
|
"""Test that invalid configs raise appropriate errors"""
|
|
with pytest.raises(AssertionError):
|
|
config = Arc1DConfig(min_size=0)
|
|
config.validate()
|
|
|
|
with pytest.raises(AssertionError):
|
|
config = Arc1DConfig(min_size=30, max_size=20)
|
|
config.validate()
|
|
|
|
with pytest.raises(AssertionError):
|
|
config = Arc1DConfig(num_train=0)
|
|
config.validate()
|
|
|
|
|
|
def test_arc_1d_deterministic():
|
|
"""Test that dataset generates same items with same seed"""
|
|
config = Arc1DConfig(seed=42, size=10)
|
|
dataset1 = Arc1DDataset(config)
|
|
dataset2 = Arc1DDataset(config)
|
|
|
|
for i in range(len(dataset1)):
|
|
assert dataset1[i] == dataset2[i]
|
|
|
|
|
|
def test_arc_1d_items():
|
|
"""Test basic properties of generated items"""
|
|
config = Arc1DConfig(min_size=10, max_size=15, num_train=2, size=50, seed=42)
|
|
dataset = Arc1DDataset(config)
|
|
|
|
for i in range(len(dataset)):
|
|
item = dataset[i]
|
|
assert isinstance(item, dict)
|
|
assert "question" in item
|
|
assert "answer" in item
|
|
assert "metadata" in item
|
|
assert "difficulty" in item["metadata"]
|
|
|
|
# Check metadata contents
|
|
metadata = item["metadata"]
|
|
assert "task_name" in metadata
|
|
assert "size" in metadata
|
|
assert "train_examples" in metadata
|
|
assert "test_example" in metadata
|
|
|
|
# Verify size constraints
|
|
assert config.min_size <= metadata["size"] <= config.max_size
|
|
|
|
# Check training examples
|
|
train_examples = metadata["train_examples"]
|
|
assert len(train_examples) == config.num_train
|
|
for example in train_examples:
|
|
assert "input" in example
|
|
assert "output" in example
|
|
assert len(example["input"]) == metadata["size"]
|
|
assert len(example["output"]) == metadata["size"]
|
|
|
|
# Check test example
|
|
test_example = metadata["test_example"]
|
|
assert "input" in test_example
|
|
assert "output" in test_example
|
|
assert len(test_example["input"]) == metadata["size"]
|
|
assert len(test_example["output"]) == metadata["size"]
|
|
|
|
|
|
def test_arc_1d_iteration():
|
|
"""Test that iteration respects dataset size"""
|
|
config = Arc1DConfig(size=100, seed=42) # Small size for testing
|
|
dataset = Arc1DDataset(config)
|
|
|
|
# Test manual iteration
|
|
items = []
|
|
for item in dataset:
|
|
items.append(item)
|
|
assert len(items) == config.size, "Iterator should yield exactly size items"
|
|
|
|
# Test list conversion
|
|
items = list(dataset)
|
|
assert len(items) == config.size, "Iterator should yield exactly size items"
|
|
|
|
# Test multiple iterations
|
|
first_items = list(dataset)
|
|
second_items = list(dataset)
|
|
assert first_items == second_items, "Multiple iterations should yield same items"
|
|
|
|
|
|
def test_arc_1d_scoring():
|
|
"""Test answer scoring logic"""
|
|
config = Arc1DConfig(size=1, seed=42)
|
|
dataset = Arc1DDataset(config)
|
|
entry = dataset[0]
|
|
|
|
# Test exact match
|
|
assert dataset.score_answer(entry["answer"], entry) == 1.0
|
|
|
|
# Test partial match (answer contained within response)
|
|
assert dataset.score_answer(f"The answer is: {entry['answer']}", entry) > 0.5
|
|
|
|
# Test incorrect answer
|
|
assert dataset.score_answer("wrong answer", entry) == 0.0
|
|
|
|
# Test None answer
|
|
assert dataset.score_answer(None, entry) == 0.0
|
|
|
|
|
|
@pytest.mark.parametrize("board_size", [8, 9, 10, 12, 15, 20])
|
|
def test_arc_1d_sizes(board_size: int):
|
|
config = Arc1DConfig(size=1000, seed=42 + board_size, min_size=board_size, max_size=board_size)
|
|
dataset = Arc1DDataset(config)
|
|
for entry in dataset:
|
|
assert len(entry["metadata"]["test_example"]["input"]) == board_size
|
|
assert len(entry["metadata"]["test_example"]["output"]) == board_size
|
|
assert dataset.score_answer(entry["answer"], entry) == 1.0
|
|
|
|
|
|
@pytest.mark.parametrize("min_size,max_size", [(8, 10), (9, 13), (10, 12), (12, 20)])
|
|
def test_arc_1d_size_ranges(min_size: int, max_size: int):
|
|
config = Arc1DConfig(size=1000, seed=42, min_size=min_size, max_size=max_size)
|
|
dataset = Arc1DDataset(config)
|
|
for entry in dataset:
|
|
assert min_size <= len(entry["metadata"]["test_example"]["input"]) <= max_size
|
|
assert min_size <= len(entry["metadata"]["test_example"]["output"]) <= max_size
|
|
assert dataset.score_answer(entry["answer"], entry) == 1.0
|
|
|
|
|
|
def test_arc_1d_generate_all_tasks():
|
|
config = Arc1DConfig(size=100, seed=17, min_size=8, max_size=10)
|
|
dataset = Arc1DDataset(config)
|
|
tasks = dataset.ARC_1D_TASKS
|
|
rng = Random(999)
|
|
for task_name, (generator_fn, args) in tasks.items():
|
|
for j in range(3):
|
|
for i in range(20):
|
|
x = generator_fn(rng=rng, size=10, **args)
|
|
if x is not None:
|
|
break
|
|
assert i < 20
|
|
print(task_name, j, i, x)
|
|
|
|
|
|
def test_arc_1d_curriculum():
|
|
"""Test the curriculum for complex arithmetic."""
|
|
curriculum = Arc1DCurriculum()
|
|
base_value = {"size": 150, "seed": 1}
|
|
|
|
base_cfg: Arc1DCurriculum = curriculum.generate_configuration(base_value)
|
|
|
|
assert base_cfg.seed == 1
|
|
assert base_cfg.size == 150
|
|
assert base_cfg.min_size == 10
|
|
assert base_cfg.max_size == 10
|
|
|
|
# Test and validate increase in levels
|
|
curriculum.increment_attr_level("size")
|
|
|
|
increased_cfg: Arc1DCurriculum = curriculum.generate_configuration(base_value)
|
|
assert increased_cfg.min_size == 10
|
|
assert increased_cfg.max_size == 25
|
|
|
|
# Test and validate decrease in levels
|
|
curriculum.decrement_attr_level("size")
|
|
|
|
decreased_cfg: Arc1DCurriculum = curriculum.generate_configuration(base_value)
|
|
assert decreased_cfg.min_size == 10
|
|
assert decreased_cfg.max_size == 10
|
|
|
|
# Test upper bound boundary condition
|
|
for _ in range(10):
|
|
curriculum.increment_attr_level("size")
|
|
upper_bound_cfg: Arc1DCurriculum = curriculum.generate_configuration(base_value)
|
|
assert upper_bound_cfg.min_size == 10
|
|
assert upper_bound_cfg.max_size == 100
|
|
|
|
# Test lower bound boundary condition
|
|
for _ in range(10):
|
|
curriculum.decrement_attr_level("size")
|
|
lower_bound_cfg: Arc1DCurriculum = curriculum.generate_configuration(base_value)
|
|
assert lower_bound_cfg.min_size == 10
|
|
assert lower_bound_cfg.max_size == 10
|