mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
125 lines
4.4 KiB
Python
125 lines
4.4 KiB
Python
"""Number sorting task generator"""
|
|
|
|
import random
|
|
from dataclasses import dataclass
|
|
from random import Random
|
|
from typing import List, Optional, Tuple
|
|
|
|
|
|
@dataclass
|
|
class NumberSortingConfig:
|
|
"""Configuration for number sorting task generation"""
|
|
|
|
min_numbers: int = 3 # Minimum numbers to sort
|
|
max_numbers: int = 10 # Maximum numbers to sort
|
|
min_decimals: int = 0 # Minimum decimal places
|
|
max_decimals: int = 2 # Maximum decimal places
|
|
min_value: float = -100.0 # Minimum value
|
|
max_value: float = 100.0 # Maximum value
|
|
seed: Optional[int] = None
|
|
size: int = 500 # Virtual dataset size
|
|
|
|
def validate(self):
|
|
"""Validate configuration parameters"""
|
|
assert self.min_numbers > 0, "min_numbers must be positive"
|
|
assert self.min_numbers <= self.max_numbers, "max_numbers must be >= min_numbers"
|
|
assert self.min_decimals >= 0, "min_decimals must be non-negative"
|
|
assert self.min_decimals <= self.max_decimals, "max_decimals must be >= min_decimals"
|
|
assert self.min_value < self.max_value, "max_value must be > min_value"
|
|
|
|
|
|
class NumberSortingDataset:
|
|
"""Generates number sorting tasks"""
|
|
|
|
def __init__(self, config: NumberSortingConfig):
|
|
self.config = config
|
|
self.config.validate()
|
|
self.seed = config.seed if config.seed is not None else Random().randint(0, 2**32)
|
|
|
|
def __len__(self) -> int:
|
|
return self.config.size
|
|
|
|
def __iter__(self):
|
|
self._current_idx = 0
|
|
return self
|
|
|
|
def __next__(self):
|
|
if self._current_idx >= self.config.size:
|
|
raise StopIteration
|
|
item = self[self._current_idx]
|
|
self._current_idx += 1
|
|
return item
|
|
|
|
def _format_number(self, num: float, decimals: int) -> str:
|
|
"""Format number with specified decimal places"""
|
|
formatted = f"{num:.{decimals}f}"
|
|
# Reparse to ensure exact decimal representation
|
|
return f"{float(formatted):.{decimals}f}"
|
|
|
|
def _generate_numbers(self, rng: Random) -> Tuple[List[float], List[str]]:
|
|
"""Generate list of numbers and their string representations"""
|
|
count = rng.randint(self.config.min_numbers, self.config.max_numbers)
|
|
decimals = rng.randint(self.config.min_decimals, self.config.max_decimals)
|
|
|
|
numbers = []
|
|
number_strs = []
|
|
|
|
for _ in range(count):
|
|
num = rng.uniform(self.config.min_value, self.config.max_value)
|
|
num_str = self._format_number(num, decimals)
|
|
# Reparse to ensure exact value
|
|
num = float(num_str)
|
|
numbers.append(num)
|
|
number_strs.append(num_str)
|
|
|
|
return numbers, number_strs
|
|
|
|
def __getitem__(self, idx: int) -> dict:
|
|
"""Generate a single sorting task"""
|
|
rng = Random(self.seed + idx)
|
|
|
|
numbers, number_strs = self._generate_numbers(rng)
|
|
|
|
# Generate both ascending and descending answers
|
|
asc_numbers = sorted(numbers)
|
|
desc_numbers = sorted(numbers, reverse=True)
|
|
|
|
# Format answers as string lists
|
|
decimals = len(number_strs[0].split(".")[-1]) if "." in number_strs[0] else 0
|
|
asc_answer = [self._format_number(n, decimals) for n in asc_numbers]
|
|
desc_answer = [self._format_number(n, decimals) for n in desc_numbers]
|
|
|
|
# Randomly choose ascending or descending
|
|
is_ascending = rng.choice([True, False])
|
|
direction = "ascending" if is_ascending else "descending"
|
|
answer = asc_answer if is_ascending else desc_answer
|
|
|
|
return {
|
|
"question": f"Sort these numbers in {direction} order: {', '.join(number_strs)}",
|
|
"answer": str(answer),
|
|
"metadata": {"original_numbers": number_strs, "direction": direction, "sorted_numbers": answer},
|
|
}
|
|
|
|
|
|
def number_sorting_dataset(
|
|
min_numbers: int = 3,
|
|
max_numbers: int = 10,
|
|
min_decimals: int = 0,
|
|
max_decimals: int = 2,
|
|
min_value: float = -100.0,
|
|
max_value: float = 100.0,
|
|
seed: Optional[int] = None,
|
|
size: int = 500,
|
|
) -> NumberSortingDataset:
|
|
"""Create a NumberSortingDataset with the given configuration."""
|
|
config = NumberSortingConfig(
|
|
min_numbers=min_numbers,
|
|
max_numbers=max_numbers,
|
|
min_decimals=min_decimals,
|
|
max_decimals=max_decimals,
|
|
min_value=min_value,
|
|
max_value=max_value,
|
|
seed=seed,
|
|
size=size,
|
|
)
|
|
return NumberSortingDataset(config)
|