Add tool-schema support to SFT tokenization

[hamishivi]

Replacement for #1734, rebased onto current main and pushed from hamishivi/open-instruct.

Summary

• Pass the dataset tools column to apply_chat_template during SFT tokenization so tool schemas are rendered into the prompt. A new _normalize_tools_for_chat_template helper accepts a list/dict/JSON-string (and treats empty/None as no tools), rejecting tool-name-string lists.
• Derive assistant labels from token offset mappings against the rendered conversation (prefix-stable per-assistant-turn spans) instead of mask_labels, so labels stay correct when tools are present.
• Consume the tools column during SFT tokenization rather than persisting it to the cached dataset (_SFT_TOKENIZE_FNS).
• Add sft_tulu_tokenize_without_truncation_v1 and register last_turn_tulu_tokenize_and_truncate_v1 (which now also forwards tools).
• Bump DATASET_CACHE_VERSION to v7 to invalidate stale caches.

Notes

• For the no-tools tulu path the new tokenizer is byte-identical to the previous mask_labels path: the existing GOLD_SFT golden-hash test passes unchanged.
• This PR intentionally excludes the unrelated dataset_config_name plumbing and the sft_filter_v1 rename that were bundled alongside this work on the source branch.

Test plan

• uv run pytest open_instruct/test_dataset_transformation.py (20 passed locally), including:
  • existing GOLD_SFT/preference/rlvr golden-hash tests (unchanged),
  • new TestToolNormalization (JSON parsing, dict-wrapping, empty handling, rejection cases),
  • new TestSFTTuluTokenizeLabels (assistant-only trainable labels, tools column consumed, no-truncation variant).

GPU_TESTS=bypass

Made with Cursor

[gemini-code-assist]

Code Review

This pull request introduces tool-schema support to SFT tokenization by parsing the tools column, passing it to apply_chat_template, and deriving assistant labels from offset mappings using fast tokenizers. It also ensures the tools column is consumed rather than persisted, and adds comprehensive unit tests. The reviewer feedback suggests two valuable improvements: stripping whitespace from the tools string before parsing to handle whitespace-only strings safely, and adding an explicit check to verify that add_generation_prompt actually appends the assistant header to prevent silent mis-masking when the template does not support it.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

When parsing the tools column from a string, it is safer to strip any leading or trailing whitespace first. This prevents potential parsing issues with whitespace-only strings (e.g., from CSV/TSV exports) and allows explicitly treating them as empty/None.

Suggested change

	if isinstance(tools, str):
	try:
	tools = json.loads(tools)
	except json.JSONDecodeError as exc:
	raise ValueError(f"{TOOLS_COLUMN_KEY} must be a JSON-encoded tool schema list, got: {tools!r}") from exc
	# Re-check after parsing: a JSON "null" or "" decodes to None / "".
	if tools is None or tools == "":
	return None
	if isinstance(tools, str):
	tools_str = tools.strip()
	if not tools_str or tools_str.lower() == "null":
	return None
	try:
	tools = json.loads(tools_str)
	except json.JSONDecodeError as exc:
	raise ValueError(f"{TOOLS_COLUMN_KEY} must be a JSON-encoded tool schema list, got: {tools!r}") from exc
	# Re-check after parsing: a JSON "null" or "" decodes to None / "".
	if tools is None or tools == "":
	return None

[gemini-code-assist]

If the chat template does not support add_generation_prompt (i.e., it silently ignores the argument), header (rendered with add_generation_prompt=True) will be identical to the prefix rendered with add_generation_prompt=False. In this case, the prefix-stability check won't fail, but the assistant header will be silently trained because the trainable span will start before the assistant content. Adding an explicit check to ensure add_generation_prompt actually appended the assistant header prevents this silent mis-masking.

Suggested change

	if message_idx == 0:
	header = ""
	else:
	header = tokenizer.apply_chat_template(
	conversation=messages[:message_idx], tools=tools, tokenize=False, add_generation_prompt=True
	)
	through = tokenizer.apply_chat_template(
	conversation=messages[: message_idx + 1], tools=tools, tokenize=False, add_generation_prompt=False
	)
	if message_idx == 0:
	header = ""
	else:
	header = tokenizer.apply_chat_template(
	conversation=messages[:message_idx], tools=tools, tokenize=False, add_generation_prompt=True
	)
	# Ensure add_generation_prompt actually appended the assistant header to avoid silent mis-masking
	header_without_prompt = tokenizer.apply_chat_template(
	conversation=messages[:message_idx], tools=tools, tokenize=False, add_generation_prompt=False
	)
	if header == header_without_prompt:
	raise ValueError(
	"Cannot compute assistant label spans: the chat template does not support "
	"add_generation_prompt or failed to append the assistant header."
	)
	through = tokenizer.apply_chat_template(
	conversation=messages[: message_idx + 1], tools=tools, tokenize=False, add_generation_prompt=False
	)