Chat-template repair: warn-by-default, AST classification, dict support

[danielhanchen]

Staging mirror of unslothai#5049

Original PR: unslothai#5049
Author: danielhanchen

This is a staging copy for review and editing. Once finalized, changes will be pushed back to the original PR.

---

Original description

Follow-up hardening on top of PR unslothai#4426 (which fixed the unslothai#4150 RuntimeError for ChatML LoRA reloads). Addresses the six follow-ups identified during the unslothai#4426 investigation.

Behavior changes

1. Warn-by-default instead of RuntimeError

Before: if the tokenizer's chat_template ignored add_generation_prompt and _fix_chat_template could not repair it, fix_chat_template raised a hard RuntimeError that blocked model loading entirely. In the unslothai#4150 scenario this meant a tokenizer saved by a downstream tool (LlamaFactory, Axolotl) could not be reloaded at all.

After: repair failure emits a logger.warning_once and returns the original template. Model loading continues; the user can either manually set tokenizer.chat_template or supply their own prompt for inference. UNSLOTH_STRICT_CHAT_TEMPLATE=1 restores the pre-warn hard fail for users who want it.

2. Local path vs HF hub distinguished in the message

Before: Please file a bug report to the maintainers of <name_or_path>. In the unslothai#4150 case <name_or_path> was the user's own saves/Llama-3.1-8B-Instruct/lora/train_.../ folder -- the user could not file a bug against their own directory.

After: messages are branched on os.path.isdir(name_or_path). Local paths get a message pointing at the likely downstream tool that re-serialized the tokenizer; HF hub IDs get a message pointing at the upstream model maintainers.

3. Dict / list chat_template now handled

Before: passing a tokenizer whose chat_template is a dict (e.g. Hermes-3 {default, tool_use}) or list to _fix_chat_template crashed with AttributeError: 'dict' object has no attribute 'find'. Surfaced during unslothai#4426 Phase 3 testing.

After: fix_chat_template detects dict / list shape and fixes each variant independently, preserving the outer structure.

Internals

• _find_end_position now matches all four Jinja whitespace-control variants ({% %}, {%- %}, {% -%}, {%- -%}) and returns the rightmost endfor/endif. Previously {%- endfor -%} (both-side dash, used by Qwen3-Guard) was silently bypassed because the old str.find was only looking for two specific dash styles.

• _has_add_generation_prompt_block uses Jinja AST (jinja2.nodes.If / Name walk) instead of substring matching. Templates that mention add_generation_prompt only inside a comment are correctly classified as "no block"; templates that use whitespace-control variants on the if/endif tags are correctly classified as "has block".

• _template_ends_with_toplevel_for gates the GH#4150 ChatML repair on the AST: only fires when the last structural top-level node is a For loop (the standard ChatML shape), ignoring trailing pure-whitespace Output nodes. Templates wrapped in an outer If (Qwen3-Guard) are now explicitly skipped at the _fix_chat_template level, not just at load_correct_tokenizer's name-based exemption. This means direct callers of _fix_chat_template also get correct behavior.

• _validate_patched_template

[danielhanchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the chat template repair logic by introducing AST-based detection of generation prompt blocks and a rendering-based method for deriving assistant prefixes. It also adds support for multi-variant templates and enhances error handling. Reviewers suggested performance improvements for Jinja environment reuse and AST traversal, along with a readability refactor for format detection.

[gemini-code-assist]

For a minor performance improvement, you could create a single jinja2.Environment() instance at the module level and reuse it across _template_ends_with_toplevel_for, _has_add_generation_prompt_block, _derive_assistant_prefix_by_render, and _VariantTokenizerProxy. This avoids repeated object creation.

For example, you could define _JINJA_ENV = jinja2.Environment() at the module level and then change this line to:

ast = _JINJA_ENV.parse(chat_template)

[gemini-code-assist]

Using any with find_all traverses the entire subtree, which is slightly inefficient. For better performance, you can use node.find(...) which stops after the first match. The isinstance check is also redundant since find_all already filters by type.

        if node.find(jinja2.nodes.Output) is not None:
            return True

[gemini-code-assist]

This nested try-except block for detecting the chat format can be refactored for better readability. Consider using a small helper function to probe for each format. For example:

def _check_format(messages):
    try:
        tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False)
        return True
    except Exception:
        return False

is_sharegpt = None
if _check_format([{"role": "user", "content": "Who are you?"}]):
    is_sharegpt = False
elif _check_format([{"from": "human", "value": "Who are you?"}]):
    is_sharegpt = True

[danielhanchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request significantly enhances the chat template repair logic in unsloth/tokenizer_utils.py. It introduces AST-based template analysis to ensure structural correctness, a robust 'render-diff' approach for automatically inferring assistant prefixes, and support for multi-variant templates (dictionaries or lists). Feedback includes a recommendation to simplify a redundant node type check in _if_body_emits_content and a suggestion to update the keyword argument handling in the _VariantTokenizerProxy fallback to use conversation for better compatibility with standard API expectations.

[gemini-code-assist]

The isinstance(d, jinja2.nodes.Output) check is redundant because node.find_all(jinja2.nodes.Output) already returns an iterator of jinja2.nodes.Output nodes. The any() call can be made directly on the result of find_all().

        if any(node.find_all(jinja2.nodes.Output)):
            return True

[gemini-code-assist]

The keyword argument for apply_chat_template is conversation, not messages. While the current usage within this file is positional, using the correct keyword name makes this fallback logic robust for other potential uses.

Suggested change

	messages = args[0] if args else kwargs.get("messages", [])
	add_generation_prompt = kwargs.get("add_generation_prompt", False)
	return env.from_string(self._template).render(
	messages = messages,
	add_generation_prompt = add_generation_prompt,
	)
	messages = args[0] if args else kwargs.get("conversation", [])
	add_generation_prompt = kwargs.get("add_generation_prompt", False)
	return env.from_string(self._template).render(
	messages = messages,
	add_generation_prompt = add_generation_prompt,
	)

[danielhanchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request significantly enhances the chat template repair logic in unsloth/tokenizer_utils.py by introducing AST-based structural analysis and a render-diff heuristic to automatically detect assistant prefixes. It also adds support for multi-variant templates via a new proxy class and improves user-facing warning messages. Feedback suggests making the fallback rendering logic more robust by supporting the conversation keyword and passing extra template variables, as well as using more flexible string matching when wrapping trailing expressions. Additionally, several staging artifacts and metadata files were identified for removal.

[gemini-code-assist]

The fallback rendering logic in _VariantTokenizerProxy.apply_chat_template has two issues:

1. It uses kwargs.get("messages"), but the standard keyword argument in the Transformers apply_chat_template API is conversation. If a user calls it with conversation=msgs, the fallback will default to an empty list.
2. It only passes messages and add_generation_prompt to the Jinja render method. Any extra keyword arguments passed to apply_chat_template (which are intended for the template, such as date_string in Llama 3.1) will be lost.

While this proxy is primarily used for internal repair logic with simple messages, making it more robust ensures that validation and potential future uses work correctly for all templates.

Suggested change

	messages = args[0] if args else kwargs.get("messages", [])
	add_generation_prompt = kwargs.get("add_generation_prompt", False)
	return env.from_string(self._template).render(
	messages = messages,
	add_generation_prompt = add_generation_prompt,
	)
	messages = args[0] if args else kwargs.get("conversation", kwargs.get("messages", []))
	add_generation_prompt = kwargs.get("add_generation_prompt", False)
	render_kwargs = {**kwargs, "messages": messages, "add_generation_prompt": add_generation_prompt}
	for key in ("tokenize", "padding", "truncation", "max_length", "return_tensors", "return_dict", "return_assistant_tokens_mask", "conversation"):
	render_kwargs.pop(key, None)
	return env.from_string(self._template).render(**render_kwargs)

[gemini-code-assist]

The check for a trailing expression in Case 1 is very strict and will fail if there is any whitespace (like a newline) between the {% endfor %} tag and the {{ ... }} expression. Using .strip() or .lstrip()/.rstrip() would make this repair logic more robust for templates that are formatted with newlines.

Suggested change

	and after_endfor.startswith("{{")
	and after_endfor.endswith("}}")
	and after_endfor.lstrip().startswith("{{")
	and after_endfor.rstrip().endswith("}}")

[gemini-code-assist]

This file appears to be a staging artifact containing a diff of the changes. It should be removed from the repository before merging to keep the codebase clean.

[gemini-code-assist]

This file seems to be metadata from a staging or CI process and likely does not belong in the final pull request.

[danielhanchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors and enhances the chat template repair mechanism in unsloth/tokenizer_utils.py. Key improvements include AST-based template parsing to identify structural nodes, support for multi-variant tokenizers through a new proxy class, and a render-diff approach to automatically derive assistant prefixes. The update also introduces a strict mode for template validation and updates several subproject dependencies. Feedback suggests that the fallback Jinja rendering path should include common tokens like bos_token and eos_token to prevent potential errors in templates that depend on them.

[gemini-code-assist]

In the fallback path for read-only tokenizers, the Jinja rendering environment only receives messages and add_generation_prompt. Many real-world chat templates rely on other variables like bos_token, eos_token, or pad_token. Without these, the rendering will likely fail with an UndefinedError during the repair diagnostic. Consider passing these common tokens from self._base to the render call to improve the robustness of the fallback path.

            return env.from_string(self._template).render(
                messages = messages,
                add_generation_prompt = add_generation_prompt,
                bos_token = getattr(self._base, "bos_token", ""),
                eos_token = getattr(self._base, "eos_token", ""),
                pad_token = getattr(self._base, "pad_token", ""),
                unk_token = getattr(self._base, "unk_token", ""),
            )