Support DatasetDict and IterableDatasetDict as eval_dataset in evaluate

[albertvillanova]

This PR extends evaluate() in the stable trainers (DPO, KTO, Reward, SFT) to formally accept a DatasetDict or IterableDatasetDict as eval_dataset, mirroring the support already added to __init__ in:

• Support DatasetDict and IterableDatasetDict as eval_dataset in trainers #6322

Motivation

load_dataset returns a DatasetDict (or an IterableDatasetDict when streaming) whenever no split is specified, so users commonly hold these types and pass them straight to evaluate() for a held-out evaluation, e.g. trainer.evaluate(eval_dataset=load_dataset(...)). Because both subclass dict, this already worked at runtime in each trainer's evaluate(), since it dispatches per split via isinstance(eval_dataset, dict). But the type annotation only declared a plain dict, so the support was implicit and untested.

Solution

The eval_dataset parameter of each trainer's evaluate() now accepts DatasetDict and IterableDatasetDict explicitly, in its type annotation. Each split is still prepared independently, using the same _prepare_dataset path as passing a single dataset.

Changes

• Import DatasetDict and IterableDatasetDict in the DPO, KTO, Reward, and SFT trainers
• Extend the eval_dataset type annotation of each trainer's evaluate() method
• Add tests verifying a DatasetDict and, where supported, an IterableDatasetDict eval dataset can be passed directly to evaluate() and are prepared per split

---

Note

Low Risk
Annotation and test-only alignment with existing dict-based eval preparation; no new eval logic paths.

Overview
DPO, KTO, Reward, and SFT evaluate() now documents the same eval_dataset shapes as trainer __init__: single Dataset / IterableDataset, multi-split DatasetDict / IterableDatasetDict, or a plain dict of splits. Passing something like load_dataset(...) without a split is an intended use case; each named split is still prepared via _prepare_dataset and reported as eval_<split>_loss.

Runtime handling is unchanged (DatasetDict subclasses dict, so per-split preparation already worked). The change is mainly type hints aligned with #6322, plus new tests (map vs streaming) on all four trainers so multi-split evaluate() is covered end-to-end.

^{Reviewed by Cursor Bugbot for commit fc3067f. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[albertvillanova]

Failing test is addressed in:

• Fix KTO evaluate crash on streaming datasets with unpaired-preference data #6325

[qgallouedec]

thanks, we get a CI error, maybe we should bump the datasets version?

[qgallouedec]

Ok you're addressing it in #6326 (comment)!

[albertvillanova]

The KTO streaming test (test_evaluate_with_eval_dataset_dict[True]) surfaced a separate, pre-existing bug in KTOTrainer's unpaired-preference collators: "label" was returned as a plain list instead of a tensor, which crashes Accelerate's dispatch for IterableDataset.

That's tracked and fixed independently in (this PR needs merging first):

• Fix KTO evaluate crash on streaming datasets with unpaired-preference data #6325

[qgallouedec]

thanks, lgtm!

[qgallouedec]

ultra nit, can be addressed later: we could have something like a shared AnyDataset = Dataset | IterableDataset | DatasetDict | IterableDatasetDict | dict[str, Dataset | IterableDataset]

[albertvillanova]

Follow-up PR.