diff --git a/pyproject.toml b/pyproject.toml index 50c60cb0d..0393d4c3d 100644 --- a/pyproject.toml +++ b/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", diff --git a/scripts/prepare_data.py b/scripts/prepare_data.py index 009a0b769..8121b3aa8 100644 --- a/scripts/prepare_data.py +++ b/scripts/prepare_data.py @@ -24,6 +24,7 @@ import argparse import glob +import json import logging import shutil import sys @@ -141,6 +142,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( @@ -224,6 +247,12 @@ def main(): else: output.mkdir(parents=True) + 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, train_data_paths=args.data, @@ -237,6 +266,8 @@ def main(): minimum_valid_tokens=args.minimum_valid_tokens, allow_empty_output=args.allow_empty_output, trust_remote_code=args.trust_remote_code, + render_endpoint=args.vllm_server_url, + chat_template_kwargs=chat_template_kwargs, ) log.info("Done preparing data") diff --git a/src/speculators/data_generation/preprocessing.py b/src/speculators/data_generation/preprocessing.py index fc1574fe7..c56215dd9 100644 --- a/src/speculators/data_generation/preprocessing.py +++ b/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", @@ -656,6 +660,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. @@ -803,36 +932,56 @@ def load_and_preprocess_dataset( minimum_valid_tokens: int | None = None, allow_empty_output: bool = False, 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. allow_empty_output: If True, allow returning an empty dataset instead of - raising when no samples survive preprocessing. + raising when no samples survive preprocessing. 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( @@ -842,7 +991,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." @@ -869,18 +1022,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) diff --git a/src/speculators/data_generation/render_client.py b/src/speculators/data_generation/render_client.py new file mode 100644 index 000000000..5d0e3a9a0 --- /dev/null +++ b/src/speculators/data_generation/render_client.py @@ -0,0 +1,64 @@ +"""Client for vLLM's /v1/chat/completions/render endpoint (RFC #652).""" + +from http import HTTPStatus + +import httpx + +from speculators.data_generation.vllm_client import InvalidResponseError, with_retries + +DEFAULT_RENDER_TIMEOUT = 30 + + +class RenderError(Exception): + """Non-200 response from the render endpoint.""" + + +@with_retries +def render_conversation( + endpoint: str, + messages: list[dict], + *, + tools: list[dict] | None = None, + chat_template_kwargs: dict | None = None, + max_length: int | None = None, + timeout: float = DEFAULT_RENDER_TIMEOUT, +) -> dict: + """POST to /v1/chat/completions/render and return token_ids + loss_mask.""" + url = f"{endpoint.rstrip('/')}/v1/chat/completions/render" + + body: dict = { + "messages": messages, + "add_generation_prompt": False, + "return_assistant_tokens_mask": True, + } + if tools is not None: + body["tools"] = tools + if chat_template_kwargs is not None: + body["chat_template_kwargs"] = chat_template_kwargs + if max_length is not None: + body["truncate_prompt_tokens"] = max_length + body["truncation_side"] = "right" + + resp = httpx.post(url, json=body, timeout=timeout) + + if HTTPStatus.BAD_REQUEST <= resp.status_code < HTTPStatus.INTERNAL_SERVER_ERROR: + # Deterministic client error (bad request, wrong URL) -- retrying wastes + # requests without changing the outcome. InvalidResponseError short- + # circuits @with_retries (see vllm_client._handle_retry_error). + raise InvalidResponseError( + f"Render endpoint returned {resp.status_code}: {resp.text[:500]}" + ) + if resp.status_code != HTTPStatus.OK: + raise RenderError( + f"Render endpoint returned {resp.status_code}: {resp.text[:500]}" + ) + + data = resp.json() + if "token_ids" not in data: + raise RenderError( + f"Render endpoint response missing 'token_ids': {str(data)[:500]}" + ) + return { + "token_ids": data["token_ids"], + "loss_mask": data.get("assistant_tokens_mask"), + } diff --git a/tests/integration/datagen/test_render_multimodal.py b/tests/integration/datagen/test_render_multimodal.py new file mode 100644 index 000000000..5488cb508 --- /dev/null +++ b/tests/integration/datagen/test_render_multimodal.py @@ -0,0 +1,150 @@ +"""Faithful multimodal test for the render fallback — real Qwen3-VL render server. + +Heavy/local test (skips in CI). Reproduces a real failure: vLLM does NOT truncate +multimodal prompts to ``truncate_prompt_tokens``, so for a conversation that exceeds +``seq_length`` the server returns more ``token_ids`` than ``max_length`` while the +local regex fallback truncates the mask to ``max_length`` -> ``input_ids`` and +``loss_mask`` end up different lengths. + +Set ``RENDER_ENDPOINT`` to reuse an already-running Qwen3-VL render server; otherwise +one is launched from the sibling vLLM venv. +""" + +import os +import shutil +import socket +import subprocess +import time + +import numpy as np +import pytest + +try: + import httpx +except ImportError: + pytest.skip("httpx not available", allow_module_level=True) + +from datasets import Dataset as HFDataset +from PIL import Image + +from speculators.data_generation.preprocessing import ( + build_dataset_from_render, + load_processor, +) +from tests.e2e.utils import VLLM_PYTHON + +MODEL = "Qwen/Qwen3-VL-2B-Instruct" + + +def _find_free_port(): + with socket.socket() as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +@pytest.fixture(scope="module") +def test_dir(tmp_path_factory): + """Module-scoped isolated temp dir: holds the generated image and doubles + as the server's --allowed-local-media-path root.""" + return tmp_path_factory.mktemp("speculators_render") + + +@pytest.fixture(scope="module") +def image(test_dir): + img_path = test_dir / "render_mm_test.png" + arr = np.zeros((128, 128, 3), dtype=np.uint8) + arr[:64, :, 0] = 200 + arr[64:, :, 2] = 200 + arr[32:96, 32:96, 1] = 180 + Image.fromarray(arr).save(img_path) + return str(img_path) + + +@pytest.fixture(scope="module") +def render_server(test_dir): + env_endpoint = os.environ.get("RENDER_ENDPOINT") + if env_endpoint: + yield env_endpoint.rstrip("/") + return + + if shutil.which(VLLM_PYTHON) is None: + pytest.skip(f"vLLM python not found at {VLLM_PYTHON}") + + port = _find_free_port() + proc = subprocess.Popen( # noqa: S603 + [ + VLLM_PYTHON, + "-m", + "vllm.entrypoints.cli.main", + "launch", + "render", + "--model", + MODEL, + "--port", + str(port), + "--host", + "127.0.0.1", + "--max-model-len", + "4096", + "--gpu-memory-utilization", + "0.3", + "--allowed-local-media-path", + str(test_dir), + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + env={**os.environ, "HF_HUB_OFFLINE": "1"}, + ) + url = f"http://127.0.0.1:{port}" + for _ in range(120): + try: + if httpx.get(f"{url}/health", timeout=2).status_code == 200: + break + except httpx.ConnectError: + pass + time.sleep(2) + else: + proc.kill() + pytest.skip("render server failed to start") + yield url + proc.terminate() + proc.wait(timeout=10) + + +@pytest.mark.sanity +def test_multimodal_over_length_not_misaligned(render_server, image): + """A multimodal conversation that exceeds max_length must never be emitted as + a misaligned (input_ids vs loss_mask) pair. vLLM won't truncate the image + prompt, so the render path drops the row -- parity with the default pipeline, + which drops over-length multimodal rows too. (Before the fix this emitted a + 86-vs-48 misaligned sample.)""" + processor = load_processor(MODEL) + conv = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe the colors."}, + {"type": "image", "path": image}, + ], + }, + { + "role": "assistant", + "content": [{"type": "text", "text": "Red, blue, green."}], + }, + ] + # max_length below the image's own token count: vLLM can't fit it. + ds = build_dataset_from_render( + HFDataset.from_dict({"conversations": [conv]}), + render_server, + processor, + max_length=48, + ) + + for i in range(len(ds)): + n_ids = len(ds[i]["input_ids"]) + n_mask = len(ds[i]["loss_mask"]) + assert n_ids == n_mask, ( + f"misaligned sample emitted: input_ids={n_ids} loss_mask={n_mask}" + ) + # over-length multimodal row is dropped (parity with the default pipeline) + assert len(ds) == 0 diff --git a/tests/integration/datagen/test_render_truncation.py b/tests/integration/datagen/test_render_truncation.py new file mode 100644 index 000000000..3baec633b --- /dev/null +++ b/tests/integration/datagen/test_render_truncation.py @@ -0,0 +1,123 @@ +"""Test that render endpoint truncation matches HF truncation. + +Heavy/local test — spins up a real vLLM render server. Skips in CI +(requires vllm installed in a sibling venv). + +Proves the render path is a drop-in replacement for the HF path: +same conversation + same seq_length → same truncated token_ids. +""" + +import os +import shutil +import socket +import subprocess +import time + +import pytest + +try: + import httpx +except ImportError: + pytest.skip("httpx not available", allow_module_level=True) + +from tests.e2e.utils import VLLM_PYTHON + +MODEL = "Qwen/Qwen3-0.6B" +# Short enough to force truncation on a multi-turn conversation. +SEQ_LENGTH = 30 + + +def _find_free_port(): + with socket.socket() as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +@pytest.fixture(scope="module") +def render_server(): + if shutil.which(VLLM_PYTHON) is None: + pytest.skip(f"vLLM python not found at {VLLM_PYTHON}") + + port = _find_free_port() + proc = subprocess.Popen( # noqa: S603 + [ + VLLM_PYTHON, + "-m", + "vllm.entrypoints.cli.main", + "launch", + "render", + "--model", + MODEL, + "--port", + str(port), + "--host", + "127.0.0.1", + "--max-model-len", + "4096", + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + env={**os.environ, "HF_HUB_OFFLINE": "1"}, + ) + url = f"http://127.0.0.1:{port}" + for _ in range(60): + try: + if httpx.get(f"{url}/health", timeout=2).status_code == 200: + break + except httpx.ConnectError: + pass + time.sleep(1) + else: + proc.kill() + pytest.skip("render server failed to start") + + yield url + proc.terminate() + proc.wait(timeout=10) + + +@pytest.mark.sanity +def test_render_truncation_matches_hf(render_server): + """Truncated token_ids from the render endpoint must equal + HF apply_chat_template with the same max_length.""" + from transformers import AutoTokenizer # noqa: PLC0415 + + tok = AutoTokenizer.from_pretrained(MODEL) + + conv = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Tell me about the history of Paris."}, + { + "role": "assistant", + "content": "Paris has a long and fascinating history.", + }, + {"role": "user", "content": "What about Rome?"}, + {"role": "assistant", "content": "Rome is the Eternal City."}, + ] + + # HF path: what speculators does today. + full_ids = tok.apply_chat_template( + conv, + tokenize=True, + add_generation_prompt=False, + return_dict=False, + ) + assert len(full_ids) > SEQ_LENGTH, "conversation too short to test truncation" + hf_truncated = list(full_ids[:SEQ_LENGTH]) + + # Render path: POST with truncate_prompt_tokens. + resp = httpx.post( + f"{render_server}/v1/chat/completions/render", + json={ + "model": MODEL, + "messages": conv, + "add_generation_prompt": False, + "truncate_prompt_tokens": SEQ_LENGTH, + "truncation_side": "right", + }, + timeout=30, + ) + assert resp.status_code == 200, resp.text + render_ids = resp.json()["token_ids"] + + assert render_ids == hf_truncated diff --git a/tests/unit/data_generation/test_build_dataset_from_render.py b/tests/unit/data_generation/test_build_dataset_from_render.py new file mode 100644 index 000000000..4ede9c981 --- /dev/null +++ b/tests/unit/data_generation/test_build_dataset_from_render.py @@ -0,0 +1,189 @@ +"""Tests for build_dataset_from_render: server-provided and regex-fallback masks.""" + +from unittest.mock import patch + +import pytest +from datasets import Dataset as HFDataset +from transformers import AutoTokenizer + +from speculators.data_generation.preprocessing import ( + build_dataset_from_render, + load_processor, +) +from speculators.data_generation.render_client import RenderError + +MODEL = "Qwen/Qwen2.5-0.5B-Instruct" + +CONV = [ + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "The capital of France is Paris."}, +] + + +@pytest.fixture(scope="module") +def processor(): + return load_processor(MODEL) + + +@pytest.fixture(scope="module") +def tokenizer(): + return AutoTokenizer.from_pretrained(MODEL) + + +def _make_dataset(): + return HFDataset.from_dict({"conversations": [CONV]}) + + +def _mock_render(token_ids, loss_mask=None): + """Patch render_conversation to return fixed token_ids + loss_mask.""" + + def _fake(endpoint, messages, **kwargs): + return {"token_ids": token_ids, "loss_mask": loss_mask} + + return patch( + "speculators.data_generation.preprocessing.render_conversation", + side_effect=_fake, + ) + + +@pytest.mark.sanity +def test_server_provided_mask_flows_through(processor, tokenizer): + """When the render endpoint returns a real loss_mask, use it directly.""" + ids = tokenizer.apply_chat_template( + CONV, tokenize=True, add_generation_prompt=False, return_dict=False + ) + mask = [0] * len(ids) + # Mark the last 5 tokens as trainable (simulates assistant span). + for i in range(len(ids) - 5, len(ids)): + mask[i] = 1 + + with _mock_render(list(ids), loss_mask=mask): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=4096 + ) + + assert len(ds) == 1 + assert ds[0]["input_ids"].tolist() == list(ids) + assert ds[0]["loss_mask"].tolist() == mask + assert ds[0]["loss_mask"].sum().item() == 5 + + +@pytest.mark.sanity +def test_regex_fallback_when_loss_mask_is_none(processor, tokenizer): + """When loss_mask=None, fall back to regex-based assistant detection.""" + ids = tokenizer.apply_chat_template( + CONV, tokenize=True, add_generation_prompt=False, return_dict=False + ) + + with _mock_render(list(ids), loss_mask=None): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=4096 + ) + + assert len(ds) == 1 + assert ds[0]["input_ids"].tolist() == list(ids) + mask = ds[0]["loss_mask"] + # The regex fallback should mark assistant tokens as trainable. + assert mask.sum().item() > 0, "regex fallback produced all-zero mask" + # Verify the masked-in tokens decode to the assistant content. + masked_ids = [t for t, m in zip(ids, mask.tolist(), strict=False) if m] + decoded = tokenizer.decode(masked_ids) + assert "Paris" in decoded + + +@pytest.mark.sanity +def test_full_server_mask_skips_local_pattern_detection(processor, tokenizer): + """When the server masks every conversation, the local regex pattern is + never detected -- so a model whose detection would raise still succeeds.""" + ids = tokenizer.apply_chat_template( + CONV, tokenize=True, add_generation_prompt=False, return_dict=False + ) + mask = [0] * len(ids) + mask[-1] = 1 + + with ( + _mock_render(list(ids), loss_mask=mask), + patch( + "speculators.data_generation.preprocessing._detect_assistant_pattern", + side_effect=ValueError("cannot detect"), + ) as detect, + ): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=4096 + ) + + detect.assert_not_called() + assert len(ds) == 1 + assert ds[0]["loss_mask"].sum().item() == 1 + + +@pytest.mark.sanity +def test_over_length_render_dropped(processor): + """vLLM does not truncate multimodal prompts, so the server can return + token_ids longer than max_length. Such a row is dropped (parity with the + default path), never emitted as an over-length / misaligned sample.""" + long_ids = list(range(100)) # 100 > max_length (50) below + with _mock_render(long_ids, loss_mask=None): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=50 + ) + assert len(ds) == 0 + + +@pytest.mark.sanity +def test_render_error_does_not_crash_the_run(processor): + """A per-conversation render failure is skipped, not propagated -- one bad + row must not abort the whole dataset build (parity with the default path).""" + with patch( + "speculators.data_generation.preprocessing.render_conversation", + side_effect=RenderError("boom"), + ): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=4096 + ) + assert len(ds) == 0 + + +@pytest.fixture(scope="module") +def mm_processor(): + return load_processor("Qwen/Qwen3-VL-2B-Instruct") + + +@pytest.mark.sanity +def test_multimodal_keeps_messages_for_extraction(mm_processor): + """Multimodal (ProcessorMixin) rows keep a `messages` column so hidden-state + extraction can re-send the conversation with images via the chat API -- + token_ids alone cannot carry the image.""" + mm_conv = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + {"type": "image", "path": "/tmp/x.png"}, + ], + }, + {"role": "assistant", "content": [{"type": "text", "text": "A cat."}]}, + ] + ids = list(range(12)) + mask = [0] * 9 + [1, 1, 1] + with _mock_render(ids, loss_mask=mask): + ds = build_dataset_from_render( + HFDataset.from_dict({"conversations": [mm_conv]}), + "http://fake", + mm_processor, + max_length=4096, + ) + assert "messages" in ds.column_names + assert len(ds) == 1 + assert ds[0]["messages"][0]["role"] == "user" + + +@pytest.mark.sanity +def test_text_model_has_no_messages_column(processor): + """Text-only processors must not add a `messages` column (parity with the + default path, which only adds it for ProcessorMixin).""" + with _mock_render(list(range(10)), loss_mask=[0] * 7 + [1, 1, 1]): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=4096 + ) + assert "messages" not in ds.column_names diff --git a/tests/unit/data_generation/test_load_and_preprocess_dataset.py b/tests/unit/data_generation/test_load_and_preprocess_dataset.py new file mode 100644 index 000000000..febfe126b --- /dev/null +++ b/tests/unit/data_generation/test_load_and_preprocess_dataset.py @@ -0,0 +1,39 @@ +"""Tests for load_and_preprocess_dataset's render-path argument validation. + +Both checks fire before target_model_path / train_data_paths are ever touched +(no processor load, no dataset load), so these run with no mocking. +""" + +import pytest + +from speculators.data_generation.preprocessing import load_and_preprocess_dataset + + +@pytest.mark.sanity +def test_assistant_pattern_rejected_with_render_endpoint(): + """assistant_pattern is silently unreachable on the render path (it always + auto-detects its own pattern) -- reject the combination instead of + dropping the user's pattern without telling them.""" + with pytest.raises(ValueError, match="assistant_pattern"): + load_and_preprocess_dataset( + target_model_path="unused", + train_data_paths=["unused"], + seq_length=10, + assistant_pattern="custom-pattern", + render_endpoint="http://localhost:8000", + ) + + +@pytest.mark.sanity +def test_chat_template_kwargs_rejected_without_render_endpoint(): + """chat_template_kwargs only reaches build_dataset_from_render; the default + HF path (build_eagle3_dataset) never receives it. Reject rather than + silently ignore a value the user explicitly set.""" + with pytest.raises(ValueError, match="chat_template_kwargs"): + load_and_preprocess_dataset( + target_model_path="unused", + train_data_paths=["unused"], + seq_length=10, + chat_template_kwargs={"enable_thinking": True}, + render_endpoint=None, + ) diff --git a/tests/unit/data_generation/test_render_client.py b/tests/unit/data_generation/test_render_client.py new file mode 100644 index 000000000..09542fee7 --- /dev/null +++ b/tests/unit/data_generation/test_render_client.py @@ -0,0 +1,133 @@ +"""Tests for render_client.py.""" + +import json + +import httpx +import pytest + +import speculators.data_generation.render_client as render_mod +from speculators.data_generation.render_client import ( + RenderError, + render_conversation, +) +from speculators.data_generation.vllm_client import InvalidResponseError + +MOCK_TOKEN_IDS = [ + 151644, + 872, + 198, + 9707, + 151645, + 198, + 151644, + 77091, + 198, + 13048, + 0, + 151645, +] + +MESSAGES = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, +] + + +def _make_transport(*, token_ids=None, assistant_tokens_mask=None): + def handler(request: httpx.Request) -> httpx.Response: + body = {"token_ids": token_ids or MOCK_TOKEN_IDS} + if assistant_tokens_mask is not None: + body["assistant_tokens_mask"] = assistant_tokens_mask + return httpx.Response(200, json=body) + + return httpx.MockTransport(handler) + + +def _capturing_transport(): + """Return (transport, captured_body_dict).""" + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured.update(json.loads(request.content)) + return httpx.Response(200, json={"token_ids": MOCK_TOKEN_IDS}) + + return httpx.MockTransport(handler), captured + + +def _counting_transport(*, status_code=200, body=None): + """Return (transport, call_counter) -- counter tracks requests received, + to assert whether a failure was retried.""" + calls = {"count": 0} + + def handler(request: httpx.Request) -> httpx.Response: + calls["count"] += 1 + if status_code != 200: + return httpx.Response(status_code, text="error from mock") + return httpx.Response(200, json=body or {"token_ids": MOCK_TOKEN_IDS}) + + return httpx.MockTransport(handler), calls + + +def _call_render(transport, *, max_retries=0, **kwargs): + original_post = httpx.post + + def mock_post(url, **kw): + return httpx.Client(transport=transport).request("POST", url, **kw) + + render_mod.httpx.post = mock_post # type: ignore[attr-defined] + try: + return render_conversation( + "http://localhost:8000", MESSAGES, max_retries=max_retries, **kwargs + ) + finally: + render_mod.httpx.post = original_post # type: ignore[attr-defined] + + +class TestRenderConversation: + def test_basic_render(self): + result = _call_render(_make_transport()) + assert result["token_ids"] == MOCK_TOKEN_IDS + assert result["loss_mask"] is None # not yet returned by endpoint + + def test_loss_mask_passthrough(self): + mask = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0] + result = _call_render(_make_transport(assistant_tokens_mask=mask)) + assert result["loss_mask"] == mask + + @pytest.mark.parametrize( + ("status_code", "body", "max_retries", "expected_exc", "expected_calls"), + [ + # 4xx is deterministic (bad request, wrong URL) -- retrying wastes + # requests without changing the outcome. + (400, None, 3, InvalidResponseError, 1), + # 5xx may be transient -- goes through the normal retry policy + # (initial attempt + max_retries). + (500, None, 2, RenderError, 3), + # 200 but missing token_ids must raise RenderError, not a bare + # KeyError that build_dataset_from_render's except clause can't + # catch -- and it's not exempt from the retry policy either. + (200, {"assistant_tokens_mask": [0, 1]}, 1, RenderError, 2), + ], + ) + def test_error_handling_and_retry_policy( + self, status_code, body, max_retries, expected_exc, expected_calls + ): + transport, calls = _counting_transport(status_code=status_code, body=body) + with pytest.raises(expected_exc): + _call_render(transport, max_retries=max_retries) + assert calls["count"] == expected_calls + + def test_request_body_forwarding(self): + transport, body = _capturing_transport() + tools = [{"type": "function", "function": {"name": "f", "parameters": {}}}] + _call_render( + transport, + tools=tools, + chat_template_kwargs={"enable_thinking": True}, + max_length=4096, + ) + + assert body["tools"] == tools + assert body["chat_template_kwargs"] == {"enable_thinking": True} + assert body["add_generation_prompt"] is False + assert body["truncate_prompt_tokens"] == 4096