mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
108 lines
4.3 KiB
Python
108 lines
4.3 KiB
Python
import random
|
|
from dataclasses import dataclass
|
|
from fractions import Fraction
|
|
from typing import Any, Dict, Optional
|
|
|
|
import sympy
|
|
|
|
from ..factory import ProceduralDataset, register_dataset
|
|
|
|
|
|
@dataclass
|
|
class SimpleIntegrationConfig:
|
|
min_terms: int = 2
|
|
max_terms: int = 5
|
|
min_degree: int = 1
|
|
max_degree: int = 10
|
|
min_bounds: int = 1
|
|
max_bounds: int = 10
|
|
operators: tuple = ("+", "-")
|
|
symbols: tuple = ("x", "X")
|
|
seed: Optional[int] = None
|
|
size: int = 500
|
|
|
|
def validate(self) -> None:
|
|
"""Validate the configuration parameters of the integral proble"""
|
|
assert self.min_bounds > 0, "min_bounds must be positive"
|
|
assert self.max_bounds >= self.min_bounds, "max_bounds must be >= min_bounds"
|
|
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_degree >= -10, "min_degree must be >= -10"
|
|
assert self.max_degree >= self.min_degree, "max_degree must be >= min_degree"
|
|
assert all(op in ("+", "-") for op in self.operators), "invalid operator specified"
|
|
|
|
|
|
class SimpleIntegrationDataset(ProceduralDataset):
|
|
"""Generates simple integration problems with one variable"""
|
|
|
|
def __init__(self, config: SimpleIntegrationConfig):
|
|
self._prompt_templates = [
|
|
"Find the indefinite integral: ∫ {integrand} dx",
|
|
"Calculate the antiderivative: ∫ {integrand} dx",
|
|
"Evaluate the indefinite integral: ∫ {integrand} dx",
|
|
]
|
|
super().__init__(config=config, seed=config.seed, size=config.size)
|
|
|
|
def _generate_coefficient(self, rng: random.Random) -> Fraction:
|
|
"""Generate a random coefficient for the polynomial"""
|
|
if rng.choice([True, False]): # 50% chance for integer
|
|
return Fraction(rng.randint(self.config.min_bounds, self.config.max_bounds), 1)
|
|
denominator = rng.randint(2, 10)
|
|
return Fraction(rng.randint(self.config.min_bounds, self.config.max_bounds), denominator)
|
|
|
|
def _generate_polynomial(self, rng: random.Random) -> tuple[sympy.Symbol, sympy.Expr]:
|
|
"""Generate a random polynomial with one variable"""
|
|
terms = []
|
|
x = sympy.Symbol(rng.choice(self.config.symbols))
|
|
|
|
for _ in range(rng.randint(self.config.min_terms, self.config.max_terms)):
|
|
coefficient = self._generate_coefficient(rng)
|
|
degree = rng.randint(self.config.min_degree, self.config.max_degree)
|
|
operator = rng.choice(self.config.operators)
|
|
term = coefficient * x**degree
|
|
if operator == "-":
|
|
term = -term
|
|
terms.append(term)
|
|
return x, sum(terms)
|
|
|
|
def __getitem__(self, idx: int) -> dict:
|
|
rng = random.Random(self.seed + idx)
|
|
symbol, polynomial = self._generate_polynomial(rng)
|
|
derivative = sympy.diff(polynomial, symbol)
|
|
|
|
return {
|
|
"question": rng.choice(self._prompt_templates).format(integrand=derivative),
|
|
"answer": str(polynomial) + " + C",
|
|
"metadata": {
|
|
"integrand": str(derivative),
|
|
"variable": str(symbol),
|
|
"expected_answer_expression": polynomial,
|
|
},
|
|
}
|
|
|
|
def score_answer(self, answer: Optional[str], metadata: Dict[str, Any]) -> float:
|
|
"""Determine if the solution provided solves the problem"""
|
|
reward = 0.0
|
|
if answer is not None:
|
|
try:
|
|
var = metadata["variable"]
|
|
x = sympy.Symbol(var)
|
|
# Parse answer while allowing integration constant 'C'
|
|
user_expr = sympy.parse_expr(answer, local_dict={var: x, "C": sympy.Symbol("C")})
|
|
# Compute derivative of student's answer
|
|
derivative = sympy.diff(user_expr, x)
|
|
integrand = sympy.parse_expr(metadata["integrand"], local_dict={var: x})
|
|
|
|
# Check mathematical equivalence through simplification
|
|
if sympy.simplify(derivative - integrand) == 0:
|
|
reward = 1.0
|
|
elif answer.strip():
|
|
reward = 0.05
|
|
else:
|
|
reward = 0.01
|
|
except:
|
|
reward = 0.01
|
|
return reward
|
|
|
|
|
|
register_dataset("simple_integration", SimpleIntegrationDataset, SimpleIntegrationConfig)
|