[v1] Support multimodal data training (#10656)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

src/llamafactory/v1/plugins/data_plugins/converter.py

@@ -14,11 +14,13 @@
 
 
 import json
+import re
 from typing import Any, Literal, NotRequired, TypedDict
 
 from ...utils import logging
+from ...utils.constants import AUDIO_PLACEHOLDER, IMAGE_PLACEHOLDER, VIDEO_PLACEHOLDER
 from ...utils.plugin import BasePlugin
-from ...utils.types import DPOSample, Sample, SFTSample, ToolCall
+from ...utils.types import Content, DPOSample, Sample, SFTSample, ToolCall
 
 
 logger = logging.get_logger(__name__)
@@ -29,6 +31,9 @@ class AlpacaSample(TypedDict, total=False):
     instruction: str
     input: NotRequired[str]
     output: str
+    images: NotRequired[list[str] | str]
+    videos: NotRequired[list[str] | str]
+    audios: NotRequired[list[str] | str]
 
 
 SharegptMessage = TypedDict(
@@ -40,6 +45,9 @@ SharegptMessage = TypedDict(
 class SharegptSample(TypedDict, total=False):
     conversations: list[SharegptMessage]
     tools: NotRequired[str]
+    images: NotRequired[list[str] | str]
+    videos: NotRequired[list[str] | str]
+    audios: NotRequired[list[str] | str]
 
 
 class OpenaiMessage(TypedDict, total=False):
@@ -54,6 +62,65 @@ class OpenaiSample(TypedDict, total=False):
 class PairSample(TypedDict, total=False):
     chosen: list[OpenaiMessage]
     rejected: list[OpenaiMessage]
+    images: NotRequired[list[str] | str]
+    videos: NotRequired[list[str] | str]
+    audios: NotRequired[list[str] | str]
+
+
+# Inline media tag -> v1 content block type, and the raw-sample column holding the paths.
+_MEDIA_SPECS: tuple[tuple[str, str, str], ...] = (
+    (IMAGE_PLACEHOLDER, "image_url", "images"),
+    (VIDEO_PLACEHOLDER, "video_url", "videos"),
+    (AUDIO_PLACEHOLDER, "audio_url", "audios"),
+)
+_TAG_TO_BLOCK = {tag: block_type for tag, block_type, _col in _MEDIA_SPECS}
+_TAG_PATTERN = re.compile("(" + "|".join(re.escape(tag) for tag, _b, _c in _MEDIA_SPECS) + ")")
+
+
+def _as_media_list(value: Any) -> list:
+    """Normalize a media column value into a list of paths/URLs (None -> [], scalar -> [scalar])."""
+    if value is None:
+        return []
+    if isinstance(value, (list, tuple)):
+        return list(value)
+    return [value]
+
+
+def _build_media_iters(raw_sample: dict[str, Any]) -> dict[str, Any]:
+    """Build per-modality path iterators from a raw sample's media columns."""
+    return {tag: iter(_as_media_list(raw_sample.get(col))) for tag, _block_type, col in _MEDIA_SPECS}
+
+
+def _to_content_blocks(text: str, media_iters: dict[str, Any]) -> list[Content]:
+    """Split ``text`` on inline media placeholders, interleaving media-url content blocks.
+
+    Each placeholder consumes the next path from its modality iterator (in document order). Plain
+    text with no placeholders yields a single text block (byte-identical to the legacy behavior).
+    Raises on an unmatched placeholder (more tags than media files).
+    """
+    if not _TAG_PATTERN.search(text):
+        return [{"type": "text", "value": text}]
+
+    blocks: list[Content] = []
+    for segment in _TAG_PATTERN.split(text):
+        block_type = _TAG_TO_BLOCK.get(segment)
+        if block_type is not None:
+            try:
+                path = next(media_iters[segment])
+            except StopIteration:
+                raise ValueError(f"More {segment} tags than provided media files.") from None
+            blocks.append({"type": block_type, "value": path})
+        elif segment:
+            blocks.append({"type": "text", "value": segment})
+    return blocks
+
+
+def _assert_media_consumed(media_iters: dict[str, Any]) -> None:
+    """Ensure every media file was referenced by a tag (fewer tags than media -> error)."""
+    for tag, media_iter in media_iters.items():
+        unused = len(list(media_iter))
+        if unused:
+            raise ValueError(f"Fewer {tag} tags than provided media files ({unused} unused).")
 
 
 class DataConverterPlugin(BasePlugin):
@@ -76,6 +143,7 @@ def alpaca_converter(raw_sample: AlpacaSample) -> SFTSample:
         SFTSample: SFT sample.
     """
     messages = []
+    media_iters = _build_media_iters(raw_sample)
     if "system" in raw_sample:
         messages.append(
             {"role": "system", "content": [{"type": "text", "value": raw_sample["system"]}], "loss_weight": 0.0}
@@ -85,9 +153,9 @@ def alpaca_converter(raw_sample: AlpacaSample) -> SFTSample:
         messages.append(
             {
                 "role": "user",
-                "content": [
-                    {"type": "text", "value": raw_sample.get("instruction", "") + raw_sample.get("input", "")}
-                ],
+                "content": _to_content_blocks(
+                    raw_sample.get("instruction", "") + raw_sample.get("input", ""), media_iters
+                ),
                 "loss_weight": 0.0,
             }
         )
@@ -97,6 +165,7 @@ def alpaca_converter(raw_sample: AlpacaSample) -> SFTSample:
             {"role": "assistant", "content": [{"type": "text", "value": raw_sample["output"]}], "loss_weight": 1.0}
         )
 
+    _assert_media_consumed(media_iters)
     return {"messages": messages}
 
 
@@ -121,6 +190,7 @@ def sharegpt_converter(raw_sample: SharegptSample) -> SFTSample:
     }
     sample = {}
     messages = []
+    media_iters = _build_media_iters(raw_sample)
     for message in raw_sample.get("conversations", []):
         tag = message["from"]
         if tag not in tag_mapping:
@@ -146,11 +216,12 @@ def sharegpt_converter(raw_sample: SharegptSample) -> SFTSample:
             messages.append(
                 {
                     "role": tag_mapping[tag],
-                    "content": [{"type": "text", "value": message["value"]}],
+                    "content": _to_content_blocks(message["value"], media_iters),
                     "loss_weight": 1.0 if tag == "gpt" else 0.0,
                 }
             )
 
+    _assert_media_consumed(media_iters)
     sample["messages"] = messages
 
     tools = raw_sample.get("tools")
@@ -178,6 +249,8 @@ def pair_converter(raw_sample: PairSample) -> DPOSample:
     """
 
     def process_message(raw_messages: list[OpenaiMessage]):
+        # chosen and rejected share the sample's media; each side consumes its own iterators.
+        media_iters = _build_media_iters(raw_sample)
         messages = []
         for message in raw_messages:
             if message["role"] == "tool":
@@ -201,11 +274,12 @@ def pair_converter(raw_sample: PairSample) -> DPOSample:
                 messages.append(
                     {
                         "role": message["role"],
-                        "content": [{"type": "text", "value": message["content"]}],
+                        "content": _to_content_blocks(message["content"], media_iters),
                         "loss_weight": 1.0 if message["role"] == "assistant" else 0.0,
                     }
                 )
 
+        _assert_media_consumed(media_iters)
         return messages
 
     sample = {}
@@ -221,3 +295,4 @@ def pair_converter(raw_sample: PairSample) -> DPOSample:
             logger.warning_rank0(f"Invalid tools format: {str(tools)}")
 
     return sample
+