mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2025-10-09 13:40:09 +03:00
d0ef136d5b
* added curriculum * readapted readme * corrected small errors * Delete eval/eval/r1/algorithmic/word_sorting.json * removed redundant argument * added spell * removed duplicated fit * changed config * added composite changes * added composite changes * updated yaml * added spell backward * updated read me * added qwen2.5 * added * Add files via upload * updated missing trainer func * updated curr * updated spell back * updated correctness score func * updated configs * added local evals * added updates * updated datasets * added fsdp to hf utility * added algorithmic qwen 3b yaml * updated read me * updated configs * added preappend token * updated with thinking token * updated test score board * resolved comments * added evaluation scripts * removed results from pr * added config * added partial reward scoring * added evaluation composites * added training configs * added games eval * added rubriks cube * resolved merge cinflicts * added games config * added latest eval configs * updated strucutre * Delete training/evaluations/eval_graphs_composite.yaml --------- Co-authored-by: joesharratt1229 <joesharrat1229@gmail.com>
272 lines
10 KiB
Python
272 lines
10 KiB
Python
from dataclasses import dataclass
|
|
from random import Random
|
|
from typing import Any, Literal, Optional
|
|
|
|
from ..coaching import BaseCurriculum, RangeAttributeDefinition, ScalarAttributeDefinition
|
|
from ..factory import ProceduralDataset, register_dataset
|
|
|
|
DATASET_NAME = "basic_arithmetic"
|
|
|
|
|
|
@dataclass
|
|
class BasicArithmeticDatasetConfig:
|
|
"""Configuration for arithmetic dataset generation"""
|
|
|
|
min_terms: int = 2
|
|
max_terms: int = 6
|
|
min_digits: int = 1
|
|
max_digits: int = 4
|
|
operators: list[str] = ("+", "-", "*", "/")
|
|
allow_parentheses: bool = True
|
|
allow_negation: bool = True
|
|
seed: Optional[int] = None
|
|
size: int = 500 # Virtual dataset size
|
|
format_style: Literal["simple", "natural"] = "simple"
|
|
whitespace: Literal["no_space", "single", "random"] = "single" # Whitespace style between terms
|
|
|
|
def validate(self) -> None:
|
|
"""Validate configuration parameters"""
|
|
assert self.min_terms > 0, "min_terms must be positive"
|
|
assert self.max_terms >= self.min_terms, "max_terms must be >= min_terms"
|
|
assert self.min_digits > 0, "min_digits must be positive"
|
|
assert self.max_digits >= self.min_digits, "max_digits must be >= min_digits"
|
|
assert len(self.operators) > 0, "must provide at least one operator"
|
|
for op in self.operators:
|
|
assert op in ["+", "-", "*", "/"], f"unsupported operator: {op}"
|
|
|
|
|
|
def find_common_divisors(a: int, b: int) -> list[int]:
|
|
# Helper function to find GCD using Euclidean algorithm
|
|
def gcd(x, y):
|
|
while y:
|
|
x, y = y, x % y
|
|
return x
|
|
|
|
# Get the GCD of the two numbers
|
|
gcd_value = gcd(abs(a), abs(b))
|
|
# Find all divisors of the GCD
|
|
divisors = []
|
|
i = 1
|
|
# We only need to check up to sqrt(gcd_value)
|
|
while i * i <= gcd_value:
|
|
if gcd_value % i == 0:
|
|
divisors.append(i)
|
|
# Don't add the same number twice for perfect squares
|
|
if i * i != gcd_value:
|
|
divisors.append(gcd_value // i)
|
|
i += 1
|
|
return divisors
|
|
|
|
|
|
def eval_floordiv(exp: str) -> int:
|
|
return eval(exp.replace("/", "//"))
|
|
|
|
|
|
class BasicArithmeticDataset(ProceduralDataset):
|
|
"""Dataset that generates basic arithmetic tasks with configurable complexity"""
|
|
|
|
def __init__(self, config: BasicArithmeticDatasetConfig):
|
|
super().__init__(config=config, seed=config.seed, size=config.size)
|
|
|
|
def __getitem__(self, idx: int) -> dict[str, Any]:
|
|
"""Generate a single arithmetic task
|
|
|
|
Args:
|
|
idx: Index of the item to generate
|
|
|
|
Returns:
|
|
dict with keys:
|
|
- question: str, the formatted arithmetic expression
|
|
- answer: str, the ground truth result
|
|
- metadata: dict with generation parameters
|
|
"""
|
|
# Create deterministic RNG from base seed and idx
|
|
rng = Random(self.seed + idx)
|
|
|
|
num_terms = rng.randint(self.config.min_terms, self.config.max_terms)
|
|
num_digits = rng.randint(self.config.min_digits, self.config.max_digits)
|
|
|
|
if self.config.allow_parentheses:
|
|
expression, result = self._generate_complex_task(rng, num_terms, num_digits)
|
|
else:
|
|
expression, result = self._generate_simple_task(rng, num_terms, num_digits)
|
|
|
|
question = self._format_question(rng, expression)
|
|
|
|
return {
|
|
"question": question,
|
|
"answer": str(result),
|
|
"metadata": {
|
|
"source_dataset": DATASET_NAME,
|
|
"source_index": idx,
|
|
"expression": expression,
|
|
"num_terms": num_terms,
|
|
"num_digits": num_digits,
|
|
"difficulty": {
|
|
"num_terms": (self.config.min_terms, self.config.max_terms),
|
|
"num_digits": (self.config.min_digits, self.config.max_digits),
|
|
},
|
|
},
|
|
}
|
|
|
|
def _generate_complex_task(self, rng: Random, num_terms: int, num_digits: int) -> tuple[str, int]:
|
|
"""Generate a complex arithmetic task with possible parentheses"""
|
|
|
|
def add_terms(remaining: int) -> list[str]:
|
|
# split terms randomly into left and right
|
|
num_left = rng.randint(1, remaining)
|
|
num_right = remaining - num_left
|
|
|
|
left_parts = []
|
|
if num_left > 1 and rng.random() > 0.5 and self.config.allow_parentheses:
|
|
if rng.random() > 0.5 and self.config.allow_negation:
|
|
left_parts.append("-(")
|
|
else:
|
|
left_parts.append("(")
|
|
left_parts.extend(add_terms(num_left))
|
|
left_parts.append(")")
|
|
else:
|
|
for i in range(num_left):
|
|
c = rng.randint(-(10**num_digits) + 1, 10**num_digits - 1)
|
|
left_parts.append(str(c))
|
|
if i + 1 < num_left:
|
|
left_parts.append(rng.choice([o for o in self.config.operators if o != "/"]))
|
|
|
|
if num_right == 0:
|
|
return left_parts
|
|
|
|
op = rng.choice(self.config.operators)
|
|
if op != "/":
|
|
left_parts.append(op)
|
|
left_parts.extend(add_terms(num_right))
|
|
else:
|
|
# left part has parantheses or no division
|
|
dividend = eval_floordiv("".join(left_parts) if left_parts[-1] == ")" else left_parts[-1])
|
|
left_parts.append(op)
|
|
|
|
if num_right > 1:
|
|
right_parts = add_terms(num_right - 1)
|
|
if right_parts[-1] == ")":
|
|
right_value = eval_floordiv("".join(right_parts))
|
|
|
|
if right_value == 0:
|
|
correction = 1
|
|
else:
|
|
target = rng.choice(find_common_divisors(dividend, right_value))
|
|
correction = target - right_value
|
|
|
|
right_parts.pop()
|
|
right_parts.append("+")
|
|
right_parts.append(str(correction))
|
|
right_parts.append(")")
|
|
|
|
else:
|
|
if dividend != 0:
|
|
divisor = rng.choice(find_common_divisors(dividend, 0))
|
|
else:
|
|
divisor = rng.randint(1, 10**num_digits - 1)
|
|
left_parts.append(str(divisor))
|
|
left_parts.append("+")
|
|
left_parts.extend(right_parts)
|
|
else:
|
|
if dividend != 0:
|
|
divisor = rng.choice(find_common_divisors(dividend, 0))
|
|
else:
|
|
divisor = rng.randint(1, 10**num_digits - 1)
|
|
left_parts.append(str(divisor))
|
|
|
|
return left_parts
|
|
|
|
parts = add_terms(num_terms)
|
|
|
|
# Add whitespace according to config
|
|
if self.config.whitespace == "no_space":
|
|
expression = "".join(parts)
|
|
elif self.config.whitespace == "single":
|
|
expression = " ".join(parts)
|
|
else: # random
|
|
space_parts = []
|
|
for p in parts:
|
|
if rng.random() < 0.15:
|
|
space_parts.append(" ")
|
|
space_parts.append(p)
|
|
expression = "".join(space_parts).strip()
|
|
result = eval_floordiv(expression) # Note: eval is safe here as we control the input
|
|
|
|
return expression, result
|
|
|
|
def _generate_simple_task(self, rng: Random, num_terms: int, num_digits: int) -> tuple[str, int]:
|
|
"""Generate a simple linear arithmetic task without parentheses"""
|
|
constants = [rng.randint(0, 10**num_digits) for _ in range(num_terms)]
|
|
operators = [rng.choice(self.config.operators) for _ in range(num_terms - 1)]
|
|
|
|
# Build expression and compute result
|
|
expression_parts = []
|
|
result = constants[0]
|
|
|
|
expression_parts.append(str(constants[0]))
|
|
for i, op in enumerate(operators):
|
|
c = constants[i + 1]
|
|
expression_parts.append(op)
|
|
expression_parts.append(str(c))
|
|
|
|
if op == "+":
|
|
result += c
|
|
elif op == "-":
|
|
result -= c
|
|
elif op == "*":
|
|
result *= c
|
|
elif op == "/":
|
|
# Find a number that divides result evenly
|
|
divisors = [d for d in range(2, min(abs(result), 10**num_digits)) if result % d == 0]
|
|
if divisors:
|
|
c = rng.choice(divisors)
|
|
result //= c
|
|
else:
|
|
# Fallback to multiplication if no clean division possible
|
|
op = "*"
|
|
c = rng.randint(1, 10**num_digits - 1)
|
|
result *= c
|
|
else:
|
|
raise RuntimeError(f"Unsupported operator: {op}")
|
|
|
|
expression = " ".join(expression_parts)
|
|
return expression, result
|
|
|
|
def _format_question(self, rng: Random, expression: str) -> str:
|
|
"""Format the the question with the arithmetic expression"""
|
|
|
|
if self.config.format_style == "simple":
|
|
return f"Calculate {expression}."
|
|
else:
|
|
templates = ["What is {0}?", "Solve {0}.", "Compute {0}.", "Evaluate: {0}."]
|
|
template = rng.choice(templates)
|
|
return template.format(expression)
|
|
|
|
|
|
class BasicArithmeticCurriculum(BaseCurriculum):
|
|
def __init__(self):
|
|
super().__init__(name=BasicArithmeticCurriculum.__name__, config_cls=BasicArithmeticDatasetConfig)
|
|
self._define_attributes(
|
|
RangeAttributeDefinition(
|
|
name="num_terms",
|
|
levels=[2, 3, 4, 5, 6],
|
|
description="Number of terms in the expression",
|
|
lower_field_name="min_terms",
|
|
upper_field_name="max_terms",
|
|
ensure_interval=False,
|
|
),
|
|
RangeAttributeDefinition(
|
|
name="num_digits",
|
|
levels=[1, 2, 3, 4],
|
|
description="Number of digits in the numbers",
|
|
lower_field_name="min_digits",
|
|
upper_field_name="max_digits",
|
|
ensure_interval=False,
|
|
),
|
|
)
|
|
|
|
|
|
# Register the dataset
|
|
register_dataset(DATASET_NAME, BasicArithmeticDataset, BasicArithmeticDatasetConfig, BasicArithmeticCurriculum)
|