Feature/tool use support

docs/cli/prepare_data.md

@@ -34,6 +34,8 @@ python scripts/prepare_data.py \
 
   Example: `--data sharegpt --data ./custom_data.jsonl`
 
+  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).
+
 - **`--seq-length`** (int, default: `8192`) Maximum sequence length for each sample. Longer samples will be truncated.
 
 - **`--max-samples`** (int, default: `None`) Maximum number of samples to process. If `None`, processes all samples.

src/speculators/data_generation/preprocessing.py

@@ -1,4 +1,5 @@
 import bisect
+import json
 import random
 import re
 from collections.abc import Callable
@@ -388,6 +389,7 @@ def _get_input_ids_loss_mask(
     max_length: int,
     assistant_pattern: str | Pattern[str] | None,
     *,
+    tools: list[dict] | None = None,
     # For logging
     conv_idx: int | None = None,
 ):
@@ -398,6 +400,7 @@ def _get_input_ids_loss_mask(
         encoded_any = processor.apply_chat_template(
             hf_conv,
             tokenize=True,
+            tools=tools,  # type: ignore[arg-type]
             add_generation_prompt=False,
             return_assistant_tokens_mask=True,
             return_dict=True,
@@ -429,6 +432,7 @@ def _get_input_ids_loss_mask(
             encoded_any = processor.apply_chat_template(
                 hf_conv,
                 tokenize=True,
+                tools=tools,
                 add_generation_prompt=False,
                 return_dict=True,
                 processor_kwargs=processor_kwargs,
@@ -437,6 +441,7 @@ def _get_input_ids_loss_mask(
             encoded_any = processor.apply_chat_template(
                 hf_conv,
                 tokenize=True,
+                tools=tools,
                 add_generation_prompt=False,
                 return_dict=True,
                 **processor_kwargs,
@@ -456,6 +461,7 @@ def _get_input_ids_loss_mask(
         formatted_text = processor.apply_chat_template(
             hf_conv,
             tokenize=False,
+            tools=tools,  # type: ignore[arg-type]
             add_generation_prompt=False,
         )
         assert isinstance(formatted_text, str)
@@ -478,6 +484,29 @@ def _get_input_ids_loss_mask(
     return input_ids, loss_mask
 
 
+def _parse_conv_tools(conv_tools: object, idx: int) -> list | None:
+    """Parse the tools JSON string for one conversation; warn and return None
+    on invalid JSON or unexpected types."""
+    if not conv_tools:
+        return None
+    if isinstance(conv_tools, list):
+        return conv_tools
+    if not isinstance(conv_tools, str):
+        log.warning(
+            f"Non-string value in tools column for conversation {idx}: "
+            f"{type(conv_tools).__name__}, proceeding without tools"
+        )
+        return None
+    try:
+        return json.loads(conv_tools)
+    except json.JSONDecodeError as e:
+        log.warning(
+            f"Invalid JSON in tools column for conversation {idx}: {e}, "
+            "proceeding without tools"
+        )
+        return None
+
+
 def _preprocess_batch(
     examples: dict,
     processor: ProcessorLike,
@@ -499,7 +528,17 @@ def _preprocess_batch(
         log.warning(f"No conversations key found. Keys: {list(examples.keys())}")
         return results
 
+    tools_col = examples.get("tools")
+    if tools_col is not None and len(tools_col) != len(conversations):
+        log.warning(
+            f"Tools column length ({len(tools_col)}) does not match "
+            f"conversations length ({len(conversations)}), proceeding without tools"
+        )
+        tools_col = None
+
     for idx, conv in enumerate(conversations):
+        conv_tools = tools_col[idx] if tools_col is not None else None
+
         if not conv or not isinstance(conv, list):
             continue
 
@@ -508,12 +547,15 @@ def _preprocess_batch(
         if not normalized_conv:
             continue
 
+        parsed_tools = _parse_conv_tools(conv_tools, idx)
+
         try:
             input_ids, loss_mask = _get_input_ids_loss_mask(
                 normalized_conv,
                 processor,
                 max_length=max_length,
                 assistant_pattern=assistant_pattern,
+                tools=parsed_tools,
                 conv_idx=idx,
             )
         except (TypeError, ValueError, KeyError, AttributeError, RuntimeError) as e:

tests/integration/datagen/test_preprocessing.py

@@ -2,12 +2,15 @@
 Unit tests for the preprocessing module in the Speculators data generation.
 """
 
+import json
 import re
+from typing import Any
 
 import pytest
 import torch
 from datasets import Dataset as HFDataset
 from PIL import Image
+from transformers import AutoTokenizer
 
 from speculators.data_generation.preprocessing import (
     _adapt_conv_for_hf,
@@ -1176,3 +1179,217 @@ def test_build_eagle3_dataset_with_custom_pattern():
     # Should successfully build dataset with custom pattern
     assert isinstance(result, HFDataset)
     assert len(result) > 0
+
+
+# Tests for tool role and tool_calls / thinking field preservation
+
+
+@pytest.mark.sanity
+def test_normalize_conversation_with_tool_role():
+    """Test that 'tool' role is normalized correctly and not skipped."""
+    conv: list[dict[str, Any]] = [
+        {"role": "user", "content": "Call the weather API"},
+        {"role": "assistant", "content": "Sure, calling now."},
+        {"role": "tool", "content": '{"temperature": 20}'},
+        {"role": "assistant", "content": "It is 20 degrees."},
+    ]
+    result = _normalize_conversation(conv)
+
+    roles = [t["role"] for t in result]
+    assert "tool" in roles, "tool role should be preserved"
+    tool_turn = next(t for t in result if t["role"] == "tool")
+    assert tool_turn["content"] == '{"temperature": 20}'
+
+
+@pytest.mark.sanity
+def test_normalize_conversation_preserves_tool_calls_field():
+    """Test that 'tool_calls' field is preserved when present."""
+    tool_calls = [
+        {"id": "call_1", "type": "function", "function": {"name": "get_weather"}}
+    ]
+    conv: list[dict[str, Any]] = [
+        {"role": "user", "content": "What's the weather?"},
+        {"role": "assistant", "content": "", "tool_calls": tool_calls},
+    ]
+    result = _normalize_conversation(conv)
+
+    assert len(result) == 2
+    assistant_turn = result[1]
+    assert "tool_calls" in assistant_turn
+    assert assistant_turn["tool_calls"] == tool_calls
+
+
+@pytest.mark.sanity
+def test_preprocess_batch_with_tools():
+    """Test that tools from the dataset are forwarded to apply_chat_template.
+
+    tools must be a list of JSON strings (one per conversation in the batch),
+    matching the HuggingFace datasets batched-column convention.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(TEXT_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
+
+    example_tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_weather",
+                "description": "Get the current weather",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"location": {"type": "string"}},
+                    "required": ["location"],
+                },
+            },
+        }
+    ]
+
+    conv = [
+        {"role": "user", "content": "What's the weather in Prague?"},
+        {"role": "assistant", "content": "Let me check."},
+    ]
+    # tools must be a list — one JSON string per conversation
+    examples_with_tools = {
+        "conversations": [conv],
+        "tools": [json.dumps(example_tools)],
+    }
+    examples_without_tools = {
+        "conversations": [conv],
+    }
+
+    assistant_pattern = _detect_assistant_pattern(tokenizer)
+    results_with = _preprocess_batch(
+        examples_with_tools,
+        tokenizer,
+        max_length=512,
+        assistant_pattern=assistant_pattern,
+    )
+    results_without = _preprocess_batch(
+        examples_without_tools,
+        tokenizer,
+        max_length=512,
+        assistant_pattern=assistant_pattern,
+    )
+
+    assert "input_ids" in results_with
+    assert "loss_mask" in results_with
+    assert len(results_with["input_ids"]) == 1
+
+    # When the template renders tool definitions the sequence must be strictly
+    # longer than without tools.  Skip the length check if the template silently
+    # ignores the tools kwarg (tool name absent from decoded output).
+    decoded_with = tokenizer.decode(results_with["input_ids"][0])
+    if "get_weather" in decoded_with:
+        assert len(results_with["input_ids"][0]) > len(
+            results_without["input_ids"][0]
+        ), "Token sequence should be longer when tool definitions are included"
+
+
+@pytest.mark.sanity
+def test_preprocess_batch_with_invalid_tools_json():
+    """Test that invalid JSON in the tools column is handled gracefully.
+
+    The pipeline should warn and continue without tools rather than raising.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(TEXT_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!"},
+            ]
+        ],
+        "tools": ["this is not valid json"],
+    }
+
+    assistant_pattern = _detect_assistant_pattern(tokenizer)
+    # Must not raise; the bad JSON entry is skipped with a warning
+    results = _preprocess_batch(
+        examples, tokenizer, max_length=512, assistant_pattern=assistant_pattern
+    )
+
+    assert "input_ids" in results
+    assert len(results["input_ids"]) == 1
+
+
+@pytest.mark.sanity
+def test_preprocess_batch_tools_with_hf_assistant_mask():
+    """Test that tools are forwarded when using the HF assistant token mask path."""
+    tokenizer = AutoTokenizer.from_pretrained(TEXT_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 not _supports_assistant_mask(tokenizer):
+        pytest.skip("Tokenizer does not support HF assistant token mask")
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    example_tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_weather",
+                "description": "Get the current weather",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"location": {"type": "string"}},
+                    "required": ["location"],
+                },
+            },
+        }
+    ]
+
+    examples = {
+        "conversations": [
+            [
+                {"role": "user", "content": "What's the weather in Prague?"},
+                {"role": "assistant", "content": "Let me check."},
+            ]
+        ],
+        "tools": [json.dumps(example_tools)],
+    }
+
+    # assistant_pattern=None selects the HF mask path
+    results = _preprocess_batch(
+        examples, tokenizer, max_length=512, assistant_pattern=None
+    )
+
+    assert "input_ids" in results
+    assert "loss_mask" in results
+    assert len(results["input_ids"]) == 1
+    assert torch.any(results["loss_mask"][0] == 1)
+
+
+@pytest.mark.sanity
+def test_normalize_conversation_tool_calls_with_empty_content():
+    """Test that an assistant turn with tool_calls and no text content is normalized."""
+    tool_calls = [
+        {"id": "call_1", "type": "function", "function": {"name": "get_weather"}}
+    ]
+    conv: list[dict[str, Any]] = [
+        {"role": "user", "content": "What's the weather?"},
+        {"role": "assistant", "content": "", "tool_calls": tool_calls},
+        {"role": "tool", "content": '{"temperature": 22}'},
+        {"role": "assistant", "content": "It is 22 degrees outside."},
+    ]
+    result = _normalize_conversation(conv)
+
+    assert len(result) == 4
+    assistant_tool_call_turn = result[1]
+    assert assistant_tool_call_turn["role"] == "assistant"
+    assert assistant_tool_call_turn["content"] == ""
+    assert assistant_tool_call_turn["tool_calls"] == tool_calls