From a57dde2d947092892f3e922ac4fee8918293ea54 Mon Sep 17 00:00:00 2001
From: natedemoss <ndemoss28@gmail.com>
Date: Fri, 28 Aug 2026 23:14:43 -0400
Subject: [PATCH 1/4] [Bugfix][Core] Fix multimodal prompt building for rows
 without media

`_build_messages` substitutes `<image>`/`<video>`/`<audio>` placeholders by
splitting the prompt on a regex alternation built from the placeholders a row
actually carries. Three problems with how that set is built:

- A row whose media column is absent (`None`) contributes nothing, so for a
  text-only row of a mixed dataset the alternation collapses to `"()"`. That
  pattern matches the empty string at every position, and the prompt comes back
  as one `{"type": "text"}` dict per character. `filter_long_prompt` already
  treats text-only rows as a supported case, so such rows do reach here.

- A media type name that `MultimodalTypes.get` does not know was skipped
  silently. A plausible typo (`{"images": "images"}` instead of
  `{"image": "images"}`) therefore dropped every image and trained text-only
  against a VLM dataset, with no warning. Reject it instead, and name the
  supported types in the `--multimodal-keys` help.

- The leftover-media check ran even when no message had string content. A row
  already in list-of-dicts form is deliberately passed through untouched
  ("no processing will be done"), so its media is never spent on a placeholder
  and the check aborted the whole dataset load with a count-mismatch assert.
  Only run it when placeholder expansion actually happened.

Also add the missing space in that assert's message.

Tests: tests/test_build_messages.py covers all three, plus the substitution
paths that must not change. Three of the nine fail on main.

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: natedemoss <ndemoss28@gmail.com>
---
 .buildkite/pipeline.yml      |   1 +
 tests/test_build_messages.py | 115 +++++++++++++++++++++++++++++++++++
 vime/utils/arguments.py      |   4 +-
 vime/utils/data.py           |  39 ++++++++----
 4 files changed, 147 insertions(+), 12 deletions(-)
 create mode 100644 tests/test_build_messages.py

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 6aeb41548..0ed21be78 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -146,6 +146,7 @@ steps:
           for test_file in \
             tests/test_advantage_whiten_cp.py \
             tests/test_block_fp8_zero_block.py \
+            tests/test_build_messages.py \
             tests/test_deep_ep_tms_patch.py \
             tests/test_discounted_returns.py \
             tests/test_eval_config.py \
