feat: carry the tools column through to the trainer

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,
         )
 

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(

tests/test_row_filters.py

@@ -38,14 +38,15 @@
 
 from __future__ import annotations
 
+import json
 import logging
 
 import pytest
 from datasets import Dataset
 
 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 ────────────────────────────────────────────────────────────────