Add SFT vision COCO captions dataset and config support

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

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"

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.",

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")

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()