diff --git a/tests/test_build_messages.py b/tests/test_build_messages.py
new file mode 100644
index 000000000..2afb2df99
--- /dev/null
+++ b/tests/test_build_messages.py
@@ -0,0 +1,115 @@
+"""CPU unit tests for ``vime.utils.data._build_messages``.
+
+``--multimodal-keys`` maps a media type name to the dataset column holding that
+media, and the prompt marks insertion points with the type's placeholder
+(``<image>``, ``<video>``, ``<audio>``). ``_build_messages`` rewrites a plain
+string prompt into the list-of-content-dicts form the VLM chat templates and
+``process_vision_info`` expect.
+
+The cases below pin down what happens at the edges of that rewrite: a row of a
+multimodal dataset that carries no media at all, a row that is already in
+list-of-dicts form, and a ``--multimodal-keys`` mapping that names a type vime
+does not support.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vime.utils.data import _build_messages
+
+
+NUM_GPUS = 0
+
+IMAGE_KEYS = {"image": "images"}
+
+
+def _row(prompt, **columns):
+    return {"text": prompt, **columns}
+
+
+def _content(messages):
+    assert len(messages) == 1, f"expected a single message, got {len(messages)}"
+    return messages[0]["content"]
+
+
+@pytest.mark.unit
+def test_media_placeholders_are_replaced():
+    messages = _build_messages(_row("What is in <image>?", images=["a.png"]), "text", True, IMAGE_KEYS)
+
+    assert _content(messages) == [
+        {"type": "text", "text": "What is in "},
+        {"type": "image", "image": "a.png"},
+        {"type": "text", "text": "?"},
+    ]
+
+
+@pytest.mark.unit
+def test_media_entries_may_be_rich_dicts():
+    item = {"type": "image", "image": "a.png", "max_pixels": 50176}
+    messages = _build_messages(_row("<image> here", images=[item]), "text", True, IMAGE_KEYS)
+
+    assert _content(messages) == [item, {"type": "text", "text": " here"}]
+
+
+@pytest.mark.unit
+def test_row_without_media_keeps_its_prompt_intact():
+    """A mixed dataset has rows with no media; those prompts must stay intact.
+
+    Building the split pattern from an empty placeholder set yields ``"()"``,
+    which matches the empty string at every position, so the prompt used to come
+    back as one ``{"type": "text"}`` dict per character.
+    """
+    prompt = "Describe this image."
+    messages = _build_messages(_row(prompt, images=None), "text", True, IMAGE_KEYS)
+
+    # Same representation a text-only dataset gets, i.e. no `--multimodal-keys`.
+    assert _content(messages) == prompt
+
+
+@pytest.mark.unit
+def test_row_with_an_empty_media_column_keeps_one_text_segment():
+    prompt = "Describe this image."
+    messages = _build_messages(_row(prompt, images=[]), "text", True, IMAGE_KEYS)
+
+    assert _content(messages) == [{"type": "text", "text": prompt}]
+
+
+@pytest.mark.unit
+def test_unknown_media_type_is_rejected():
+    """A typo like ``images`` used to be skipped silently: the media never made
+    it into the prompt and training ran text-only against a VLM dataset."""
+    with pytest.raises(ValueError, match="Unknown multimodal type 'images'"):
+        _build_messages(_row("What is in <image>?", images=["a.png"]), "text", True, {"images": "images"})
+
+
+@pytest.mark.unit
+def test_list_content_row_is_passed_through():
+    """Content already in list-of-dicts form embeds its media inline, so there
+    is no placeholder for this function to spend the row's media on."""
+    content = [{"type": "image", "image": "a.png"}, {"type": "text", "text": "What is in it?"}]
+    prompt = [{"role": "user", "content": list(content)}]
+
+    messages = _build_messages(_row(prompt, images=["a.png"]), "text", True, IMAGE_KEYS)
+
+    assert _content(messages) == content
+
+
+@pytest.mark.unit
+def test_more_media_than_placeholders_still_raises():
+    with pytest.raises(AssertionError, match="Multimodal data count mismatch"):
+        _build_messages(_row("What is in <image>?", images=["a.png", "b.png"]), "text", True, IMAGE_KEYS)
+
+
+@pytest.mark.unit
+def test_more_placeholders_than_media_still_raises():
+    with pytest.raises(AssertionError, match="Not enough image data"):
+        _build_messages(_row("<image> vs <image>", images=["a.png"]), "text", True, IMAGE_KEYS)
+
+
+@pytest.mark.unit
+def test_without_multimodal_keys_the_prompt_is_untouched():
+    prompt = "Describe this image."
+
+    assert _build_messages(_row(prompt), "text", False, None) == prompt
+    assert _content(_build_messages(_row(prompt), "text", True, None)) == prompt
diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py
index 113ab12a0..dbde57666 100644
--- a/vime/utils/arguments.py
+++ b/vime/utils/arguments.py
@@ -673,7 +673,9 @@ def add_data_arguments(parser):
                 type=json.loads,
                 default=None,
                 help=(
-                    'JSON string for multimodal data mapping media types to data keys. Example: \'{"image": "image_file"}\''
+                    "JSON string for multimodal data mapping media types to data keys. "
+                    "Supported media types are 'image', 'video' and 'audio'; each one substitutes the "
+                    'matching <image>/<video>/<audio> placeholder in the prompt. Example: \'{"image": "image_file"}\''
                 ),
             )
             parser.add_argument("--metadata-key", type=str, default="metadata", help="JSON dataset key")
