feat: carry the tools column through to the trainer - #58
Merged
Merged
Conversation
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.
Neonkraft
approved these changes
Aug 27, 2026
| truncated_span_action=mc.truncated_span_action, | ||
| ), | ||
| columns_to_keep=["messages"], | ||
| columns_to_keep=["messages", "tools"], |
Collaborator
There was a problem hiding this comment.
Move this up to a constant COLUMNS_TO_KEEP, similar to MESSAGE_FEATURES, please :)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
columns_to_keep=["messages"]insrc/post_training/methods/sft.pydrops any row-leveltoolscolumn before the trainer sees it, andMESSAGES_FEATURESdeclares onlyrole/content, so a structuredmessage.tool_callsfield is removed by the schema cast as well. Theqwen3template's tool branches — the# Toolspreamble, thetool_callsloop, thetool-role<tool_response>block — are therefore unreachable through this pipeline: no row can carry a declaration to them. This PR carries the column.TRL reads it as a JSON string (
json.loads(tools) if isinstance(tools, str) else tools,trl/trainer/sft_trainer.py), sotoolsis declaredValue("string")rather than as a nested tool schema, anddatasetsfills it with null for a source that has none — so a tools-carrying dataset still concatenates with one that has none._classify_rownow renders withtools=too, parsed the way TRL parses it. That half is not cosmetic: the row filter measures againstmax_length, and a declaration emits the entire# Toolspreamble ahead of the conversation — 5,382 tokens on the worst row of the pool we're preparing. A filter rendering without it measures a shorter row than the one truncation is applied to, and keeps rows whose assistant span the trainer then cuts. Carrying the column without this would introduce that, which is why both are in one PR.Type of change
Validation
pytest tests/— 132 pass, 4 new. Ruff 0.9.10 and black 25.1.0 clean at the pinned pre-commit versions.Three of the four new tests fail against
main; the fourth pinsconcatenate_datasets' null-filling for a mixed schema, which holds either way and is there to catch a silent regression rather than to prove this change.The load-bearing one is
test_a_tools_declaration_counts_toward_max_length: the same row with the cap set to its bare rendered length is kept without tools and rejected with them.