fix(data): stabilize mixed-modality SFT schema

[yangtao-hub]

What does this PR do?

Fixes #7266.
Fixes #9195.

Related to #5613 and #6223.

This PR stabilizes non-streaming, non-packed SFT preprocessing for datasets that mix text-only, image-text, video-text, and audio-text samples.

Hugging Face Datasets infers the Arrow schema from processed batches. If an early batch contains no media, optional media columns may be inferred as null. A later batch containing media then fails with:

TypeError: Couldn't cast array of type list<item: string> to null

The fix provides an explicit output schema to Dataset.map() for the standard supervised preprocessing path. Token columns use stable numeric types, while images, videos, and audios inherit their Arrow types from the aligned input dataset. Inheriting these types avoids narrowing supported media representations to string paths.

Streaming, packed SFT, generation-based evaluation, and other training stages retain their existing preprocessing behavior.

Validation

• Added a regression test in which text-only samples are processed before image and video samples.
• Verified that the output media columns retain stable Arrow types across batches while preserving missing media values.
• Verified the schema behavior with datasets==4.0.0.
• Ran Ruff lint and formatting checks on the modified files.
• Validated the underlying explicit-schema approach in a mixed text, image, and video SFT workload.

Before submitting

• Did you read the contributor guideline?
• Did you write any new necessary tests?

[gemini-code-assist]

Code Review

This pull request introduces a helper function _get_sft_dataset_features to enforce stable output features during non-packed supervised fine-tuning (SFT) preprocessing, alongside a unit test verifying its behavior on mixed-modality batches. Feedback suggests making the helper function more robust by dynamically checking for the presence of multimodal keys (_images, _videos, _audios) in the dataset features to prevent potential KeyError exceptions on text-only datasets.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If a custom dataset or pipeline is used where _images, _videos, or _audios are not present in the dataset features (e.g., a text-only dataset that hasn't been aligned with these specific columns), accessing dataset.features["_images"] directly will raise a KeyError.

To make this robust and prevent potential crashes, we should dynamically check if these keys exist in dataset.features before adding them to the output features.

Suggested change

	def _get_sft_dataset_features(dataset: "Dataset") -> Features:
	r"""Return stable output features for non-packed supervised preprocessing."""
	return Features(
	{
	"input_ids": Sequence(Value("int32")),
	"attention_mask": Sequence(Value("int8")),
	"labels": Sequence(Value("int64")),
	"images": dataset.features["_images"],
	"videos": dataset.features["_videos"],
	"audios": dataset.features["_audios"],
	}
	)
	def _get_sft_dataset_features(dataset: "Dataset") -> Features:
	r"""Return stable output features for non-packed supervised preprocessing."""
	features = {
	"input_ids": Sequence(Value("int32")),
	"attention_mask": Sequence(Value("int8")),
	"labels": Sequence(Value("int64")),
	}
	dataset_features = getattr(dataset, "features", None) or {}
	for key, target in [("_images", "images"), ("_videos", "videos"), ("_audios", "audios")]:
	if key in dataset_features:
	features[target] = dataset_features[key]
	return Features(features)

[yangtao-hub]

This helper is only used after align_dataset, whose built-in converters always emit _images, _videos, and _audios, using None for absent modalities. The supervised processor also always returns the corresponding output keys, so omitting them from the explicit Features would make the schema inconsistent with the map output.

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

This PR stabilizes Hugging Face datasets feature schemas during non-packed SFT preprocessing to avoid modality columns being inferred inconsistently across mixed-modality batches.

Changes:

• Add _get_sft_dataset_features() helper to force stable output Features during Dataset.map for SFT (non-packed, non-generate eval).
• Apply the stable features override in the SFT preprocessing path.
• Add a regression test ensuring mixed-modality batches keep consistent images/videos/audios feature types and values.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
tests/data/test_loader.py	Adds a regression test covering stable features across mixed-modality batches.
src/llamafactory/data/loader.py	Introduces and applies stable Features for SFT preprocessing outputs.

[yangtao-hub]

This preprocessing path does not call Dataset.with_format("torch") or set_format("torch"). Dataset rows remain Python lists and are converted by DataCollatorForSeq2Seq, which produces torch.long tensors for token IDs. Keeping the Arrow representation as int32 therefore reduces storage usage without changing the model input dtype.

[yangtao-hub]

These columns are guaranteed by the internal aligned-dataset contract: all built-in converters emit _images, _videos, and _audios, using None for absent modalities. Since the supervised processor always outputs the corresponding keys, direct access intentionally enforces this invariant. I will keep the current implementation.

[yangtao-hub]

Updated the test to use list comprehensions so each sample has an independent list.