Add SFT vision COCO captions dataset and config support

[subawocit]

Description

This PR adds support for COCO-caption dataset for Supervised Fine-Tuning.

Files and Changes

• src/maxtext/input_pipeline/hf_data_processing.py: Added add_default_prompt_if_missing to automatically inject default_prompt when the specified prompt column is missing in dataset features (e.g. COCO captioning).
• src/maxtext/configs/base.yml and src/maxtext/configs/types.py: Registered default_prompt.
• src/maxtext/configs/post_train/sft-vision-coco-captions.yml: Added configuration for COCO-caption dataset.
• tests/unit/hf_data_processing_test.py: Added test_add_default_prompt_if_missing unit test verifying prompt injection when missing and preservation when present.

Tests

Tested end-to-end multimodal SFT and decode pipeline on Gemma 3 4B:

1. Pre-SFT Baseline Inference

python3 -m maxtext.inference.decode \
    src/maxtext/configs/base.yml \
    model_name=gemma3-4b \
    tokenizer_path=google/gemma-3-4b-it \
    tokenizer_type=huggingface \
    load_parameters_path=gs://yuchenhou-maxtext-logs/checkpoints/gemma3-4b-processor/0/items \
    per_device_batch_size=1 \
    run_name=gemma3_4b_coco_pre_sft \
    max_prefill_predict_length=1024 \
    max_target_length=2048 \
    steps=1 \
    async_checkpointing=false \
    scan_layers=false \
    use_multimodal=true \
    prompt='Describe the image concisely' \
    image_path=tests/assets/test_image.jpg \
    attention=dot_product \
    skip_jax_distributed_system=true

Output:

Input `<start_of_turn>user
<start_of_image>Describe the image concisely<end_of_turn>
<start_of_turn>model
` -> `Here's a concise description of the image:

The image shows a panoramic view of Seattle, Washington. The iconic Space Needle dominates the skyline, with numerous other buildings visible. In the distance, snow-capped mountains are partially obscured by clouds.`

---

2. Vision SFT Training on COCO

python3 -m maxtext.trainers.post_train.sft.train_sft_native \
    src/maxtext/configs/post_train/sft-vision-coco-captions.yml \
    model_name=gemma3-4b \
    tokenizer_path=google/gemma-3-4b-it \
    tokenizer_type=huggingface \
    load_parameters_path=gs://yuchenhou-maxtext-logs/checkpoints/gemma3-4b-processor/0/items \
    base_output_directory=gs://yuchenhou-maxtext-logs/gemma3-4b/multimodal/sft \
    run_name=gemma3_4b_sft_coco \
    per_device_batch_size=1 \
    max_prefill_predict_length=1024 \
    max_target_length=2048 \
    steps=101 \
    checkpoint_period=20 \
    scan_layers=false \
    async_checkpointing=false \
    save_checkpoint_on_completion=true \
    log_period=1 \
    eval_interval=10 \
    eval_steps=2 \
    gcs_metrics=true \
    float32_qk_product=true \
    float32_logits=true \
    attention=dot_product \
    sharding_tolerance=0.05 \
    checkpoint_storage_use_zarr3=false \
    checkpoint_storage_use_ocdbt=false \
    enable_single_controller=false \
    grain_worker_count=0

Verified locally that both train and val losses are decreasing and no errors detected.

---

3. Post-SFT Inference Verification

python3 -m maxtext.inference.decode \
    src/maxtext/configs/base.yml \
    model_name=gemma3-4b \
    tokenizer_path=google/gemma-3-4b-it \
    tokenizer_type=huggingface \
    load_parameters_path=gs://yuchenhou-maxtext-logs/gemma3-4b/multimodal/sft/gemma3_4b_sft_coco/checkpoints/100/items \
    per_device_batch_size=1 \
    run_name=gemma3_4b_coco_post_sft \
    max_prefill_predict_length=1024 \
    max_target_length=1600 \
    steps=1 \
    async_checkpointing=false \
    scan_layers=false \
    use_multimodal=true \
    prompt='Describe the image concisely.' \
    image_path=tests/assets/test_image.jpg \
    attention=dot_product \
    skip_jax_distributed_system=true

Output:

Input `<start_of_turn>user
<start_of_image>Describe the image concisely.<end_of_turn>
<start_of_turn>model
` -> `A city skyline with a snow covered mountain`

---

4. Unit Tests

JAX_PLATFORMS=cpu pytest tests/unit/hf_data_processing_test.py

Output:

=========================== 4 passed, 1 warning in 92.40s (0:01:32) ===========================

---

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces a default_prompt configuration to inject a default instruction when the prompt column is missing in a dataset (such as COCO captioning). It adds a new configuration file for SFT vision COCO captions, updates the configuration schema, and modifies the Hugging Face data processing pipeline to handle the prompt injection. The review feedback points out a potential TypeError when checking dataset.features if it is None (e.g., in streaming datasets) and suggests a safer check.

[codecov]

Codecov Report

❌ Patch coverage is 85.71429% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/input_pipeline/hf_data_processing.py	85.71%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @subawocit, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This Pull Request introduces support for the MS COCO captions dataset for Supervised Fine-Tuning (SFT) of multimodal models. It adds a default prompt injection mechanism during the dataset preprocessing phase to handle cases where the prompt column is missing, and provides a pre-configured SFT training recipe for COCO. Overall, the PR is well-structured, follows local patterns, and includes a solid unit test for the new prompt injection logic.

🔍 General Feedback

• Positive Highlights: The addition of a clear, end-to-end multimodal SFT recipe (sft-vision-coco-captions.yml) is excellent and follows existing configs (e.g., SlideVQA/ChartQA) very consistently.
• Robustness: The proposed fallback mechanism is highly functional but can be improved to robustly support IterableDataset streams when streaming=True (which is standard for MaxText Hugging Face data pipelines).
• Consistency: Slightly mismatched default prompt punctuation across base.yml, types.py, and hf_data_processing.py can be unified to avoid any potential confusion or formatting discrepancies.

[igorts-git]

Do we still need getattr here now that we added default_prompt?

[subawocit]

Good catch - I removed this getattr

[igorts-git]

Do we need this here? looks like the AI agent is lazy and does not want to put imports at the top of the file.

[subawocit]

Agree - moved this import up