open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-04-29 17:35:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

formatting

Browse source
This commit is contained in:
Andreas Koepf 2025-02-09 00:04:42 +01:00
parent 8d8d85e6b2
commit e56316ebb2
2 changed files with 13 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

16
reasoning_gym/arc/arc_agi.py
View file

@ -99,13 +99,13 @@ class ArcAgiDataset(ProceduralDataset):
def _create_augmentation_fn(self, rng: Random) -> Callable[[Board], Board]:
"""Create a composite augmentation function from enabled options"""
fns = []
if self.config.use_rotations:
fns.append(rng.choice(ROTATION_AUGMENTATIONS))
if self.config.use_mirrors:
fns.append(rng.choice(MIRROR_AUGMENTATIONS))
if self.config.use_color_permutation:
color_table = list(range(10))
rng.shuffle(color_table)
@ -116,7 +116,7 @@ class ArcAgiDataset(ProceduralDataset):
for fn in fns:
result = fn(result)
return result
return composite_fn
def __getitem__(self, idx: int) -> dict:
@ -137,16 +137,14 @@ class ArcAgiDataset(ProceduralDataset):
# Apply augmentation to all train examples
augmented_train = []
for p in train:
augmented_train.append({
"input": augment(p["input"]),
"output": augment(p["output"])
})
augmented_train.append({"input": augment(p["input"]), "output": augment(p["output"])})
examples = [
format_board_pair(i + 1, p, formatting_options=self.config.board_format_opts)
format_board_pair(i + 1, p, formatting_options=self.config.board_format_opts)
for i, p in enumerate(augmented_train)
]
examples = "".join(examples)
# Apply augmentation to test example
augmented_test_input = augment(test["input"])
augmented_test_output = augment(test["output"])

32
tests/test_arc_agi.py
View file

@ -52,48 +52,28 @@ def test_arc_agi_items():
def test_arc_agi_augmentations():
"""Test that augmentations can be selectively enabled/disabled"""
# Test with all augmentations disabled
config = ArcAgiConfig(
seed=42,
size=10,
use_rotations=False,
use_mirrors=False,
use_color_permutation=False
)
config = ArcAgiConfig(seed=42, size=10, use_rotations=False, use_mirrors=False, use_color_permutation=False)
base_dataset = ArcAgiDataset(config)
base_items = list(base_dataset)
# Test with rotations only
rot_config = ArcAgiConfig(
seed=42,
size=10,
use_rotations=True,
use_mirrors=False,
use_color_permutation=False
)
rot_config = ArcAgiConfig(seed=42, size=10, use_rotations=True, use_mirrors=False, use_color_permutation=False)
rot_dataset = ArcAgiDataset(rot_config)
rot_items = list(rot_dataset)
# Items should differ when rotations are enabled
assert any(
base_items[i]["metadata"]["input"] != rot_items[i]["metadata"]["input"]
for i in range(len(base_items))
base_items[i]["metadata"]["input"] != rot_items[i]["metadata"]["input"] for i in range(len(base_items))
), "Rotation augmentation had no effect"
# Test with color permutation only
color_config = ArcAgiConfig(
seed=42,
size=10,
use_rotations=False,
use_mirrors=False,
use_color_permutation=True
)
color_config = ArcAgiConfig(seed=42, size=10, use_rotations=False, use_mirrors=False, use_color_permutation=True)
color_dataset = ArcAgiDataset(color_config)
color_items = list(color_dataset)
# Items should differ when color permutation is enabled
assert any(
base_items[i]["metadata"]["input"] != color_items[i]["metadata"]["input"]
for i in range(len(base_items))
base_items[i]["metadata"]["input"] != color_items[i]["metadata"]["input"] for i in range(len(base_items))
), "Color permutation had no effect"
@ -113,7 +93,7 @@ def test_arc_agi_scoring():
assert dataset.score_answer(None, entry=item) == 0.0
# Test wrong but valid grid format
wrong_answer = "0 0\n0 0"
wrong_answer = "1 0 0 0\n0 0 0 1"
assert dataset.score_answer(wrong_answer, entry=item) == 0.05