Revert "[Data] Derive off-policy loss masks from vLLM render boundaries" - #943
Conversation
|
Note Currently processing new changes in this PR. This may take a few minutes, please wait... ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (15)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Merge Protections🟢 Merge protection satisfied — ready to merge. Show 1 satisfied protection🟢 Require approval from approved reviewers listAll pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (3)
tests/integration/datagen/test_preprocessing.py (2)
319-322: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDrop the parenthesis-counting assertion.
re.compileat Line 316 already proves the pattern is a valid regular expression, which implies balanced groups. The count assertion instead compares raw(and)characters._detect_assistant_patternbuilds the pattern withre.escape(role_marker), so a chat template whose role marker contains a literal(produces\(in the pattern and breaks this assertion even though the regex is valid. Assert the structural property that matters instead: the pattern exposes group 1, which_create_loss_mask_from_offsetsreads.💚 Proposed fix
- # Pattern should contain balanced parentheses - assert pattern.count("(") == pattern.count(")") - # Pattern should have at least one capture group (may use negative lookahead) - assert "(" in pattern, "Pattern should have a capture group for content" + # The masking code reads group 1, so the pattern must expose one capture group. + assert compiled.groups >= 1, "Pattern should have a capture group for content"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/datagen/test_preprocessing.py` around lines 319 - 322, Remove the raw parenthesis-count assertion in the test around _detect_assistant_pattern, since re.compile already validates grouping and escaped literal parentheses are valid. Replace it with an assertion that the compiled pattern exposes capture group 1, matching the group consumed by _create_loss_mask_from_offsets; retain the existing requirement that the pattern contains a capture group.
683-685: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace the conditional and inequality assertions with exact assertions.
These assertions pass when zero rows survive. A regression that drops every conversation would not fail the test. The same weakness appears in several of the new tests:
- Lines 684-685:
assert len(results["input_ids"]) <= 1passes for a result of0. The valid conversation at Line 674 must survive, so assert== 1.- Lines 714-717: the
if len(results["input_ids"]) > 0:guard makes the entire truncation check vacuous. Assert that one row survives, then assert the length bound.- Lines 926-931:
assert len(result) <= len(dataset)plusif len(result) > 0:makestest_build_eagle3_dataset_basicpass on an empty dataset. Assertlen(result) == 2.- Line 980: the same guard in
test_build_eagle3_dataset_removes_original_columns. Assertlen(result) == 1first.💚 Proposed fix for Lines 683-685
- # Should only process the valid conversation - assert len(results["input_ids"]) <= 1 - assert len(results["loss_mask"]) <= 1 + # Only the valid conversation is processed; the None and empty rows are skipped. + assert len(results["input_ids"]) == 1 + assert len(results["loss_mask"]) == 1💚 Proposed fix for Lines 714-717
- if len(results["input_ids"]) > 0: - # Should be truncated to max_length - assert len(results["input_ids"][0]) <= max_length - assert len(results["loss_mask"][0]) <= max_length + assert len(results["input_ids"]) == 1 + # Should be truncated to max_length + assert len(results["input_ids"][0]) == max_length + assert len(results["loss_mask"][0]) == max_lengthAs per path instructions: "Ensure PyTest tests are clear, comprehensive, and cover edge cases specific to speculative decoding" and "Check that new code paths introduced in the PR are covered."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/datagen/test_preprocessing.py` around lines 683 - 685, Strengthen the affected preprocessing tests by replacing permissive row-count checks with exact survival assertions: require one row in the test around results["input_ids"] and results["loss_mask"], assert one row before the truncation length bound instead of guarding it with a nonempty check, require len(result) == 2 in test_build_eagle3_dataset_basic, and require len(result) == 1 in test_build_eagle3_dataset_removes_original_columns before validating columns or bounds.Source: Path instructions
scripts/prepare_data.py (1)
128-133: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winValidate
--assistant-patternat parse time.
argparseaccepts any string here. An invalid regular expression reachesre.finditerinside_create_loss_mask_from_offsets, one call per conversation._preprocess_batchcatches the resultingre.errorin its broadexcept Exceptionat Line 592 and logs one error per row. Every row is dropped, and the run then fails with "No samples remain after preprocessing". Compile the pattern during argument parsing so the user sees the real cause immediately.♻️ Proposed refactor to validate the pattern
+def _compiled_regex(value: str) -> str: + try: + re.compile(value) + except re.error as e: + raise argparse.ArgumentTypeError( + f"Invalid --assistant-pattern regex: {e}" + ) from e + return value + + def parse_args():parser.add_argument( "--assistant-pattern", - type=str, + type=_compiled_regex, default=None, help=( "Custom regex pattern for matching assistant responses. " "If not provided, auto-detected from chat template." ), )Add
import reat the top of the file.As per path instructions: "Check that scripts handle argument parsing robustly, log progress clearly, and are safe to run in multi-GPU environments."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@scripts/prepare_data.py` around lines 128 - 133, Validate the --assistant-pattern value during argument parsing by compiling it with re, using argparse’s parser error mechanism for invalid expressions. Add the required re import and update the argument definition around _create_loss_mask_from_offsets without changing valid-pattern behavior or deferred preprocessing handling.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/cli/prepare_data.md`:
- Around line 33-37: Update the --data description in prepare_data.md to use a
complete sentence and document the hf:<dataset_id>[:<subset>:<split>] dataset
form supported by _load_hf_dataset, including that the split defaults to train.
In `@docs/cli/response_regeneration.md`:
- Line 143: Align the regenerated-row schema descriptions with the actual
producer contract by documenting exactly one review-only field, including its
purpose and usage. Update docs/cli/response_regeneration.md lines 143-143 and
docs/user_guide/tutorials/response_regeneration.md lines 81-81 to use the same
field name and explanation, keeping both Markdown examples clear, complete, and
consistent with the current API.
In `@docs/user_guide/tutorials/train.md`:
- Around line 132-135: Update the --data parameter description in the
“Parameters explained” section to give “Can be supplied multiple times” an
explicit subject, while preserving its meaning that the data argument supports
combining multiple datasets.
In `@scripts/prepare_data.py`:
- Line 104: Update the --data argument help text in the argument parser to
describe that it accepts dataset preset names, HuggingFace dataset names or
specs, and local paths, rather than calling it a preprocessing data path.
In `@src/speculators/data_generation/preprocessing.py`:
- Around line 259-294: Validate the `USER_MSG_2` lookup before computing
`second_user_end` or slicing `prefix`; raise a `ValueError` when the marker is
absent, alongside the existing assistant-marker validation. Store the validated
position in a reusable `second_user_start` symbol and use it when computing
`suffix1` instead of performing a second unguarded
`formatted.find("USER_MSG_2")` call.
- Line 339: Change the regex preprocessing path’s loss_mask allocation in the
relevant preprocessing function from torch.bool to torch.long, matching the
native assistant-mask and _passthrough_pretokenized paths. Update the
integration test assertion that expects torch.bool to expect torch.long, while
preserving the existing mask value comparisons.
- Around line 443-451: The multimodal preprocessing path currently derives
loss-mask text from decoded input IDs instead of the exact rendered text used
for offset_mapping. Update the flow around encoded, input_ids, and
formatted_text to retain and use the tokenized apply_chat_template text for mask
matching, then assert that processor.decode(encoded["input_ids"][0]) matches
that text character-for-character.
In `@tests/integration/datagen/test_regex_patterns.py`:
- Around line 45-49: Broaden the exception handling in the processor-loading
fixture around load_processor to catch unavailable-model failures, including
huggingface_hub GatedRepoError, RepositoryNotFoundError, and HfHubHTTPError,
plus Transformers’ OSError, so gated, missing, or unreachable models call
pytest.skip with the existing context. Preserve handling of the currently listed
exception types.
In `@tests/unit/scripts/test_response_regeneration.py`:
- Line 857: Correct the section around “Every shared-registry preset works
on-policy” so its claim matches actual coverage: either narrow the title to
supported text presets, or include the excluded sharegpt4v_coco case in the test
with an explicit assertion while respecting REGEN_DATASETS and script.py
rejection behavior.
---
Nitpick comments:
In `@scripts/prepare_data.py`:
- Around line 128-133: Validate the --assistant-pattern value during argument
parsing by compiling it with re, using argparse’s parser error mechanism for
invalid expressions. Add the required re import and update the argument
definition around _create_loss_mask_from_offsets without changing valid-pattern
behavior or deferred preprocessing handling.
In `@tests/integration/datagen/test_preprocessing.py`:
- Around line 319-322: Remove the raw parenthesis-count assertion in the test
around _detect_assistant_pattern, since re.compile already validates grouping
and escaped literal parentheses are valid. Replace it with an assertion that the
compiled pattern exposes capture group 1, matching the group consumed by
_create_loss_mask_from_offsets; retain the existing requirement that the pattern
contains a capture group.
- Around line 683-685: Strengthen the affected preprocessing tests by replacing
permissive row-count checks with exact survival assertions: require one row in
the test around results["input_ids"] and results["loss_mask"], assert one row
before the truncation length bound instead of guarding it with a nonempty check,
require len(result) == 2 in test_build_eagle3_dataset_basic, and require
len(result) == 1 in test_build_eagle3_dataset_removes_original_columns before
validating columns or bounds.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: c5c65b32-d9f5-41dc-8586-ec6562bdb5bb
📒 Files selected for processing (15)
docs/cli/prepare_data.mddocs/cli/response_regeneration.mddocs/user_guide/tutorials/response_regeneration.mddocs/user_guide/tutorials/train.mdpyproject.tomlscripts/prepare_data.pyscripts/response_regeneration/script.pysrc/speculators/data_generation/preprocessing.pysrc/speculators/data_generation/render_client.pytests/e2e/smoke/test_render_boundary.pytests/integration/datagen/test_preprocessing.pytests/integration/datagen/test_regex_patterns.pytests/integration/datagen/test_render_boundary.pytests/unit/scripts/test_response_regeneration.pytests/unit/train/test_prepare_data.py
💤 Files with no reviewable changes (4)
- pyproject.toml
- tests/e2e/smoke/test_render_boundary.py
- tests/integration/datagen/test_render_boundary.py
- src/speculators/data_generation/render_client.py
| - **`--data`** (str, required, repeatable) Path to training data. Can be a HuggingFace dataset name or local path. Use multiple times to specify multiple datasets. | ||
|
|
||
| Example: `--data ./target_responses.jsonl --data hf:my-org/more-target-responses` | ||
| Example: `--data sharegpt --data ./custom_data.jsonl` | ||
|
|
||
| Natural-language input uses a `conversations` column and requires `--render-endpoint`. Assistant responses must already have been produced by the target model. Tool-calling datasets may also include a separate `tools` column. Speculator-format input uses `input_ids` and `loss_mask`. | ||
| The input conversation should be provided in the `conversations` column. Tool-calling datasets that include separate columns for tools are also supported, as demonstrated in [llamafactory/reason-tool-use-demo-1500](https://huggingface.co/datasets/llamafactory/reason-tool-use-demo-1500) and [interstellarninja/hermes_reasoning_tool_use](https://huggingface.co/datasets/interstellarninja/hermes_reasoning_tool_use). |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document the hf: dataset spec and complete the sentence.
load_raw_dataset supports a third input form that this section omits: _load_hf_dataset parses hf:<dataset_id>[:<subset>:<split>] and defaults the split to train. A reader following this doc cannot discover it. The first sentence is also a fragment, which LanguageTool flags.
🐛 Proposed fix for the `--data` description
-- **`--data`** (str, required, repeatable) Path to training data. Can be a HuggingFace dataset name or local path. Use multiple times to specify multiple datasets.
-
- Example: `--data sharegpt --data ./custom_data.jsonl`
+- **`--data`** (str, required, repeatable) Training data source. The value can be a built-in dataset preset name, an `hf:<dataset_id>[:<subset>:<split>]` spec for any HuggingFace dataset already in conversations format, or a local JSON/JSONL path. Use the flag multiple times to combine datasets.
+
+ Example: `--data sharegpt --data hf:my-org/my-dataset:train --data ./custom_data.jsonl`As per path instructions: "Check for clarity, accuracy, and completeness."
🧰 Tools
🪛 LanguageTool
[style] ~33-~33: To form a complete sentence, be sure to include a subject or ‘there’.
Context: ...red, repeatable) Path to training data. Can be a HuggingFace dataset name or local ...
(MISSING_IT_THERE)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/cli/prepare_data.md` around lines 33 - 37, Update the --data description
in prepare_data.md to use a complete sentence and document the
hf:<dataset_id>[:<subset>:<split>] dataset form supported by _load_hf_dataset,
including that the split defaults to train.
Sources: Path instructions, Linters/SAST tools
| ## Output Format | ||
|
|
||
| Rows are in speculator format and ready for training: one row per target generation, holding the prompt the target conditioned on followed by the tokens it generated. The endpoint must support `return_token_ids`, which the script uses to read the generation boundary directly instead of re-tokenizing the text and recovering the boundary with a regex. | ||
| Rows are pre-tokenized and ready for training: one row per target generation, holding the prompt the target conditioned on followed by the tokens it generated. The endpoint must support `return_token_ids`, which the script uses to read the generation boundary directly instead of re-tokenizing the text and recovering the boundary with a regex. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Use one documented review-only field for regenerated rows.
The two output descriptions show different schemas. Align them with the actual producer contract so users do not build incompatible downstream parsers.
docs/cli/response_regeneration.md#L143-L143: document whether the review-only field istext,conversations, or both.docs/user_guide/tutorials/response_regeneration.md#L81-L81: use the same field and explanation as the CLI documentation.
As per path instructions, keep both Markdown examples clear, accurate, complete, and aligned with the current API.
📍 Affects 2 files
docs/cli/response_regeneration.md#L143-L143(this comment)docs/user_guide/tutorials/response_regeneration.md#L81-L81
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/cli/response_regeneration.md` at line 143, Align the regenerated-row
schema descriptions with the actual producer contract by documenting exactly one
review-only field, including its purpose and usage. Update
docs/cli/response_regeneration.md lines 143-143 and
docs/user_guide/tutorials/response_regeneration.md lines 81-81 to use the same
field name and explanation, keeping both Markdown examples clear, complete, and
consistent with the current API.
Source: Path instructions
| **Parameters explained:** | ||
|
|
||
| - `--model` - The target model you want to accelerate | ||
| - `--data` - On-policy target-model data, either natural-language `conversations` or speculator-format `input_ids` and `loss_mask`. Can be supplied multiple times to combine datasets. | ||
| - `--render-endpoint` - Target model's vLLM base URL; required only for natural-language conversations. | ||
| - `--data` - Dataset to use (built-in support for `sharegpt`, `ultrachat`. Otherwise provide a custom path to a jsonl file). Can be supplied multiple times to combine multiple datasets. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Line 135: Add a subject to the final sentence.
Can be supplied multiple times is unclear because it has no subject. Use explicit wording.
Proposed wording
-- `--data` - Dataset to use (built-in support for `sharegpt`, `ultrachat`. Otherwise provide a custom path to a jsonl file). Can be supplied multiple times to combine multiple datasets.
+- `--data` - Dataset to use (built-in support for `sharegpt`, `ultrachat`. Otherwise provide a custom path to a JSONL file). You can supply this option multiple times to combine datasets.As per path instructions, documentation must be clear, accurate, and complete.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| **Parameters explained:** | |
| - `--model` - The target model you want to accelerate | |
| - `--data` - On-policy target-model data, either natural-language `conversations` or speculator-format `input_ids` and `loss_mask`. Can be supplied multiple times to combine datasets. | |
| - `--render-endpoint` - Target model's vLLM base URL; required only for natural-language conversations. | |
| - `--data` - Dataset to use (built-in support for `sharegpt`, `ultrachat`. Otherwise provide a custom path to a jsonl file). Can be supplied multiple times to combine multiple datasets. | |
| **Parameters explained:** | |
| - `--model` - The target model you want to accelerate | |
| - `--data` - Dataset to use (built-in support for `sharegpt`, `ultrachat`. Otherwise provide a custom path to a JSONL file). You can supply this option multiple times to combine datasets. |
🧰 Tools
🪛 LanguageTool
[style] ~135-~135: To form a complete sentence, be sure to include a subject.
Context: ...provide a custom path to a jsonl file). Can be supplied multiple times to combine m...
(MISSING_IT_THERE)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/user_guide/tutorials/train.md` around lines 132 - 135, Update the --data
parameter description in the “Parameters explained” section to give “Can be
supplied multiple times” an explicit subject, while preserving its meaning that
the data argument supports combining multiple datasets.
Sources: Path instructions, Linters/SAST tools
| "speculator-format input_ids/loss_mask rows. Assistant responses " | ||
| "must come from the target model; this command does not generate them." | ||
| ), | ||
| help="Path to training data (same as used in preprocessing)", |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the --data help text.
prepare_data.py is the preprocessing step, so "same as used in preprocessing" is misleading. The value is also not restricted to a path: the usage example at Line 20 passes sharegpt, a dataset preset name, and docs/cli/prepare_data.md Line 33 documents HuggingFace dataset names, hf: specs, and local paths.
🐛 Proposed fix for the help text
- help="Path to training data (same as used in preprocessing)",
+ help=(
+ "Training data source: a dataset preset name, an hf: spec, or a local "
+ "JSON/JSONL path. Repeat to combine multiple datasets."
+ ),📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| help="Path to training data (same as used in preprocessing)", | |
| help=( | |
| "Training data source: a dataset preset name, an hf: spec, or a local " | |
| "JSON/JSONL path. Repeat to combine multiple datasets." | |
| ), |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/prepare_data.py` at line 104, Update the --data argument help text in
the argument parser to describe that it accepts dataset preset names,
HuggingFace dataset names or specs, and local paths, rather than calling it a
preprocessing data path.
| first_start = formatted.find("ASSISTANT_MSG_1") | ||
| first_end = first_start + len("ASSISTANT_MSG_1") | ||
| second_start = formatted.find("ASSISTANT_MSG_2") | ||
| second_end = second_start + len("ASSISTANT_MSG_2") | ||
|
|
||
| if first_start == -1 or second_start == -1: | ||
| raise ValueError("Could not detect assistant messages in chat template") | ||
|
|
||
| # Extract role marker from before the second assistant message | ||
| second_user_end = formatted.find("USER_MSG_2") + len("USER_MSG_2") | ||
| prefix = formatted[second_user_end:second_start] | ||
|
|
||
| # Find where the assistant role marker starts | ||
| assistant_pos = prefix.rfind("assistant") | ||
| if assistant_pos != -1: | ||
| # Search for a tag start ('<' or '[') before 'assistant' | ||
| role_start = -1 | ||
| for char in ["<", "["]: | ||
| pos = prefix.rfind(char, 0, assistant_pos) | ||
| role_start = max(role_start, pos) | ||
| if role_start != -1: | ||
| role_marker = prefix[role_start:] | ||
| else: | ||
| role_marker = prefix[assistant_pos:] | ||
| else: | ||
| role_marker = prefix | ||
|
|
||
| # Strip <think>...</think> blocks from the role marker. Thinking model | ||
| # templates wrap assistant content in these tags, but the test messages | ||
| # can produce empty blocks (e.g. "<think>\n\n</think>\n") with reasoning models, | ||
| # which then get baked into the regex as literals. Removing them ensures | ||
| # that reasoning stays within the assistant content group. | ||
| role_marker = re.sub(r"<think>.*?</think>\s*", "", role_marker, flags=re.DOTALL) | ||
|
|
||
| # Determine the stable TURN-LEVEL suffix | ||
| suffix1 = formatted[first_end : formatted.find("USER_MSG_2")] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Guard the USER_MSG_2 lookup before you use it as an index.
Lines 264-265 guard the assistant markers, but the USER_MSG_2 lookups at Line 268 and Line 294 are unguarded. If a chat template does not render the user content verbatim, formatted.find("USER_MSG_2") returns -1. Then second_user_end becomes 9 and prefix is sliced from the wrong position, and suffix1 is sliced to len(formatted) - 1. The function still returns a pattern, so the run continues and produces a loss mask over the wrong token span. Fail loudly instead.
🐛 Proposed fix to validate the user marker
- if first_start == -1 or second_start == -1:
- raise ValueError("Could not detect assistant messages in chat template")
-
- # Extract role marker from before the second assistant message
- second_user_end = formatted.find("USER_MSG_2") + len("USER_MSG_2")
- prefix = formatted[second_user_end:second_start]
+ if first_start == -1 or second_start == -1:
+ raise ValueError("Could not detect assistant messages in chat template")
+
+ second_user_start = formatted.find("USER_MSG_2")
+ if second_user_start == -1:
+ raise ValueError("Could not detect user messages in chat template")
+
+ # Extract role marker from before the second assistant message
+ second_user_end = second_user_start + len("USER_MSG_2")
+ prefix = formatted[second_user_end:second_start]Then reuse second_user_start at Line 294:
- suffix1 = formatted[first_end : formatted.find("USER_MSG_2")]
+ suffix1 = formatted[first_end:second_user_start]As per path instructions: "Verify that shift-based alignment of hidden states to target tokens is correct (off-by-one errors are a common bug here)."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| first_start = formatted.find("ASSISTANT_MSG_1") | |
| first_end = first_start + len("ASSISTANT_MSG_1") | |
| second_start = formatted.find("ASSISTANT_MSG_2") | |
| second_end = second_start + len("ASSISTANT_MSG_2") | |
| if first_start == -1 or second_start == -1: | |
| raise ValueError("Could not detect assistant messages in chat template") | |
| # Extract role marker from before the second assistant message | |
| second_user_end = formatted.find("USER_MSG_2") + len("USER_MSG_2") | |
| prefix = formatted[second_user_end:second_start] | |
| # Find where the assistant role marker starts | |
| assistant_pos = prefix.rfind("assistant") | |
| if assistant_pos != -1: | |
| # Search for a tag start ('<' or '[') before 'assistant' | |
| role_start = -1 | |
| for char in ["<", "["]: | |
| pos = prefix.rfind(char, 0, assistant_pos) | |
| role_start = max(role_start, pos) | |
| if role_start != -1: | |
| role_marker = prefix[role_start:] | |
| else: | |
| role_marker = prefix[assistant_pos:] | |
| else: | |
| role_marker = prefix | |
| # Strip <think>...</think> blocks from the role marker. Thinking model | |
| # templates wrap assistant content in these tags, but the test messages | |
| # can produce empty blocks (e.g. "<think>\n\n</think>\n") with reasoning models, | |
| # which then get baked into the regex as literals. Removing them ensures | |
| # that reasoning stays within the assistant content group. | |
| role_marker = re.sub(r"<think>.*?</think>\s*", "", role_marker, flags=re.DOTALL) | |
| # Determine the stable TURN-LEVEL suffix | |
| suffix1 = formatted[first_end : formatted.find("USER_MSG_2")] | |
| first_start = formatted.find("ASSISTANT_MSG_1") | |
| first_end = first_start + len("ASSISTANT_MSG_1") | |
| second_start = formatted.find("ASSISTANT_MSG_2") | |
| second_end = second_start + len("ASSISTANT_MSG_2") | |
| if first_start == -1 or second_start == -1: | |
| raise ValueError("Could not detect assistant messages in chat template") | |
| second_user_start = formatted.find("USER_MSG_2") | |
| if second_user_start == -1: | |
| raise ValueError("Could not detect user messages in chat template") | |
| # Extract role marker from before the second assistant message | |
| second_user_end = second_user_start + len("USER_MSG_2") | |
| prefix = formatted[second_user_end:second_start] | |
| # Find where the assistant role marker starts | |
| assistant_pos = prefix.rfind("assistant") | |
| if assistant_pos != -1: | |
| # Search for a tag start ('<' or '[') before 'assistant' | |
| role_start = -1 | |
| for char in ["<", "["]: | |
| pos = prefix.rfind(char, 0, assistant_pos) | |
| role_start = max(role_start, pos) | |
| if role_start != -1: | |
| role_marker = prefix[role_start:] | |
| else: | |
| role_marker = prefix[assistant_pos:] | |
| else: | |
| role_marker = prefix | |
| # Strip <think>...</think> blocks from the role marker. Thinking model | |
| # templates wrap assistant content in these tags, but the test messages | |
| # can produce empty blocks (e.g. "<think>\n\n</think>\n") with reasoning models, | |
| # which then get baked into the regex as literals. Removing them ensures | |
| # that reasoning stays within the assistant content group. | |
| role_marker = re.sub(r"<think>.*?</think>\s*", "", role_marker, flags=re.DOTALL) | |
| # Determine the stable TURN-LEVEL suffix | |
| suffix1 = formatted[first_end:second_user_start] |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/speculators/data_generation/preprocessing.py` around lines 259 - 294,
Validate the `USER_MSG_2` lookup before computing `second_user_end` or slicing
`prefix`; raise a `ValueError` when the marker is absent, alongside the existing
assistant-marker validation. Store the validated position in a reusable
`second_user_start` symbol and use it when computing `suffix1` instead of
performing a second unguarded `formatted.find("USER_MSG_2")` call.
Source: Path instructions
| max_length: int | None = None, | ||
| ) -> torch.Tensor: | ||
| """Create loss mask by finding assistant response spans in formatted text.""" | ||
| loss_mask = torch.zeros(len(offsets), dtype=torch.bool) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Return a torch.long mask so the loss_mask column dtype is uniform across paths.
This function returns a torch.bool tensor, and _preprocess_batch appends it unchanged at Line 618. The two other paths produce torch.long: the native assistant-mask path at Line 409 and _passthrough_pretokenized at Line 529. load_and_preprocess_dataset calls concatenate_datasets at Line 921 over one dataset per entry in train_data_paths. If one entry resolves to the regex path and another resolves to the native-mask path or the pre-tokenized path, the loss_mask column carries two different Arrow types and concatenation is not guaranteed to succeed.
🐛 Proposed fix to align the dtype
- loss_mask = torch.zeros(len(offsets), dtype=torch.bool)
+ loss_mask = torch.zeros(len(offsets), dtype=torch.long)Update tests/integration/datagen/test_preprocessing.py Line 437, which asserts mask.dtype == torch.bool. The mask comparisons mask == 1 and mask == 0 used at Lines 1155-1156 keep working.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| loss_mask = torch.zeros(len(offsets), dtype=torch.bool) | |
| loss_mask = torch.zeros(len(offsets), dtype=torch.long) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/speculators/data_generation/preprocessing.py` at line 339, Change the
regex preprocessing path’s loss_mask allocation in the relevant preprocessing
function from torch.bool to torch.long, matching the native assistant-mask and
_passthrough_pretokenized paths. Update the integration test assertion that
expects torch.bool to expect torch.long, while preserving the existing mask
value comparisons.
| encoded = cast("BatchFeature", encoded_any) | ||
|
|
||
| # Remove batch dimension | ||
| (input_ids,) = encoded["input_ids"] | ||
| (offsets,) = encoded["offset_mapping"] | ||
|
|
||
| # MM placeholder tokens are inserted separate from chat template | ||
| formatted_text = processor.decode(input_ids) | ||
| assert isinstance(formatted_text, str) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Locate multimodal model repos used in datagen tests and the offset/decode usage.
set -euo pipefail
fd -t f 'test_preprocessing.py|test_regex_patterns.py' tests | while IFS= read -r f; do
rg -n 'MM_MODEL_REPO|MODELS\s*=|gemma|Qwen3\.5' "$f"
done
# Show every place offset_mapping is produced or consumed.
rg -n -C 6 'offset_mapping|return_offsets_mapping|\.decode\(' src/speculators/data_generation/preprocessing.pyRepository: vllm-project/speculators
Length of output: 2346
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== preprocessing relevant sections =="
sed -n '380,490p' src/speculators/data_generation/preprocessing.py
echo
echo "== helper implementation =="
rg -n -C 12 'def _create_loss_mask_from_offsets|def .*assistant.*offset|assistant_pattern' src/speculators/data_generation/preprocessing.py
echo
echo "== imports and target models =="
sed -n '1,45p' tests/datagen/test_preprocessing.pyRepository: vllm-project/speculators
Length of output: 18198
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate test files =="
fd -t f 'test_.*preprocessing|.*regex.*pattern|_mm|multimodal' tests . | sed -n '1,120p'
echo
echo "== inspect test_preprocessing if present =="
if [ -f tests/datagen/test_preprocessing.py ]; then
sed -n '1,260p' tests/datagen/test_preprocessing.py
else
echo "tests/datagen/test_preprocessing.py not present"
fi
echo
echo "== searches for multilingual/model constants across tests =="
rg -n --hidden 'MM_MODEL_REPO|MODELS\s*=|unsloth/gemma|google/gemma|Qwen3\.5|Qwen/Qwen3\.5|return_offsets_mapping|offset_mapping|formatted_text|decode\(' .Repository: vllm-project/speculators
Length of output: 9642
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== regex pattern tests around model coverage =="
sed -n '1,120p' tests/integration/datagen/test_regex_patterns.py
echo
echo "== offset mapping assertions/tests around mm preprocessing =="
sed -n '1080,1195p' tests/integration/datagen/test_preprocessing.py
sed -n '500,545p' tests/integration/datagen/test_preprocessing.py
sed -n '180,205p' tests/integration/datagen/test_preprocessing.py
echo
echo "== inspect imports in preprocessing =="
sed -n '1,90p' src/speculators/data_generation/preprocessing.pyRepository: vllm-project/speculators
Length of output: 12879
Keep formatted_text and offset_mapping based on the same rendered string.
The multimodal flow builds offset_mapping from apply_chat_template(..., tokenize=True, ...) but then computes the loss mask against formatted_text = processor.decode(input_ids). If decode(input_ids) renders placeholders, special tokens, or whitespace differently from the template text, bias.finditer(..., formatted_text) and offset_mapping will target misaligned character spans. Use the tokenized template text for the mask as well, and add an assertion that processor.decode(encoded["input_ids"][0]) character-for-character matches the text used to compute offset_mapping.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/speculators/data_generation/preprocessing.py` around lines 443 - 451, The
multimodal preprocessing path currently derives loss-mask text from decoded
input IDs instead of the exact rendered text used for offset_mapping. Update the
flow around encoded, input_ids, and formatted_text to retain and use the
tokenized apply_chat_template text for mask matching, then assert that
processor.decode(encoded["input_ids"][0]) matches that text
character-for-character.
Source: Path instructions
| try: | ||
| # Using trust_remote_code=True for variety of templates | ||
| return load_processor(model_id, trust_remote_code=True) | ||
| except (TypeError, ValueError, KeyError, AttributeError, RuntimeError) as e: | ||
| pytest.skip(f"Failed to load processor for {model_id}: {e}") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Broaden the fixture's exception handling so unavailable models skip instead of error.
The except clause lists TypeError, ValueError, KeyError, AttributeError, and RuntimeError. AutoProcessor.from_pretrained raises different types when a repository is gated, missing, or unreachable: huggingface_hub raises GatedRepoError, RepositoryNotFoundError, and HfHubHTTPError, and Transformers raises OSError for a missing local or remote artifact. None of those subclass the listed types. Several entries in MODELS are gated on the Hub, so this fixture errors the test run instead of skipping it.
💚 Proposed fix for the fixture
`@pytest.fixture`(scope="module", params=MODELS)
def processor(request):
model_id = request.param
try:
# Using trust_remote_code=True for variety of templates
return load_processor(model_id, trust_remote_code=True)
- except (TypeError, ValueError, KeyError, AttributeError, RuntimeError) as e:
+ except Exception as e: # noqa: BLE001 - any load failure must skip, not fail
pytest.skip(f"Failed to load processor for {model_id}: {e}")As per path instructions: "Ensure PyTest tests are clear, comprehensive, and cover edge cases specific to speculative decoding."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| # Using trust_remote_code=True for variety of templates | |
| return load_processor(model_id, trust_remote_code=True) | |
| except (TypeError, ValueError, KeyError, AttributeError, RuntimeError) as e: | |
| pytest.skip(f"Failed to load processor for {model_id}: {e}") | |
| try: | |
| # Using trust_remote_code=True for variety of templates | |
| return load_processor(model_id, trust_remote_code=True) | |
| except Exception as e: # noqa: BLE001 - any load failure must skip, not fail | |
| pytest.skip(f"Failed to load processor for {model_id}: {e}") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/integration/datagen/test_regex_patterns.py` around lines 45 - 49,
Broaden the exception handling in the processor-loading fixture around
load_processor to catch unavailable-model failures, including huggingface_hub
GatedRepoError, RepositoryNotFoundError, and HfHubHTTPError, plus Transformers’
OSError, so gated, missing, or unreachable models call pytest.skip with the
existing context. Preserve handling of the currently listed exception types.
Source: Path instructions
|
|
||
| # --------------------------------------------------------------------------- | ||
| # 6. Every text-only shared-registry preset works in on-policy regeneration. | ||
| # 6. Every shared-registry preset works on-policy (off-policy parity). |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 8 'Every shared-registry preset|DATASET_CONFIGS|TEXT|sharegpt4v_coco|MULTIMODAL_DATASETS|_dataset_choice' \
tests/unit/scripts/test_response_regeneration.py \
scripts/response_regeneration/script.pyRepository: vllm-project/speculators
Length of output: 11025
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant test section and the regen entry list used by prepare-data.
sed -n '830,902p' tests/unit/scripts/test_response_regeneration.py
printf '\n--- regen script dataset choice ---\n'
sed -n '27,45p' scripts/response_regeneration/script.py
printf '\n--- datagen config outline/imports for DATASET_CONFIGS ---\n'
fd -a 'configs.py|_configs.py' . | sed 's#^\./##'
sed -n '1,220p' speculators/data_generation/configs.pyRepository: vllm-project/speculators
Length of output: 3823
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- locate data_generation configs ---'
fd -a configs.py . | sed 's#^\./##'
printf '%s\n' '--- DATASET_CONFIGS definitions ---'
rg -n -C 3 'DATASET_CONFIGS|_DATASETS|TEXT|MULTIMODAL|sharegpt4v_coco|NAME' . -g '*.py'
printf '%s\n' '--- module path check ---'
python3 -c 'import pathlib, sys; p=pathlib.Path("speculators/data_generation/configs.py"); print(p.exists(), p)'Repository: vllm-project/speculators
Length of output: 46950
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- configs.py lines 1-190 ---'
sed -n '1,190p' src/speculators/data_generation/configs.py
printf '%s\n' '--- unit tests around multi-parameter data generation tests ---'
rg -n -C 5 'parametrize|DATASET_CONFIGS|text only|text presets|on-policy|off-policy|prepare_data|response_regeneration' src tests/unit -g '*.py'Repository: vllm-project/speculators
Length of output: 50380
Correct the dataset coverage claim.
DATASET_CONFIGS includes sharegpt4v_coco, but on-policy regeneration excludes it via REGEN_DATASETS = [name for name in DATASET_CONFIGS if name not in MULTIMODAL_DATASETS]; scripts/response_regeneration/script.py also rejects it. Either rename the section title to cover only supported text presets or add the missing off-policy multimodal case with an explicit assertion.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/scripts/test_response_regeneration.py` at line 857, Correct the
section around “Every shared-registry preset works on-policy” so its claim
matches actual coverage: either narrow the title to supported text presets, or
include the excluded sharegpt4v_coco case in the test with an explicit assertion
while respecting REGEN_DATASETS and script.py rejection behavior.
Source: Path instructions
3 participants
Reverts #794 since all our e2e tests still rely on the old preprocessing format.