formatting

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

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