Function run_dataset

examples/run_aime_reflection_experiment.py:483–531

(
    *,
    dataset_name: str,
    engine: tg.EngineLM,
    model: Module,
    model_name: str,
    max_step: int,
    max_samples: Optional[int],
    result_slot: Dict[str, Any],
    results_root: Dict[str, Any],
    output_path: Path,
)

def run_dataset(
    *,
    dataset_name: str,
    engine: tg.EngineLM,
    model: Module,
    model_name: str,
    max_step: int,
    max_samples: Optional[int],
    result_slot: Dict[str, Any],
    results_root: Dict[str, Any],
    output_path: Path,
) -> Dict[str, Any]:
    dataset_cls = DATASET_REGISTRY[dataset_name]
    dataset = dataset_cls(split="all")

    evaluation_module = build_evaluation_module(engine)

    samples: List[Dict[str, Any]] = result_slot.setdefault("samples", [])
    result_slot["metrics"] = compute_metrics(samples)
    save_results(results_root, output_path)

    total_available = len(dataset)
    target_count = total_available if max_samples is None else min(total_available, max_samples)
    already_processed = min(len(samples), target_count)
    start_idx = already_processed
    end_idx = target_count
    progress_desc = f"{model_name} · {dataset_name}"

    if start_idx >= end_idx:
        return result_slot

    for idx in tqdm(range(start_idx, end_idx), desc=progress_desc, unit="sample"):
        question_text, ground_truth_text = dataset[idx]
        sample_log = run_sample(
            question_text=question_text,
            ground_truth_text=str(ground_truth_text),
            model=model,
            evaluation_module=evaluation_module,
            engine=engine,
            model_name=model_name,
            max_step=max_step,
        )
        samples.append(sample_log)
        result_slot["metrics"] = compute_metrics(samples)
        save_results(results_root, output_path)

    result_slot["metrics"] = compute_metrics(samples)
    save_results(results_root, output_path)
    return result_slot


def compute_metrics(samples: List[Dict[str, Any]]) -> Dict[str, Any]: