Feature/tool use support

[olisicky]

This PR adds tool-use support for data generations. Many applications uses tool-use and training speculators with such data might be beneficial for the model.

[Copilot]

Pull request overview

This PR adds comprehensive tool-use support to the data generation pipeline for training speculators. The implementation enables handling of tool-based conversations, which is becoming increasingly important for modern language model applications.

Changes:

• Added support for the 'tool' role in conversation normalization
• Implemented preservation of 'tool_calls' field in conversation turns
• Added tools parameter forwarding to tokenizer's apply_chat_template method

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
src/speculators/data_generation/preprocessing.py	Added tool role handling, tool_calls field preservation, tools column parsing and forwarding to apply_chat_template
tests/datagen/test_preprocessing.py	Added comprehensive test coverage for tool role, tool_calls preservation, tools parameter forwarding, and error handling

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[fynnsu]

Hey, thanks for the pr. This would be great to support!

Could you clarify what dataset you are using as the basis for this? Or provide more information on the format this will support (and why)?

It seems like you are expecting something along the lines of:

{
  "conversations": [
    [{"role": ..., "content": ...}, {"role": ..., "content": ...}, ...]
    [{"role": ..., "content": ...}, {"role": ..., "content": ...}, ...]
    ...
  ]
  "tools": [
    json.dump([{"type": "function", "function": {...}}]), 
    json.dump([{"type": "function", "function": {...}}]), 
    ... # same number as conversations
  ]
}

I don't have a ton of experience with tool calling datasets but the first one I found when googling was this one: https://huggingface.co/datasets/Nanbeige/ToolMind/blob/main/open_datasets/ToolACE-query.jsonl

