From bd7d1157b6c2ec88973903b0d3472917b44f6537 Mon Sep 17 00:00:00 2001 From: Yuchen Hou Date: Wed, 12 Aug 2026 21:47:33 +0000 Subject: [PATCH] Add SFT vision COCO captions dataset and config support --- src/maxtext/configs/base.yml | 1 + .../post_train/sft-vision-coco-captions.yml | 41 +++++++++++++++++++ src/maxtext/configs/types.py | 4 ++ .../input_pipeline/hf_data_processing.py | 13 ++++++ tests/unit/hf_data_processing_test.py | 21 ++++++++++ 5 files changed, 80 insertions(+) create mode 100644 src/maxtext/configs/post_train/sft-vision-coco-captions.yml diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 227c7e36f1..ae65f93acd 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -719,6 +719,7 @@ train_data_columns: ['text'] # for DPO dataset containing "chosen" and "rejected train_image_column: 'image' eval_data_columns: ['text'] # for DPO dataset containing "chosen" and "rejected" eval_image_column: 'image' +default_prompt: 'Describe this image' packing: true num_epoch: 1 generate_padding_batch_train: false diff --git a/src/maxtext/configs/post_train/sft-vision-coco-captions.yml b/src/maxtext/configs/post_train/sft-vision-coco-captions.yml new file mode 100644 index 0000000000..a986923a52 --- /dev/null +++ b/src/maxtext/configs/post_train/sft-vision-coco-captions.yml @@ -0,0 +1,41 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +base_config: "base.yml" + +use_sft: true +use_tunix_gradient_accumulation: true +use_multimodal: true +# For vision captioning, prompt contains the image, model trains only on caption/completion tokens +sft_train_on_completion_only: true +packing: false # multimodal packing is not supported yet +freeze_vision_encoder_params: true +learning_rate: 2.e-5 + +# -------------- HF pipeline -------------- +dataset_type: hf +hf_path: 'Fhrozen/coco-narratives' +train_split: 'train' +hf_eval_split: 'val' + +# Image column +train_image_column: 'image' +eval_image_column: 'image' + +# Text columns: ['prompt', 'captions'] +# 'prompt' is auto-injected as the user instruction (e.g. "Describe this image.") +# 'captions' uses the standard COCO ground truth captions +train_data_columns: ['prompt', 'captions'] +eval_data_columns: ['prompt', 'captions'] +default_prompt: "Describe the image concisely" diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 4865557b5a..d0bbc18a20 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1385,6 +1385,10 @@ class DatasetGeneral(BaseModel): train_image_column: str | list[str] = Field("image", description="Column name(s) for images in the training data.") eval_data_columns: list[str] = Field(["text"], description="Column(s) to use from the evaluation data.") eval_image_column: str | list[str] = Field("image", description="Column name(s) for images in evaluation data.") + default_prompt: str = Field( + "Describe this image", + description="Default prompt injected into the dataset when the prompt column is missing.", + ) packing: bool = Field( True, description="Whether to pack multiple short examples into a single sequence.", diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py index c835a27a3e..6cdc7ddd36 100644 --- a/src/maxtext/input_pipeline/hf_data_processing.py +++ b/src/maxtext/input_pipeline/hf_data_processing.py @@ -42,6 +42,17 @@ def _get_pad_id(tokenizer): return pad_id +def add_default_prompt_if_missing(dataset, text_columns, config): + """Inject default prompt if prompt column is missing in dataset (e.g. COCO captioning).""" + prompt_col = text_columns[0] + # Check if features is None (e.g. streaming IterableDataset) or prompt column is absent + features = getattr(dataset, "features", None) or getattr(dataset, "column_names", None) + # If so, populate each example with default_prompt + if features is None or prompt_col not in features: + dataset = dataset.map(lambda ex: {**ex, prompt_col: config.default_prompt}) + return dataset + + def vision_sft_preprocessing_pipeline( dataset, config, @@ -89,6 +100,8 @@ def vision_sft_preprocessing_pipeline( ) image_column = "images" + dataset = add_default_prompt_if_missing(dataset, text_columns, config) + dataset = dataset.select_columns(text_columns + [image_column]) if image_column != "images": dataset = dataset.rename_column(image_column, "images") diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py index 262c56ff9b..2e012556a2 100644 --- a/tests/unit/hf_data_processing_test.py +++ b/tests/unit/hf_data_processing_test.py @@ -17,7 +17,9 @@ import sys import unittest import os.path +from unittest.mock import MagicMock +import datasets import jax from jax.sharding import Mesh from jax.experimental import mesh_utils @@ -123,6 +125,25 @@ def get_first_batch(iterator): self.assertTrue((train_batch1["inputs"] == train_batch2["inputs"]).all()) # pytype: disable=unsupported-operands self.assertTrue((train_batch1["targets"] == train_batch2["targets"]).all()) # pytype: disable=unsupported-operands + def test_add_default_prompt_if_missing(self): + config = MagicMock(default_prompt="Describe this image") + + # When prompt column is missing, default_prompt is added + ds = datasets.Dataset.from_dict({"image": [b"img"], "captions": ["a cat"]}) + ds = hf_data_processing.add_default_prompt_if_missing(ds, ["prompt", "captions"], config) + self.assertEqual(ds[0]["prompt"], "Describe this image") + + # When prompt column already exists, existing prompt is preserved + ds_with_prompt = datasets.Dataset.from_dict({"prompt": ["Custom prompt"], "captions": ["a cat"]}) + ds_with_prompt = hf_data_processing.add_default_prompt_if_missing(ds_with_prompt, ["prompt", "captions"], config) + self.assertEqual(ds_with_prompt[0]["prompt"], "Custom prompt") + + # When streaming IterableDataset is used, default_prompt is added + iterable_ds = ds.to_iterable_dataset() + iterable_ds = hf_data_processing.add_default_prompt_if_missing(iterable_ds, ["prompt", "captions"], config) + first_item = next(iter(iterable_ds)) + self.assertEqual(first_item["prompt"], "Describe this image") + if __name__ == "__main__": unittest.main()