diff --git a/reasoning_gym/code/__init__.py b/reasoning_gym/code/__init__.py index 36525df8..89d039c1 100644 --- a/reasoning_gym/code/__init__.py +++ b/reasoning_gym/code/__init__.py @@ -2,7 +2,14 @@ Code reasing tasks """ -from .bf import BFConfig, BFDataset -from .codeio import CodeIOConfig, CodeIODataset +from .bf import BFConfig, BFCurriculum, BFDataset +from .codeio import CodeIOConfig, CodeIOCurriculum, CodeIODataset -__all__ = ["BFConfig", "BFDataset", "CodeIOConfig", "CodeIODataset"] +__all__ = [ + "BFConfig", + "BFDataset", + "BFCurriculum", + "CodeIOConfig", + "CodeIODataset", + "CodeIOCurriculum", +] diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index d8c36d51..390b22a8 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -6,6 +6,7 @@ from typing import Any, Optional import zss +from ..coaching import BaseCurriculum, ScalarAttributeDefinition from ..data import get_data_file_path from ..factory import ProceduralDataset, register_dataset @@ -59,10 +60,13 @@ class CodeIOConfig: seed: Optional[int] = None size: int = 500 input_prediction_probability: float = 0.5 + difficulty: Optional[int] = None def validate(self) -> None: """Validate configuration parameters""" assert 0.0 <= self.input_prediction_probability <= 1.0, "input_prediction_probability must be in [0, 1]" + if self.difficulty is not None: + assert 1 <= self.difficulty <= 10, "difficulty must be in [1, 10]" class CodeIODataset(ProceduralDataset): @@ -80,18 +84,30 @@ class CodeIODataset(ProceduralDataset): self._data_path = get_data_file_path("codeio.jsonl.gz") with gzip.open(self._data_path, "rt", encoding="utf-8") as f: - CodeIODataset._jsonl_data = [json.loads(line) for line in f] + data = [json.loads(line) for line in f] + if self.config.difficulty is not None: + data = [entry for entry in data if entry.get("difficulty", -1) == self.config.difficulty] + assert len(data) > 0, "No data found for the specified difficulty level" + CodeIODataset._jsonl_data = data + + def _generate_io_pair(self, main_code: str, input_generator_code: str, rng: Random, max_retries: int = 1): + local_vars = {"Random": Random} + + full_code = f"{main_code}\n\n{input_generator_code}" + try: + exec(full_code, local_vars, local_vars) + except Exception as e: + print(f"Error executing code:\n{full_code}") + print(f"---------------------\nException: {e}\n---------------------") + return {}, {} - def _generate_io_pair(self, main_code: str, input_generator_code: str, rng: Random, max_retries: int = 3): - local_vars = {} - exec(main_code, {"Random": Random}, local_vars) - exec(input_generator_code, {"Random": Random}, local_vars) for _ in range(max_retries): try: inputs = local_vars["generate_inputs"](rng) outputs = local_vars["main_solution"](**inputs) - except Exception: + except Exception as e: # Retry + print(f"Error generating I/O pair: {e}") continue return inputs, outputs return {}, {} @@ -124,6 +140,7 @@ class CodeIODataset(ProceduralDataset): "source_index": idx, "input_data": input_data, "output_data": output_data, + "difficulty": {"difficulty": self.config.difficulty}, }, } @@ -237,5 +254,19 @@ class CodeIODataset(ProceduralDataset): return reward +class CodeIOCurriculum(BaseCurriculum): + def __init__(self): + super().__init__(CodeIOCurriculum.__name__, CodeIOConfig) + + self._define_attributes( + ScalarAttributeDefinition( + name="difficulty", + field_name="difficulty", + levels=[6, 7, 8, 9], + description="Difficulty level of the task", + ), + ) + + # Register the dataset -register_dataset(DATASET_NAME, CodeIODataset, CodeIOConfig) +register_dataset(DATASET_NAME, CodeIODataset, CodeIOConfig, CodeIOCurriculum) diff --git a/reasoning_gym/data/codeio.jsonl.gz b/reasoning_gym/data/codeio.jsonl.gz index e4d0891b..960e17cf 100644 Binary files a/reasoning_gym/data/codeio.jsonl.gz and b/reasoning_gym/data/codeio.jsonl.gz differ diff --git a/tests/test_codeio.py b/tests/test_codeio.py index 34478d9a..b7544e6f 100644 --- a/tests/test_codeio.py +++ b/tests/test_codeio.py @@ -1,6 +1,6 @@ import pytest -from reasoning_gym.code.codeio import CodeIOConfig, CodeIODataset +from reasoning_gym.code.codeio import CodeIOConfig, CodeIOCurriculum, CodeIODataset def test_codeio_dataset(): @@ -40,3 +40,24 @@ def test_codeio_config(): CodeIOConfig(size=10, seed=42, input_prediction_probability=0.1).validate() CodeIOConfig(size=10, seed=42, input_prediction_probability=0.9).validate() + + +def test_codeio_curriculum(): + curriculum = CodeIOCurriculum() + + base_value = {"size": 150, "seed": 1} + + base_cfg: CodeIOConfig = curriculum.generate_configuration(base_value) + assert base_cfg.seed == 1 + assert base_cfg.size == 150 + assert base_cfg.difficulty == 6 + + # test incrementing attribute level + curriculum.increment_attr_level("difficulty") + increased_cfg = curriculum.generate_configuration(base_value) + assert increased_cfg.difficulty == 7 + + # test decrementing attribute level + curriculum.decrement_attr_level("difficulty") + partially_decreased_cfg = curriculum.generate_configuration(base_value) + assert partially_decreased_cfg.difficulty == 6