diff --git a/vime/utils/data.py b/vime/utils/data.py
index 984348cad..a7d28f684 100644
--- a/vime/utils/data.py
+++ b/vime/utils/data.py
@@ -150,20 +150,36 @@ def _build_messages(data: dict, prompt_key: str, as_conversation: bool, multimod
         else:
             prompt = [{"role": "user", "content": prompt}]
 
+    multimodals = {}
     if multimodal_keys:
         # Build mapping: placeholder -> (MultimodalType, content_list)
-        multimodals = {}
         for type_name, data_key in multimodal_keys.items():
             mt = MultimodalTypes.get(type_name)
-            if mt:
-                multimodal_data = data.get(data_key)
-                if multimodal_data is not None:
-                    multimodals[mt.placeholder] = (mt, list(multimodal_data))
-
+            if mt is None:
+                raise ValueError(
+                    f"Unknown multimodal type '{type_name}' in --multimodal-keys; "
+                    f"supported types are {[m.name for m in MultimodalTypes.all()]}."
+                )
+            multimodal_data = data.get(data_key)
+            if multimodal_data is not None:
+                multimodals[mt.placeholder] = (mt, list(multimodal_data))
+
+    # Only rows that actually carry media need placeholder substitution. Running
+    # the split with an empty `multimodals` would build the pattern "()", which
+    # matches the empty string everywhere and shatters the prompt into one text
+    # segment per character.
+    if multimodals:
         pattern = "(" + "|".join(re.escape(p) for p in multimodals.keys()) + ")"
 
+        # Media is only consumed from messages whose content is a plain string.
+        # A row already in list-of-dicts form embeds its media inline, so the
+        # leftover check below would fire on media this function never had a
+        # placeholder to spend.
+        expanded_any = False
+
         for message in prompt:
             if isinstance(message["content"], str):
+                expanded_any = True
                 content_list = []
                 for segment in re.split(pattern, message["content"]):
                     if not segment:
@@ -203,11 +219,12 @@ def _build_messages(data: dict, prompt_key: str, as_conversation: bool, multimod
                     f"Unsupported content type: {type(message['content'])}, expected str or list of dicts"
                 )
 
-        for placeholder, (mt, remaining) in multimodals.items():
-            assert len(remaining) == 0, (
-                f"Multimodal data count mismatch: {len(remaining)} more {mt.name}(s)"
-                f"than '{placeholder}' placeholders in prompt"
-            )
+        if expanded_any:
+            for placeholder, (mt, remaining) in multimodals.items():
+                assert len(remaining) == 0, (
+                    f"Multimodal data count mismatch: {len(remaining)} more {mt.name}(s) "
+                    f"than '{placeholder}' placeholders in prompt"
+                )
 
     return prompt
 

From e6918acf6c57350562f79df18e72a6aba3eb4cf3 Mon Sep 17 00:00:00 2001
From: natedemoss <ndemoss28@gmail.com>
Date: Fri, 28 Aug 2026 23:58:19 -0400
Subject: [PATCH 2/4] Update vime/utils/data.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Signed-off-by: natedemoss <ndemoss28@gmail.com>
---
 vime/utils/data.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/vime/utils/data.py b/vime/utils/data.py
index a7d28f684..66ba99870 100644
--- a/vime/utils/data.py
+++ b/vime/utils/data.py
@@ -162,6 +162,8 @@ def _build_messages(data: dict, prompt_key: str, as_conversation: bool, multimod
                 )
             multimodal_data = data.get(data_key)
             if multimodal_data is not None:
+                if isinstance(multimodal_data, (str, dict)):
+                    multimodal_data = [multimodal_data]
                 multimodals[mt.placeholder] = (mt, list(multimodal_data))
 
     # Only rows that actually carry media need placeholder substitution. Running

From d886b69f562451b156a0bf1bc158df228044cb6c Mon Sep 17 00:00:00 2001
From: natedemoss <ndemoss28@gmail.com>
Date: Fri, 28 Aug 2026 23:58:24 -0400
Subject: [PATCH 3/4] Update tests/test_build_messages.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Signed-off-by: natedemoss <ndemoss28@gmail.com>
---
 tests/test_build_messages.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tests/test_build_messages.py b/tests/test_build_messages.py
index 2afb2df99..30946b007 100644
--- a/tests/test_build_messages.py
+++ b/tests/test_build_messages.py
@@ -52,6 +52,22 @@ def test_media_entries_may_be_rich_dicts():
     assert _content(messages) == [item, {"type": "text", "text": " here"}]
 
 
+@pytest.mark.unit
+def test_single_media_entry_as_string_or_dict_is_wrapped():
+    # Test single string
+    messages = _build_messages(_row("What is in <image>?", images="a.png"), "text", True, IMAGE_KEYS)
+    assert _content(messages) == [
+        {"type": "text", "text": "What is in "},
+        {"type": "image", "image": "a.png"},
+        {"type": "text", "text": "?"},
+    ]
+
+    # Test single dict
+    item = {"type": "image", "image": "a.png", "max_pixels": 50176}
+    messages = _build_messages(_row("<image> here", images=item), "text", True, IMAGE_KEYS)
+    assert _content(messages) == [item, {"type": "text", "text": " here"}]
+
+
 @pytest.mark.unit
 def test_row_without_media_keeps_its_prompt_intact():
     """A mixed dataset has rows with no media; those prompts must stay intact.

From 9574840e4479d328b536fa80b982dec7626a6b90 Mon Sep 17 00:00:00 2001
From: natedemoss <ndemoss28@gmail.com>
Date: Fri, 28 Aug 2026 23:58:31 -0400
Subject: [PATCH 4/4] Update vime/utils/data.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Signed-off-by: natedemoss <ndemoss28@gmail.com>
---
 vime/utils/data.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/vime/utils/data.py b/vime/utils/data.py
index 66ba99870..c4c66f67c 100644
--- a/vime/utils/data.py
+++ b/vime/utils/data.py
@@ -223,10 +223,11 @@ def _build_messages(data: dict, prompt_key: str, as_conversation: bool, multimod
 
         if expanded_any:
             for placeholder, (mt, remaining) in multimodals.items():
-                assert len(remaining) == 0, (
-                    f"Multimodal data count mismatch: {len(remaining)} more {mt.name}(s) "
-                    f"than '{placeholder}' placeholders in prompt"
-                )
+                if len(remaining) != 0:
+                    raise AssertionError(
+                        f"Multimodal data count mismatch: {len(remaining)} more {mt.name}(s) "
+                        f"than '{placeholder}' placeholders in prompt"
+                    )
 
     return prompt