sentence reordering curriculum (#326)

reasoning_gym/algorithmic/__init__.py

@@ -34,7 +34,7 @@ from .pool_matrix import PoolMatrixConfig, PoolMatrixCurriculum, PoolMatrixDatas
 from .ransom_note import RansomNoteConfig, RansomNoteCurriculum, RansomNoteDataset
 from .rotate_matrix import RotateMatrixConfig, RotateMatrixCurriculum, RotateMatrixDataset
 from .rotten_oranges import RottenOrangesConfig, RottenOrangesCurriculum, RottenOrangesDataset
-from .sentence_reordering import SentenceReorderingConfig, SentenceReorderingDataset
+from .sentence_reordering import SentenceReorderingConfig, SentenceReorderingCurriculum, SentenceReorderingDataset
 from .spell_backward import SpellBackwardConfig, SpellBackwardDataset
 from .spiral_matrix import SpiralMatrixConfig, SpiralMatrixCurriculum, SpiralMatrixDataset
 from .string_insertion import StringInsertionConfig, StringInsertionCurriculum, StringInsertionDataset
@@ -77,6 +77,7 @@ __all__ = [
     "NumberSortingCurriculum",
     "SentenceReorderingConfig",
     "SentenceReorderingDataset",
+    "SentenceReorderingCurriculum",
     "WordSequenceReversalConfig",
     "WordSequenceReversalDataset",
     "WordSequenceReversalCurriculum",

reasoning_gym/algorithmic/sentence_reordering.py

@@ -5,6 +5,7 @@ from dataclasses import dataclass
 from random import Random
 from typing import Any, Optional
 
+from ..coaching import AttributeType, BaseCurriculum, RangeAttributeDefinition
 from ..data import read_data_file
 from ..factory import ProceduralDataset, register_dataset
 
@@ -89,7 +90,7 @@ class SentenceReorderingDataset(ProceduralDataset):
         return {
             "question": f"Restore the correct order of words in the following sentence: {question}",
             "answer": solved_sentence,
-            "metadata": {"word_count": word_count},
+            "metadata": {"word_count": word_count, "difficulty": {"words_in_sentence": word_count}},
         }
 
     def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:
@@ -114,4 +115,25 @@ class SentenceReorderingDataset(ProceduralDataset):
         return reward
 
 
-register_dataset("sentence_reordering", SentenceReorderingDataset, SentenceReorderingConfig)
+class SentenceReorderingCurriculum(BaseCurriculum):
+    def __init__(self):
+        super().__init__(SentenceReorderingCurriculum.__name__, SentenceReorderingConfig)
+
+        # Define attributes
+        self._define_attributes(
+            RangeAttributeDefinition(
+                name="words_in_sentence",
+                levels=[5, 20, 50, 100],
+                default_level=1,
+                description="Number of words in the sentence",
+                attr_type=AttributeType.APPEND,
+                min_value=3,
+                lower_field_name="min_words_in_sentence",
+                upper_field_name="max_words_in_sentence",
+            ),
+        )
+
+
+register_dataset(
+    "sentence_reordering", SentenceReorderingDataset, SentenceReorderingConfig, SentenceReorderingCurriculum
+)