Support multi-modal datasets

docs/cli/prepare_data.md

@@ -26,6 +26,8 @@ python scripts/prepare_data.py \
 
   Example: `meta-llama/Llama-3.1-8B-Instruct`
 
+- **`--trust-remote-code`** (flag) Allow executing code from HF Hub when loading the target model's processor.
+
 ### Data Arguments
 
 - **`--data`** (str, required, repeatable) Path to training data. Can be a HuggingFace dataset name or local path. Use multiple times to specify multiple datasets.

docs/cli/train.md

@@ -32,6 +32,8 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \
 
 - **`--verifier-name-or-path`** (str, required) HuggingFace model ID or local path for the verifier/target model.
 
+- **`--trust-remote-code`** (flag) Allow executing code from HF Hub when loading the verifier's tokenizer.
+
 - **`--speculator-type`** (str, default: `"eagle3"`) Type of speculator model to train. Options: `eagle3`, `dflash`
 
 - **`--from-pretrained`** (str, default: `""`) Path to a pretrained draft model to finetune.

pyproject.toml

@@ -250,6 +250,7 @@ select = [
     "INP001", # allow implicit namespace packages in scripts
     "PTH", # os.path is acceptable in scripts
     "T201", # print statements are acceptable in scripts
+    "PLR0915", # allow long parse_args functions
 ]
 
 "examples/**/*.py" = [

scripts/data_generation_offline.py

@@ -277,6 +277,7 @@ async def worker(
             continue
 
         input_ids = item["input_ids"].tolist()
+        messages = item.get("_vllm_messages")
 
         target_hidden_states_path = hidden_states_output_dir / f"hs_{idx}.safetensors"
 
@@ -286,6 +287,7 @@ async def worker(
                     client,
                     model,
                     input_ids,
+                    messages=messages,
                     timeout=request_timeout,
                     max_retries=max_retries,
                 )

scripts/prepare_data.py

@@ -49,6 +49,14 @@ def parse_args():
         required=True,
         help="HuggingFace model ID or local path for target model",
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help=(
+            "Allow executing code from HF Hub when loading the target model's "
+            "processor."
+        ),
+    )
 
     # Data arguments
     parser.add_argument(
@@ -75,7 +83,7 @@ def parse_args():
         type=str,
         default=None,
         help=(
-            "Path to save token frequency distribution"
+            "Path to save token frequency distribution "
             "(default: args.output / 'token_freq.pt')"
         ),
     )
@@ -177,6 +185,7 @@ def main():
         assistant_pattern=args.assistant_pattern,
         turn_dropout=args.turn_dropout,
         minimum_valid_tokens=args.minimum_valid_tokens,
+        trust_remote_code=args.trust_remote_code,
     )
 
     log.info("Done preparing data")

scripts/train.py

@@ -259,6 +259,7 @@ def main(args: argparse.Namespace):
         args.verifier_name_or_path,
         transformer_layer_config.vocab_size,
         args.mask_token_id,
+        trust_remote_code=args.trust_remote_code,
     )
 
     registry = SpeculatorModel.registry
@@ -393,6 +394,11 @@ def _checkpoint_freq(value: str) -> int:
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--verifier-name-or-path", type=str, required=True)
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Allow executing code from HF Hub when loading the verifier's tokenizer.",
+    )
     parser.add_argument(
         "--speculator-type",
         type=str,

src/speculators/data_generation/configs.py

@@ -1,5 +1,6 @@
 """Configuration registries for data generation pipeline."""
 
+import os
 from collections.abc import Callable
 from dataclasses import dataclass
 
@@ -9,22 +10,77 @@
 ]
 
 
-@dataclass
+@dataclass(kw_only=True)
 class DatasetConfig:
     """Configuration for loading a dataset"""
 
     name: str
     hf_path: str
+    hf_name: str | None = None
     split: str
+    filter_fn: Callable[[dict], bool] | None = None
     normalize_fn: Callable[[dict], dict] | None = None
 
 
+def get_coco_dir():
+    return os.getenv("COCO_DIR") or "coco/"
+
+
 def _normalize_ultrachat(example: dict) -> dict:
     if "messages" in example:
         return {"conversations": example["messages"]}
     return example
 
 
+def _parse_sharegpt4v_part(part: str, image_path: str):
+    if part == "<image>":
+        return {"type": "image", "path": image_path}
+
+    return {"type": "text", "text": part}
+
+
+def _parse_sharegpt4v_user_content(content: str, image_path: str):
+    return [_parse_sharegpt4v_part(part, image_path) for part in content.split("\n")]
+
+
+def _parse_sharegpt4v_assistant_content(content: str):
+    return [{"type": "text", "text": content}]
+
+
+def _filter_sharegpt4v_coco(example: dict) -> bool:
+    return example["image"].startswith("coco/")
+
+
+def _normalize_sharegpt4v_coco(example: dict) -> dict:
+    coco_dir = get_coco_dir()
+    image_path = os.path.join(coco_dir, example["image"].removeprefix("coco/"))
+
+    if not os.path.exists(image_path):
+        state_str = "set to" if os.getenv("COCO_DIR") else "default"
+
+        raise ValueError(
+            f"Please download COCO 2017 Train Images from "
+            f"<http://images.cocodataset.org/zips/train2017.zip> and place the "
+            f"extracted folder under `COCO_DIR` ({state_str}: `{coco_dir}`)."
+        )
+
+    messages = [
+        (
+            turn
+            | {
+                "value": (
+                    _parse_sharegpt4v_user_content(turn["value"], image_path)
+                    if turn["from"] in ("human", "user")
+                    else _parse_sharegpt4v_assistant_content(turn["value"])
+                )
+            }
+        )
+        for turn in example["conversations"]
+    ]
+
+    return {"conversations": messages}
+
+
 DATASET_CONFIGS: dict[str, DatasetConfig] = {
     "sharegpt": DatasetConfig(
         name="sharegpt",
@@ -37,4 +93,13 @@ def _normalize_ultrachat(example: dict) -> dict:
         split="train_sft",
         normalize_fn=_normalize_ultrachat,
     ),
+    # NOTE: You need to serve vLLM with `--allowed-local-media-path /path/to/coco`
+    "sharegpt4v_coco": DatasetConfig(
+        name="sharegpt4v_coco",
+        hf_path="Lin-Chen/ShareGPT4V",
+        hf_name="ShareGPT4V",
+        split="train",
+        filter_fn=_filter_sharegpt4v_coco,
+        normalize_fn=_normalize_sharegpt4v_coco,
+    ),
 }