open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-04-23 16:55:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(env): Prime Factorization (#344)

Browse source 
* prime factorization

* lint
This commit is contained in:
Zafir Stojanovski 2025-03-13 21:02:20 +01:00 • committed by GitHub
parent 36a7ca1150
commit 4f45c8d655
3 changed files with 52 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

26
reasoning_gym/arithmetic/prime_factorization.py
View file

@ -5,6 +5,7 @@ from dataclasses import dataclass
from random import Random
from typing import Any, Optional
from ..coaching import AttributeType, BaseCurriculum, RangeAttributeDefinition
from ..factory import ProceduralDataset, register_dataset
@ -82,8 +83,29 @@ class PrimeFactorizationDataset(ProceduralDataset):
f"(Example: for 12 the answer would be: 2 × 2 × 3)"
),
"answer": answer,
"metadata": {"number": number, "factors": factors},
"metadata": {"number": number, "factors": factors, "difficulty": {"value": number}},
}
register_dataset("prime_factorization", PrimeFactorizationDataset, PrimeFactorizationConfig)
class PrimeFactorizationCurriculum(BaseCurriculum):
def __init__(self):
super().__init__(PrimeFactorizationCurriculum.__name__, PrimeFactorizationConfig)
# Define attributes
self._define_attributes(
RangeAttributeDefinition(
name="value",
levels=[10, 1_000, 10_000, 50_000],
default_level=1,
description="Number to factorize",
attr_type=AttributeType.APPEND,
min_value=2,
lower_field_name="min_value",
upper_field_name="max_value",
)
)
register_dataset(
"prime_factorization", PrimeFactorizationDataset, PrimeFactorizationConfig, PrimeFactorizationCurriculum
)