Add --render-endpoint to prepare_data.py for vLLM-delegated tokenization (#652)

pyproject.toml

@@ -38,6 +38,7 @@ keywords = [
 dependencies = [
     "click",
     "datasets>=4.0.0,<=5.0.0",
+    "httpx",
     "huggingface-hub",
     "loguru>=0.7.2,<=0.7.3",
     "numpy>=2.0.0,<=2.4.6",

scripts/prepare_data.py

@@ -24,6 +24,7 @@
 
 import argparse
 import glob
+import json
 import logging
 import shutil
 import sys
@@ -105,6 +106,28 @@ def parse_args():
             "per conversation for data augmentation."
         ),
     )
+    parser.add_argument(
+        "--vllm-server-url",
+        type=str,
+        default=None,
+        help=(
+            "Base URL of the vLLM server, typically the same instance used "
+            "for hidden-state extraction (e.g. http://localhost:8000). "
+            "When set, tokenization is delegated to the server's "
+            "/v1/chat/completions/render endpoint instead of local HF "
+            "apply_chat_template."
+        ),
+    )
+    parser.add_argument(
+        "--chat-template-kwargs",
+        type=str,
+        default=None,
+        help=(
+            "JSON object of extra kwargs forwarded to the chat template, "
+            'e.g. \'{"enable_thinking": true}\'. Only supported on the render '
+            "path (requires --vllm-server-url)."
+        ),
+    )
 
     # Output arguments
     parser.add_argument(
@@ -177,6 +200,11 @@ def main():
         if args.token_freq_path is None
         else Path(args.token_freq_path)
     )
+    chat_template_kwargs = (
+        json.loads(args.chat_template_kwargs)
+        if args.chat_template_kwargs is not None
+        else None
+    )
 
     dataset, _ = load_and_preprocess_dataset(
         target_model_path=args.model,
@@ -190,6 +218,8 @@ def main():
         turn_dropout=args.turn_dropout,
         minimum_valid_tokens=args.minimum_valid_tokens,
         trust_remote_code=args.trust_remote_code,
+        render_endpoint=args.vllm_server_url,
+        chat_template_kwargs=chat_template_kwargs,
     )
 
     log.info("Done preparing data")

src/speculators/data_generation/preprocessing.py

@@ -8,6 +8,7 @@
 from re import Pattern
 from typing import cast
 
+import httpx
 import torch
 from datasets import Dataset as HFDataset
 from datasets import concatenate_datasets, load_dataset
@@ -23,10 +24,13 @@
 
 from speculators.data_generation.configs import DATASET_CONFIGS
 from speculators.data_generation.logging_utils import PipelineLogger
+from speculators.data_generation.render_client import RenderError, render_conversation
 from speculators.data_generation.torch_utils import set_default_torch_num_threads
+from speculators.data_generation.vllm_client import InvalidResponseError
 from speculators.train.vocab_mapping import save_token_frequency_distribution
 
 __all__ = [
+    "build_dataset_from_render",
     "build_eagle3_dataset",
     "load_and_preprocess_dataset",
     "load_raw_dataset",
@@ -651,6 +655,131 @@ def build_eagle3_dataset(
     return dataset
 
 
+def build_dataset_from_render(
+    dataset: HFDataset,
+    render_endpoint: str,
+    processor: ProcessorLike,
+    max_length: int = 2048,
+    chat_template_kwargs: dict | None = None,
+    turn_dropout: bool = False,
+    minimum_valid_tokens: int | None = None,
+) -> HFDataset:
+    """Tokenize via vLLM render endpoint instead of local HF apply_chat_template."""
+    results: dict[str, list] = {"input_ids": [], "loss_mask": [], "seq_len": []}
+    has_tools = "tools" in dataset.column_names
+    tokenizer = get_tokenizer(processor)
+
+    # Multimodal models extract hidden states via the Chat Completions API, which
+    # needs the original messages -- token_ids alone cannot carry the images.
+    # Mirror the default path and keep messages for ProcessorMixin processors so
+    # render-path multimodal rows can be extracted downstream.
+    if isinstance(processor, ProcessorMixin):
+        results["messages"] = []
+
+    # Regex fallback for conversations the render endpoint returns no mask for
+    # (template lacks {% generation %} tags). Detected lazily, so a run the
+    # server masks completely never fails here.
+    assistant_pattern: str | None = None
+    server_mask_count = 0
+    regex_fallback_count = 0
+    dropped_count = 0
+
+    for idx in range(len(dataset)):
+        row = dataset[idx]
+        conv = row.get("conversations")
+        if not conv or not isinstance(conv, list):
+            continue
+
+        normalized = _normalize_conversation(conv, turn_dropout)
+        if not normalized:
+            continue
+
+        vllm_messages = _adapt_conv_for_vllm(normalized)
+        conv_tools = _parse_conv_tools(row.get("tools"), idx) if has_tools else None
+
+        # Resilience parity with the default path (_preprocess_batch wraps each
+        # conversation in try/except): one conversation that fails to render must
+        # not abort the whole run.
+        try:
+            rendered = render_conversation(
+                render_endpoint,
+                vllm_messages,
+                tools=conv_tools,
+                chat_template_kwargs=chat_template_kwargs,
+                max_length=max_length,
+            )
+        except (RenderError, InvalidResponseError, httpx.HTTPError) as exc:
+            log.error(f"Render failed for conversation {idx}: {exc}")
+            dropped_count += 1
+            continue
+
+        token_ids = rendered["token_ids"]
+        loss_mask_raw = rendered["loss_mask"]
+
+        # vLLM does not truncate multimodal prompts -- it cannot split an image's
+        # placeholder block -- so token_ids can exceed max_length. The default
+        # pipeline drops these rows; match it instead of emitting an over-length
+        # (and, after local masking, misaligned) sample.
+        if len(token_ids) > max_length:
+            dropped_count += 1
+            continue
+
+        if loss_mask_raw is not None:
+            loss_mask = torch.as_tensor(loss_mask_raw, dtype=torch.long)
+            server_mask_count += 1
+        else:
+            if assistant_pattern is None:
+                assistant_pattern = _detect_assistant_pattern(processor)
+            # token_ids is already within max_length (over-length rows dropped
+            # above), so re-tokenize without truncation -- it round-trips to the
+            # same tokens and the offsets line up with the server's token_ids.
+            text = tokenizer.decode(token_ids)
+            encoded = tokenizer(
+                text,
+                return_offsets_mapping=True,
+                add_special_tokens=False,
+            )
+            loss_mask = _create_loss_mask_from_offsets(
+                text,
+                encoded["offset_mapping"],
+                assistant_pattern,
+                conv_idx=idx,
+                max_length=max_length,
+            )
+            loss_mask = torch.as_tensor(loss_mask, dtype=torch.long)
+            regex_fallback_count += 1
+
+        # Parity with the default path's `assert len(input_ids) == len(loss_mask)`:
+        # never emit a misaligned pair -- drop it rather than corrupt the sample.
+        if len(loss_mask) != len(token_ids):
+            log.error(
+                f"Render mask/token length mismatch for conversation {idx} "
+                f"(input_ids={len(token_ids)}, loss_mask={len(loss_mask)}); dropping"
+            )
+            dropped_count += 1
+            continue
+
+        if minimum_valid_tokens is not None:
+            if int(loss_mask.sum().item()) < minimum_valid_tokens:
+                continue
+
+        results["input_ids"].append(torch.tensor(token_ids, dtype=torch.long))
+        results["loss_mask"].append(loss_mask)
+        results["seq_len"].append(len(token_ids))
+        if "messages" in results:
+            results["messages"].append(vllm_messages)
+
+    log.info(
+        f"Render masking: {server_mask_count} used the vLLM server mask, "
+        f"{regex_fallback_count} used the local regex fallback, "
+        f"{dropped_count} dropped (render error / over-length / misaligned)."
+    )
+
+    out = HFDataset.from_dict(results)
+    out.set_format(type="torch")
+    return out
+
+
 def _load_hf_dataset(spec: str) -> tuple[HFDataset, None]:
     """Load an arbitrary HuggingFace dataset from an ``hf:`` spec.
 
@@ -797,34 +926,54 @@ def load_and_preprocess_dataset(
     turn_dropout: bool = False,
     minimum_valid_tokens: int | None = None,
     trust_remote_code: bool = False,
+    render_endpoint: str | None = None,
+    chat_template_kwargs: dict | None = None,
 ) -> tuple[HFDataset, ProcessorLike]:
     """Load, tokenize, and preprocess a dataset for EAGLE3 training.
 
-    Uses the processor's built-in chat template via apply_chat_template.
-    Caching is handled automatically by HuggingFace datasets.
+    Uses the processor's built-in chat template via apply_chat_template, or
+    delegates tokenization to a running vLLM render server when render_endpoint
+    is set. Caching is handled automatically by HuggingFace datasets.
 
     Args:
-        target_model_path: HuggingFace model ID or local path
-        train_data_path: Dataset name or path to JSON/JSONL file
-        seq_length: Maximum sequence length
-        build_dataset_num_proc: Number of processes for dataset building
-        seed: Random seed for shuffling
-        max_samples: Optional limit on number of samples
-        token_freq_path: Path to save token frequency distribution
-        cache_dir: Directory to cache HuggingFace datasets (optional)
+        target_model_path: HuggingFace model ID or local path.
+        train_data_paths: Dataset names or paths to JSON/JSONL files.
+        seq_length: Maximum sequence length.
+        build_dataset_num_proc: Number of processes for dataset building.
+        seed: Random seed for shuffling.
+        max_samples: Optional limit on number of samples.
+        token_freq_path: Path to save token frequency distribution.
         assistant_pattern: Optional custom regex pattern for matching assistant
-                          responses. If None, pattern will be auto-detected from
-                          chat template.
+            responses. If None, pattern will be auto-detected from the chat
+            template. Mutually exclusive with render_endpoint -- the render
+            path always auto-detects its own pattern.
         turn_dropout: If True, randomly keeps first N consecutive turns per
-                     conversation
-        minimum_valid_tokens: Number of tokens to consider for a valid sample
+            conversation.
+        minimum_valid_tokens: Number of tokens to consider for a valid sample.
         trust_remote_code: If True, allows executing code from HF Hub.
+        render_endpoint: URL of a running vLLM server. When set, tokenization
+            is delegated to the server's /v1/chat/completions/render endpoint
+            instead of local HF apply_chat_template.
+        chat_template_kwargs: Extra kwargs forwarded to the chat template
+            (e.g. {"enable_thinking": True}). Only supported on the render
+            path (requires render_endpoint).
 
     Returns:
-        Tuple of (preprocessed_dataset, processor)
+        Tuple of (preprocessed_dataset, processor).
     """
     if minimum_valid_tokens is not None and minimum_valid_tokens < 0:
         raise ValueError("minimum_valid_tokens must be >= 0")
+    if render_endpoint is not None and assistant_pattern is not None:
+        raise ValueError(
+            "assistant_pattern has no effect when render_endpoint is set -- the "
+            "render path always auto-detects its own pattern from the chat "
+            "template. Pass only one of assistant_pattern / render_endpoint."
+        )
+    if chat_template_kwargs is not None and render_endpoint is None:
+        raise ValueError(
+            "chat_template_kwargs is only supported on the render path; pass "
+            "render_endpoint or drop chat_template_kwargs."
+        )
     log.section("Starting dataset preprocessing")
     if minimum_valid_tokens is not None:
         log.info(
@@ -834,7 +983,11 @@ def load_and_preprocess_dataset(
     log.subsection("Loading processor")
     processor = load_processor(target_model_path, trust_remote_code=trust_remote_code)
 
-    if not hasattr(processor, "apply_chat_template") or processor.chat_template is None:
+    if render_endpoint is not None:
+        log.info(f"Using render endpoint: {render_endpoint}")
+    elif (
+        not hasattr(processor, "apply_chat_template") or processor.chat_template is None
+    ):
         raise ValueError(
             f"Processor for {target_model_path} does not support chat templates. "
             "Please use a model with a pre-configured chat template."
@@ -861,18 +1014,29 @@ def load_and_preprocess_dataset(
 
         log.info(f"Loaded {len(raw_dataset)} samples")
 
-        if turn_dropout:
-            log.info("Turn dropout enabled: randomly keeping N consecutive turns")
-
-        preprocessed_dataset = build_eagle3_dataset(
-            dataset=raw_dataset,
-            processor=processor,
-            max_length=seq_length,
-            num_proc=build_dataset_num_proc,
-            assistant_pattern=assistant_pattern,
-            turn_dropout=turn_dropout,
-            minimum_valid_tokens=minimum_valid_tokens,
-        )
+        if render_endpoint is not None:
+            preprocessed_dataset = build_dataset_from_render(
+                dataset=raw_dataset,
+                render_endpoint=render_endpoint,
+                processor=processor,
+                max_length=seq_length,
+                chat_template_kwargs=chat_template_kwargs,
+                turn_dropout=turn_dropout,
+                minimum_valid_tokens=minimum_valid_tokens,
+            )
+        else:
+            if turn_dropout:
+                log.info("Turn dropout enabled: randomly keeping N consecutive turns")
+
+            preprocessed_dataset = build_eagle3_dataset(
+                dataset=raw_dataset,
+                processor=processor,
+                max_length=seq_length,
+                num_proc=build_dataset_num_proc,
+                assistant_pattern=assistant_pattern,
+                turn_dropout=turn_dropout,
+                minimum_valid_tokens=minimum_valid_tokens,
+            )
         if minimum_valid_tokens is not None:
             log.info(f"Kept {len(preprocessed_dataset)} samples after filtering")
         processed_datasets.append(preprocessed_dataset)