which uses a different format (list of dicts containing conversations + tools together). I'm not saying we have to support this (random) dataset, but would like to understand what kinds of dataset this will support. And then I think we should add documentation indicating what's supported and/or automatic warnings if we detect unsupported tool dataset types (for common formats that we don't support).

[olisicky]

Hi @fynnsu.

Could you clarify what dataset you are using as the basis for this? Or provide more information on the format this will support (and why)?

It seems like you are expecting something along the lines of:

{
  "conversations": [
    [{"role": ..., "content": ...}, {"role": ..., "content": ...}, ...]
    [{"role": ..., "content": ...}, {"role": ..., "content": ...}, ...]
    ...
  ]
  "tools": [
    json.dump([{"type": "function", "function": {...}}]), 
    json.dump([{"type": "function", "function": {...}}]), 
    ... # same number as conversations
  ]
}

I don't have a ton of experience with tool calling datasets but the first one I found when googling was this one: https://huggingface.co/datasets/Nanbeige/ToolMind/blob/main/open_datasets/ToolACE-query.jsonl

We have mix of multiple datasets. However, the tools column if very often used so we adopted it too. As the tools are of various scheme, it can't be easily specified for the pyarrow (datasets) thus it is dumped by json.dump so you can actually have the tools in the datasets.

This format can be seen e.g., in llamafactory/reason-tool-use-demo-1500, interstellarninja/hermes_reasoning_tool_use (sorted tool use datasets in HF by most downloads).

I am not sure if it can be robust enough to support all the datasets as it depends on the authors. For example Nvidia's nice dataset add the tools into metadata field.

It is similar to the requirement of conversations you are using. Multiple datasets use messages and some pre-processing is required anyway.

which uses a different format (list of dicts containing conversations + tools together). I'm not saying we have to support this (random) dataset, but would like to understand what kinds of dataset this will support. And then I think we should add documentation indicating what's supported and/or automatic warnings if we detect unsupported tool dataset types (for common formats that we don't support).

Sure, the format requirements in the documentation sounds like a good idea. I can add it. Do you have any suggestion where it would fit the best?

Thank you for reply,

Ondřej

[fynnsu]

We have mix of multiple datasets. However, the tools column if very often used so we adopted it too. As the tools are of various scheme, it can't be easily specified for the pyarrow (datasets) thus it is dumped by json.dump so you can actually have the tools in the datasets.

This format can be seen e.g., in llamafactory/reason-tool-use-demo-1500, interstellarninja/hermes_reasoning_tool_use (sorted tool use datasets in HF by most downloads).

Thank you for the examples!

I am not sure if it can be robust enough to support all the datasets as it depends on the authors. For example Nvidia's nice dataset add the tools into metadata field.

It is similar to the requirement of conversations you are using. Multiple datasets use messages and some pre-processing is required anyway.

Sure, the format requirements in the documentation sounds like a good idea. I can add it. Do you have any suggestion where it would fit the best?

Yeah, it's definitely okay if this doesn't support all types of datasets, especially because as you noted our current handling for messages isn't super general either.

I think if you could add a note to the data generation section in scripts/README.md on tool calling that would be sufficient. Even something short along the lines of "Data gen supports tool calling datasets with separate tool columns as seen in llamafactory/reason-tool-use-demo-1500 and interstellarninja/hermes_reasoning_tool_use". That README also gets automatically included in our docs.

[olisicky]

Okey, nice. I added a note to the readme and suggestions from copilot here.

Thank you!

[fynnsu]

Looks like the quality checks are failing. Could you run this locally and fix the type issues?

uv pip install -e .[dev]
make quality

The link checks are also failing but that maybe be unrelated. I can look into it.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 66a49100-f516-4609-8b8c-e1c00a54f317

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

This PR adds tool-related field support to conversation preprocessing. It extends role normalization to handle tool roles and tool_calls, adds batch-level reading and validation of tools columns, parses tool data as JSON, and forwards tools to the tokenizer's chat template. Documentation and comprehensive tests are included.

Changes

Tool Support in Conversation Preprocessing

Layer / File(s)	Summary
Tool role and tool_calls normalization src/speculators/data_generation/preprocessing.py	Adds json import, extends _normalize_conversation to recognize and map role="tool" to standard "tool" role, and updates normalized turn construction to preserve optional tool_calls field from assistant turns.
Batch tools column handling and tokenizer integration src/speculators/data_generation/preprocessing.py	_preprocess_batch now reads optional tools column, validates length matches conversations, extracts per-conversation tool data, parses tool JSON strings with warnings on errors, and passes parsed tools to tokenizer's apply_chat_template via tools= parameter in both the HF assistant token-mask and regex-based fallback paths.
Tool preprocessing test coverage tests/datagen/test_preprocessing.py	Adds imports and new unit and integration tests covering _normalize_conversation preservation of tool roles and tool_calls, _preprocess_batch forwarding of tools to chat template with token sequence comparison, graceful error handling for invalid JSON, and loss_mask validation in the HF assistant token-mask path.
Documentation updates scripts/README.md	Clarifies that input conversations should be provided in the conversations column and notes support for tool-calling datasets with separate tool columns.

Sequence Diagram

sequenceDiagram
  participant Dataset
  participant PreprocessBatch
  participant Tokenizer
  participant Output

  Dataset->>PreprocessBatch: batch with conversations & tools columns
  PreprocessBatch->>PreprocessBatch: validate tools column length
  PreprocessBatch->>PreprocessBatch: extract & parse tools JSON per conversation
  PreprocessBatch->>Tokenizer: apply_chat_template(..., tools=parsed_tools)
  Tokenizer->>Tokenizer: render template with tool context
  Tokenizer->>Output: input_ids, loss_mask with tool-aware tokens

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested labels

enhancement, data-generation

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Feature/tool use support' accurately captures the main objective of the PR, which adds tool-use support for data generation as described in the changeset.
Description check	✅ Passed	The description is related to the changeset, explaining that the PR adds tool-use support for data generation and notes potential benefits for model training.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @olisicky.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[coderabbitai]

Actionable comments posted: 1

♻️ Duplicate comments (1)

src/speculators/data_generation/preprocessing.py (1)

292-307: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Fix indentation: else block is paired with wrong statement.

The else at line 303 is currently indented as part of the try/except block, which means it executes when json.loads() succeeds (no exception raised). However, the warning message says "Non-string value in tools column", which should execute when conv_tools is not a string.

This causes two problems:

1. The "Non-string value" warning fires when JSON parsing succeeds (incorrect)
2. No warning fires when conv_tools is actually a non-string value (missing warning)

The else should be paired with the if isinstance(conv_tools, str): check, not the try/except.

🐛 Proposed fix

         # Parse tools JSON string if present; warn and skip tools on invalid JSON
         parsed_tools = None
         if conv_tools:
             if isinstance(conv_tools, str):
                 try:
                     parsed_tools = json.loads(conv_tools)
                 except json.JSONDecodeError as e:
                     log.warning(
                         f"Invalid JSON in tools column for conversation {idx}: {e}, "
                         "proceeding without tools"
                     )
-                else:
-                    log.warning(
-                        f"Non-string value in tools column for conversation {idx}: "
-                        f"{type(conv_tools).__name__}, proceeding without tools"
-                    )
+            else:
+                log.warning(
+                    f"Non-string value in tools column for conversation {idx}: "
+                    f"{type(conv_tools).__name__}, proceeding without tools"
+                )

🤖 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 292 - 307, The
"Non-string value" warning is currently inside the try/except (paired with
json.loads) so it runs on successful parse; move/dedent that else so it pairs
with the isinstance(conv_tools, str) check instead. Concretely, inside the block
handling conv_tools, keep the try/except around json.loads(conv_tools) to set
parsed_tools and log a JSONDecodeError via log.warning, and add an else branch
paired with the isinstance(...) check (or an explicit elif/else) that logs the
"Non-string value in tools column" message when conv_tools is not a string;
reference parsed_tools, conv_tools, json.loads, and log.warning to locate and
fix the code.

🧹 Nitpick comments (1)

tests/datagen/test_preprocessing.py (1)

688-840: ⚡ Quick win

Add an explicit test for tools/conversations batch-alignment validation.

These tests cover valid tools JSON, invalid JSON, and HF-mask forwarding well, but they don’t assert behavior when the tools column length does not match conversations length (a key new branch in this PR). Please add one test for this mismatch path (raise/skip behavior as intended) to lock in correctness.

Suggested test addition

+@pytest.mark.sanity
+def test_preprocess_batch_tools_length_mismatch():
+    """Tools column must align 1:1 with conversations in batched preprocessing."""
+    tokenizer = AutoTokenizer.from_pretrained(TEST_MODEL_REPO, trust_remote_code=True)
+
+    if not hasattr(tokenizer, "apply_chat_template") or tokenizer.chat_template is None:
+        pytest.skip("Tokenizer does not support chat templates")
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    examples = {
+        "conversations": [
+            [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi"}],
+            [{"role": "user", "content": "Bye"}, {"role": "assistant", "content": "Cya"}],
+        ],
+        # Deliberate mismatch: only one tools entry for two conversations
+        "tools": [json.dumps([{"type": "function", "function": {"name": "get_weather"}}])],
+    }
+
+    assistant_pattern = _detect_assistant_pattern(tokenizer)
+    # Adjust assertion based on intended contract in preprocessing:
+    with pytest.raises(ValueError):
+        _preprocess_batch(
+            examples, tokenizer, max_length=512, assistant_pattern=assistant_pattern
+        )

As per coding guidelines, "Ensure PyTest tests are clear, comprehensive, and cover edge cases specific to speculative decoding ... Verify 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/datagen/test_preprocessing.py` around lines 688 - 840, Add a new pytest
that supplies an examples dict where "conversations" and "tools" lengths differ
and assert the new mismatch branch in _preprocess_batch is exercised: create a
test function (e.g., test_preprocess_batch_tools_conversations_mismatch) that
builds a single conversation but a tools list of different length, loads the
tokenizer as in the other tests, and use pytest.raises(ValueError) around the
call to _preprocess_batch to verify the function raises on mismatched batch
lengths (reference the _preprocess_batch symbol and follow the same tokenizer
setup/skip guards used in test_preprocess_batch_with_tools).

🤖 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 `@scripts/README.md`:
- Around line 32-33: Update the README to explicitly state that the data
generation preprocessing expects the separate "tools" column to contain
JSON-dumped strings per row (i.e., a JSON array/object encoded as a string), not
arbitrary list formats; mention the "conversations" column requirement too, give
a short example of the exact expected row (conversations: "<conversation text>",
tools: '["toolA", {"name":"toolB","args":{}}]') and add a one-line note on how
to convert common alternate layouts (e.g., Python: json.dumps(row["tools"]) or
CSV: ensure proper quoting) so users know to preprocess other formats before
running the script.

---

Duplicate comments:
In `@src/speculators/data_generation/preprocessing.py`:
- Around line 292-307: The "Non-string value" warning is currently inside the
try/except (paired with json.loads) so it runs on successful parse; move/dedent
that else so it pairs with the isinstance(conv_tools, str) check instead.
Concretely, inside the block handling conv_tools, keep the try/except around
json.loads(conv_tools) to set parsed_tools and log a JSONDecodeError via
log.warning, and add an else branch paired with the isinstance(...) check (or an
explicit elif/else) that logs the "Non-string value in tools column" message
when conv_tools is not a string; reference parsed_tools, conv_tools, json.loads,
and log.warning to locate and fix the code.

---

Nitpick comments:
In `@tests/datagen/test_preprocessing.py`:
- Around line 688-840: Add a new pytest that supplies an examples dict where
"conversations" and "tools" lengths differ and assert the new mismatch branch in
_preprocess_batch is exercised: create a test function (e.g.,
test_preprocess_batch_tools_conversations_mismatch) that builds a single
conversation but a tools list of different length, loads the tokenizer as in the
other tests, and use pytest.raises(ValueError) around the call to
_preprocess_batch to verify the function raises on mismatched batch lengths
(reference the _preprocess_batch symbol and follow the same tokenizer setup/skip
guards used in test_preprocess_batch_with_tools).

🪄 Autofix (Beta)

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

Run ID: c86880b4-cc52-4b58-9dee-485354de9571

📥 Commits

Reviewing files that changed from the base of the PR and between e8d9da1 and 2ef74d1.

📒 Files selected for processing (4)

• scripts/README.md
• src/speculators/data_generation/preprocessing.py
• tests/datagen/test_preprocessing.py
• tests/unit/convert/test_eagle3_converter.py

[olisicky]

Hi @fynnsu . I totally missed that two commits were not signed before and the branch is blocked. I fixed it and rebased, hopefully everything is fine.

Thanks,

Ondřej

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @olisicky.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to address the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/speculators/blob/main/CONTRIBUTING.md

[fynnsu]

@olisicky Hey we just landed a bigger change to preprocessing to support multi-modal. Would you be able to rebase one more time?

[olisicky]

Hi @fynnsu . Sure, just rebased.

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to address the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/speculators/blob/main/CONTRIBUTING.md

[fynnsu]

Looks good! Thanks for adding this

[fynnsu]

Okay, looks like you have a type error. This seems like it should pass and mypy just doesn't accept it for some reason, so you can add # type: ignore [arg-type] to the problem lines to just ignore the issue.

And then run make quality again just to check everything is working.

[olisicky]

My bad. I rebased but didn't synced new pinned versions of ruff and it missed checks which were indicated here. Sorry.

[olisicky]

Hi @fynnsu . Partial implementation of this PR was merged 2 weeks ago #561 but they didn't add the tools in the apply_chat_template so the model can't see the tool definitions.

[fynnsu]

Hi @olisicky, could you rebase/merge main? I think there may also be some duplicated logic in _normalize_conversation now.

Can you also add an early return to _parse_conv_tools to handle already parsed datasets where the tools column is already a list? Something like:

 if isinstance(conv_tools, list):
      return conv_tools

I think this should be good to merge then. Sorry for the confusion with the other tool calling pr.

[fynnsu]

Hi @olisicky, I just merged the pr. Sorry this took so long, but glad we were able to get it in.