open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-05-01 17:45:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add more config params

Browse source
This commit is contained in:
Zafir Stojanovski 2025-02-10 22:30:36 +01:00
parent 696fdf8be7
commit f255831f1c
2 changed files with 18 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

21
reasoning_gym/algorithmic/manipulate_matrix.py
View file

@ -28,8 +28,11 @@ def num_cols(matrix: list[list[int]]) -> int:
class ManipulateMatrixConfig: class ManipulateMatrixConfig:
"""Configuration for Manipulate Matrix dataset generation""" """Configuration for Manipulate Matrix dataset generation"""
min_rows: int = 1 # Minimum number of rows
min_cols: int = 1 # Minimum number of columns
max_rows: int = 10 # Maximum number of rows max_rows: int = 10 # Maximum number of rows
max_cols: int = 10 # Maximum number of columns max_cols: int = 10 # Maximum number of columns
max_transforms: int = 5 # Maximum number of transformations to apply
p_rotate: float = 0.2 # Probability of rotating the matrix p_rotate: float = 0.2 # Probability of rotating the matrix
p_hmirror: float = 0.2 # Probability of horizontally mirroring the matrix p_hmirror: float = 0.2 # Probability of horizontally mirroring the matrix
p_vmirror: float = 0.2 # Probability of vertically mirroring the matrix p_vmirror: float = 0.2 # Probability of vertically mirroring the matrix
@ -46,8 +49,11 @@ class ManipulateMatrixConfig:
def validate(self): def validate(self):
"""Validate configuration parameters""" """Validate configuration parameters"""
assert 1 <= self.max_rows, "max_rows must be at least 1" assert 1 <= self.min_rows, "min_rows must be at least 1"
assert 1 <= self.max_cols, "max_cols must be at least 1" assert 1 <= self.min_cols, "min_cols must be at least 1"
assert self.min_rows <= self.max_rows, "max_rows must be at least min_rows"
assert self.min_cols <= self.max_cols, "max_cols must be at least min_cols"
assert 0 <= self.max_transforms, "max_transforms must be non-negative"
assert 0 <= self.p_rotate <= 1, "p_rotate must be between 0 and 1" assert 0 <= self.p_rotate <= 1, "p_rotate must be between 0 and 1"
assert 0 <= self.p_hmirror <= 1, "p_hmirror must be between 0 and 1" assert 0 <= self.p_hmirror <= 1, "p_hmirror must be between 0 and 1"
assert 0 <= self.p_vmirror <= 1, "p_vmirror must be between 0 and 1" assert 0 <= self.p_vmirror <= 1, "p_vmirror must be between 0 and 1"
@ -86,8 +92,8 @@ class ManipulateMatrixDataset(ProceduralDataset):
def _get_matrix(self, rng: Random) -> list[list[int]]: def _get_matrix(self, rng: Random) -> list[list[int]]:
"""Generate a random matrix""" """Generate a random matrix"""
rows = rng.randint(1, self.config.max_rows) rows = rng.randint(self.config.min_rows, self.config.max_rows)
cols = rng.randint(1, self.config.max_cols) cols = rng.randint(self.config.min_cols, self.config.max_cols)
numbers = [rng.randint(0, 9) for _ in range(rows * cols)] numbers = [rng.randint(0, 9) for _ in range(rows * cols)]
matrix = [numbers[i * cols : (i + 1) * cols] for i in range(rows)] matrix = [numbers[i * cols : (i + 1) * cols] for i in range(rows)]
return matrix return matrix
@ -157,14 +163,13 @@ class ManipulateMatrixDataset(ProceduralDataset):
matrix = self._get_matrix(rng) matrix = self._get_matrix(rng)
matrix_str = self._matrix_to_str(matrix) matrix_str = self._matrix_to_str(matrix)
# Shuffle the order of operations (make sure to copy the list to guarantee same order) num_transforms = rng.randint(0, self.config.max_transforms)
all_transforms = deepcopy(self._all_transforms) transforms = rng.sample(self._all_transforms, num_transforms)
rng.shuffle(all_transforms)
operations = [] operations = []
answer = deepcopy(matrix) answer = deepcopy(matrix)
for transform in all_transforms: for transform in transforms:
# Rotate # Rotate
if transform == "rotate" and rng.random() < self.config.p_rotate: if transform == "rotate" and rng.random() < self.config.p_rotate:
rotation = rng.choice(list(self._rotations.keys())) rotation = rng.choice(list(self._rotations.keys()))

6
tests/test_manipulate_matrix.py
View file

@ -8,8 +8,12 @@ from reasoning_gym.algorithmic.manipulate_matrix import ManipulateMatrixConfig,
def test_manipulate_matrix_config_validation(): def test_manipulate_matrix_config_validation():
"""Test that invalid configs raise appropriate errors""" """Test that invalid configs raise appropriate errors"""
with pytest.raises(AssertionError):
config = ManipulateMatrixConfig(max_transforms=-1) # max_transforms should be non-negative
config.validate()
invalid_dims = [-1, 0] # Dimensions should be positive integers invalid_dims = [-1, 0] # Dimensions should be positive integers
dim_fields = ["max_rows", "max_cols"] dim_fields = ["min_rows", "min_cols", "max_rows", "max_cols"]
for field in dim_fields: for field in dim_fields:
for dim in invalid_dims: for dim in invalid_dims: