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:
Zafir Stojanovski 2025-03-20 12:12:14 +01:00 committed by GitHub
parent c6d01541aa
commit 4c47527130
104 changed files with 549 additions and 146 deletions

View file

@ -16,6 +16,8 @@ PSO_DIFFICULTY_RANGES = [
(PSO_DIFFICULTY_LEVELS[i], PSO_DIFFICULTY_LEVELS[i + 1]) for i in range(len(PSO_DIFFICULTY_LEVELS) - 1)
]
DATASET_NAME = "rearc"
@dataclass
class ReArcConfig:
@ -114,6 +116,8 @@ class ReArcDataset(ProceduralDataset):
"question": input_prompt,
"answer": answer,
"metadata": {
"source_dataset": DATASET_NAME,
"source_index": idx,
"input": task["input"],
"output": task["output"],
"task_id": task_id,
@ -178,4 +182,4 @@ class ReArcCurriculum(BaseCurriculum):
)
register_dataset("rearc", ReArcDataset, ReArcConfig, ReArcCurriculum)
register_dataset(DATASET_NAME, ReArcDataset, ReArcConfig, ReArcCurriculum)