mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-23 16:55:05 +00:00
fix(envs): Add source dataset and index to metadata (#388)
* add source dataset and index to metadata * fix typo * fix coach class and its test
This commit is contained in:
parent
c6d01541aa
commit
4c47527130
104 changed files with 549 additions and 146 deletions
|
|
@ -8,6 +8,8 @@ from typing import Any, Optional
|
|||
from ..coaching import BaseCurriculum, RangeAttributeDefinition
|
||||
from ..factory import ProceduralDataset, register_dataset
|
||||
|
||||
DATASET_NAME = "prime_factorization"
|
||||
|
||||
|
||||
@dataclass
|
||||
class PrimeFactorizationConfig:
|
||||
|
|
@ -84,6 +86,8 @@ class PrimeFactorizationDataset(ProceduralDataset):
|
|||
),
|
||||
"answer": answer,
|
||||
"metadata": {
|
||||
"source_dataset": DATASET_NAME,
|
||||
"source_index": idx,
|
||||
"number": number,
|
||||
"factors": factors,
|
||||
"difficulty": {
|
||||
|
|
@ -110,6 +114,4 @@ class PrimeFactorizationCurriculum(BaseCurriculum):
|
|||
)
|
||||
|
||||
|
||||
register_dataset(
|
||||
"prime_factorization", PrimeFactorizationDataset, PrimeFactorizationConfig, PrimeFactorizationCurriculum
|
||||
)
|
||||
register_dataset(DATASET_NAME, PrimeFactorizationDataset, PrimeFactorizationConfig, PrimeFactorizationCurriculum)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue