From 387bce041cc1ba8aaa2e39aae8ef43981eed6d15 Mon Sep 17 00:00:00 2001 From: Konstantin Nikolaou Date: Tue, 25 Aug 2026 13:13:59 +0200 Subject: [PATCH] feat: carry the tools column through to the trainer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit columns_to_keep dropped it and MESSAGES_FEATURES never declared it, so no row could reach the template's tool branches. TRL reads the column as a JSON string, so Value("string") is enough — no nested tool schema in Features, and datasets fills null for sources that have none. The filter renders with tools too. A declaration emits the template's whole "# Tools" preamble, so rendering without it would measure a shorter row than the one max_length is applied to. --- src/post_training/methods/sft.py | 22 +++++++++-- tests/test_data_loader.py | 52 +++++++++++++++++++++++++ tests/test_row_filters.py | 66 +++++++++++++++++++++++++++++--- 3 files changed, 131 insertions(+), 9 deletions(-) diff --git a/src/post_training/methods/sft.py b/src/post_training/methods/sft.py index 9fe5e59..1debf2f 100644 --- a/src/post_training/methods/sft.py +++ b/src/post_training/methods/sft.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import logging from collections import Counter from functools import partial @@ -46,6 +47,7 @@ def _classify_row( messages: list[dict], tokenizer: PreTrainedTokenizerBase, max_length: int | None, + tools: str | list[dict] | None = None, ) -> int: """Say what ``max_length`` does to one row's supervised span. @@ -73,6 +75,12 @@ def _classify_row( """ rendered = tokenizer.apply_chat_template( messages, + # Rendered WITH the tools, because the trainer renders with them. A tool + # declaration adds the template's whole "# Tools" preamble — hundreds of + # tokens — so a filter that left it out would measure a different length + # than the one max_length is applied to, and its verdict would be wrong + # for exactly the rows that carry tools. + tools=json.loads(tools) if isinstance(tools, str) else tools, return_dict=True, return_assistant_tokens_mask=True, ) @@ -96,7 +104,13 @@ def _classify_row( "content": Value("string"), "role": Value("string"), } - ) + ), + # A JSON string, matching what TRL reads: it does + # `json.loads(tools) if isinstance(tools, str) else tools` + # (trl/trainer/sft_trainer.py). Keeping it a string avoids declaring an + # arbitrary nested tool schema in Features, and datasets fills it with + # null for a dataset that has no tools, so mixed sources concatenate. + "tools": Value("string"), } ) @@ -174,7 +188,9 @@ def _filter_sft_rows( # remove_columns keeps the map cache tiny: it holds the verdict alone, not a copy of the data. verdicts = ds.map( - lambda row: {"verdict": _classify_row(row["messages"], tokenizer, max_length)}, + lambda row: { + "verdict": _classify_row(row["messages"], tokenizer, max_length, tools=row.get("tools")) + }, num_proc=num_proc, remove_columns=ds.column_names, desc="computing assistant loss masks", @@ -339,7 +355,7 @@ def build_sft_trainer(config: PostTrainingConfig, run_dir: Path) -> SFTTrainer: max_length=mc.max_seq_length, truncated_span_action=mc.truncated_span_action, ), - columns_to_keep=["messages"], + columns_to_keep=["messages", "tools"], features=MESSAGES_FEATURES, ) diff --git a/tests/test_data_loader.py b/tests/test_data_loader.py index 3726a7f..8cabfc6 100644 --- a/tests/test_data_loader.py +++ b/tests/test_data_loader.py @@ -178,6 +178,58 @@ def test_native_dataset_preserves_extra_message_fields(monkeypatch): assert mixed[0]["messages"][1]["function_calls"] == 'lookup(query="x")' +def test_tools_column_survives_mixing_and_is_null_where_absent(monkeypatch): + """A tool-declaring dataset mixed with one that has no `tools` column. + + `select_columns` leaves the two with different schemas, and + `concatenate_datasets` reconciles that by filling null rather than by + dropping the column — so no normalisation step is needed. Pinned because the + alternative would be silent: losing the column would disable tool support + for the whole mix without an error. + """ + _patch_load_dataset( + monkeypatch, + { + "dataset-a": Dataset.from_dict( + { + "messages": [ + [ + {"role": "user", "content": "look it up"}, + {"role": "assistant", "content": "done"}, + ] + ], + "tools": ['[{"name":"lookup"}]'], + "unused": ["drop me"], + } + ), + "dataset-b": Dataset.from_dict( + { + "messages": [ + [ + {"role": "user", "content": "no tools here"}, + {"role": "assistant", "content": "fine"}, + ] + ] + } + ), + }, + ) + + mixed = loader.load_and_mix_datasets( + _config( + DatasetEntry(name="a", path="dataset-a"), + DatasetEntry(name="b", path="dataset-b"), + ), + columns_to_keep=["messages", "tools"], + ) + + assert set(mixed.column_names) == {"messages", "tools"} + tools = list(mixed["tools"]) + assert len(tools) == 2 + assert '[{"name":"lookup"}]' in tools + assert None in tools + + def test_native_plain_chat_dataset_skips_schema_cast(monkeypatch): _patch_load_dataset(monkeypatch, {"dataset-a": _dataset("a", 2)}) features = Features( diff --git a/tests/test_row_filters.py b/tests/test_row_filters.py index 284b156..b0c733b 100644 --- a/tests/test_row_filters.py +++ b/tests/test_row_filters.py @@ -38,6 +38,7 @@ from __future__ import annotations +import json import logging import pytest @@ -45,7 +46,7 @@ from post_training.chat_templates.registry import get_chat_template from post_training.methods.dpo import _filter_dpo_rows -from post_training.methods.sft import _filter_sft_rows +from post_training.methods.sft import _classify_row, _filter_sft_rows SFT_LOGGER = "post_training.methods.sft" @@ -73,28 +74,31 @@ def __init__(self, template_name: str) -> None: def chat_template(self) -> str: return get_chat_template(self.template_name) - def _render(self, conversation: list[dict]) -> tuple[str, list[tuple[int, int]]]: + def _render( + self, conversation: list[dict], tools: list[dict] | None = None + ) -> tuple[str, list[tuple[int, int]]]: utils = pytest.importorskip("transformers.utils.chat_template_utils") if self.template_name not in _COMPILED: _COMPILED[self.template_name] = utils._compile_jinja_template(self.chat_template) return utils._render_with_assistant_indices( - _COMPILED[self.template_name], conversation, None, None, False + _COMPILED[self.template_name], conversation, tools, None, False ) - def render(self, conversation: list[dict]) -> str: - rendered, _ = self._render(conversation) + def render(self, conversation: list[dict], tools: list[dict] | None = None) -> str: + rendered, _ = self._render(conversation, tools) return rendered def apply_chat_template( self, conversation: list[dict], *, + tools: list[dict] | None = None, return_dict: bool = False, return_assistant_tokens_mask: bool = False, truncation: bool = False, max_length: int | None = None, ) -> dict: - rendered, indices = self._render(conversation) + rendered, indices = self._render(conversation, tools) input_ids = [ord(character) for character in rendered] assistant_masks = [0] * len(rendered) @@ -516,6 +520,56 @@ def test_reports_each_distinct_role_sequence_once(tokenizer, caplog) -> None: assert sorted(body.splitlines()) == ["user", "user -> assistant -> user"] +# ── SFT: tool declarations ───────────────────────────────────────────── + +TOOL_SCHEMA = [{"type": "function", "function": {"name": "lookup", "description": "d"}}] + + +def test_a_tools_declaration_counts_toward_max_length(tokenizer) -> None: + """Why carrying the column and rendering it have to land together. + + A tool declaration emits the template's whole "# Tools" preamble ahead of the + conversation, which pushes the assistant span later. A filter that rendered + without tools would measure a shorter row than the one max_length is applied + to, and would keep rows whose supervised span the trainer truncates away. + """ + row = _exchange(0) + bare = len(tokenizer.render(row)) + with_tools = len(tokenizer.render(row, TOOL_SCHEMA)) + assert with_tools > bare, "the tools preamble must lengthen the render" + + # A cap that fits the bare row but not the same row once tools are declared. + cap = bare + assert ( + len(_filter_sft_rows(_sft_dataset([row]), num_proc=1, tokenizer=tokenizer, max_length=cap)) + == 1 + ) + + ds = Dataset.from_dict({"messages": [row], "tools": [json.dumps(TOOL_SCHEMA)]}) + with pytest.raises(ValueError, match="zero loss"): + _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer, max_length=cap) + + +def test_tools_are_parsed_from_a_json_string(tokenizer) -> None: + """TRL reads the column with + `json.loads(tools) if isinstance(tools, str) else tools`, so the filter must + too. Left as a string, the template would iterate over its characters and + render something else entirely. + """ + row = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}] + as_string = _classify_row(row, tokenizer, None, tools=json.dumps(TOOL_SCHEMA)) + as_list = _classify_row(row, tokenizer, None, tools=TOOL_SCHEMA) + + assert as_string == as_list + + +def test_a_row_without_tools_is_unaffected(tokenizer) -> None: + """The column is null for most rows, and null must render exactly as before.""" + row = _exchange(0) + assert tokenizer.render(row, None) == tokenizer.render(row) + assert _classify_row(row, tokenizer, None, tools=None) == _classify_row(row, tokenizer, None) + + # ── DPO ────────────────────────────────────────────────────────────────