mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-29 17:35:16 +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
7475a20700
commit
ce0a6c4878
104 changed files with 549 additions and 146 deletions
|
|
@ -14,6 +14,8 @@ from ..factory import ProceduralDataset, register_dataset
|
|||
|
||||
QUESTION_TEMPLATE = """Count how many prime numbers there are between {start} and {end} (inclusive) ?"""
|
||||
|
||||
DATASET_NAME = "count_primes"
|
||||
|
||||
|
||||
@dataclass
|
||||
class CountPrimesConfig:
|
||||
|
|
@ -60,6 +62,8 @@ class CountPrimesDataset(ProceduralDataset):
|
|||
"question": QUESTION_TEMPLATE.format(start=start, end=end),
|
||||
"answer": str(answer),
|
||||
"metadata": {
|
||||
"source_dataset": DATASET_NAME,
|
||||
"source_index": idx,
|
||||
"start": start,
|
||||
"end": end,
|
||||
"primes": primes,
|
||||
|
|
@ -88,4 +92,4 @@ class CountPrimesCurriculum(BaseCurriculum):
|
|||
)
|
||||
|
||||
|
||||
register_dataset("count_primes", CountPrimesDataset, CountPrimesConfig, CountPrimesCurriculum)
|
||||
register_dataset(DATASET_NAME, CountPrimesDataset, CountPrimesConfig, CountPrimesCurriculum)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue