Fix chat template marker auto-detection for Zephyr's context-dependent tokenization - #899
Conversation
get_chat_template_parts validated candidate markers by tokenizing them
standalone and searching the token core in the rendered 3-turn probe.
Zephyr's role tags are not special tokens, and SentencePiece tokenizes
the leading < as a word-start piece standalone but as a bare piece when
the tag follows </s> plus a newline mid-conversation, so:
- response_part '<|assistant|>\n' never matched the probe and the
detector raised, forcing users back to manual markers that fail the
same way.
- instruction_part '<|user|>\n' validated through its single text-start
occurrence (the first user turn only), which at masking time misses
every later user turn and would train on user content after the first
assistant header.
Fix: probe a leading-newline candidate as well, and prefer candidates
whose token core matches at least twice in the probe (the probe holds
three turns per role, so a reliable marker matches from the second turn
on). Single-match acceptance is kept for the original candidates so
previously passing templates keep their exact detection.
Validated over the 35-template Studio catalog with real tokenizers:
Zephyr now auto-detects ('\n<|user|>\n', '\n<|assistant|>\n') with
system and user turns masked, all assistant turns trained and the final
EOS trained; every other template's detected parts, labels and verdict
are byte-identical to before the change.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request improves the auto-detection of role markers (such as Zephyr's <|user|> and <|assistant|>) whose tokenization depends on context. It updates the validation logic in unsloth_zoo/dataset_utils.py to probe a leading-newline candidate and prioritize candidates that match at least twice in the multi-turn probe, preventing issues with single-context matches. A comprehensive test suite has also been added in tests/test_zephyr_marker_context.py. The review feedback suggests a minor refactoring of the is_trained helper function in the new test file to avoid semicolons and use a generator expression for better readability and PEP 8 compliance.
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.
| def is_trained(sub): | ||
| i = text.index(sub); s, e = i, i + len(sub) | ||
| return all(k in un for k in [j for j, (a, b) in enumerate(offs) if b > a and a < e and b > s]) |
There was a problem hiding this comment.
We can simplify the helper function is_trained to be more readable and idiomatic:
- Avoid using a semicolon to separate statements on a single line, adhering to PEP 8 guidelines.
- Simplify the list comprehension inside
all(...)to a direct generator expression, which avoids creating an intermediate list in memory.
| def is_trained(sub): | |
| i = text.index(sub); s, e = i, i + len(sub) | |
| return all(k in un for k in [j for j, (a, b) in enumerate(offs) if b > a and a < e and b > s]) | |
| def is_trained(sub): | |
| i = text.index(sub) | |
| s, e = i, i + len(sub) | |
| return all(j in un for j, (a, b) in enumerate(offs) if b > a and a < e and b > s) |
References
- PEP 8 discourages multiple statements on the same line separated by semicolons. (link)
There was a problem hiding this comment.
Style-only suggestion with no behavior change; the helper is a five-line test-local closure and the current form matches the sibling tests in this directory. Keeping as is.
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. Nice work! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: eba55d740d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| from transformers import AutoTokenizer | ||
| from tokenizers import Tokenizer |
There was a problem hiding this comment.
When transformers or tokenizers is not installed, these imports raise before the surrounding try block, so _setup() never reaches the intended pytest.skip path. I verified python -m pytest tests/test_zephyr_marker_context.py -q in this environment fails all three new tests with ModuleNotFoundError: No module named 'transformers'; move these optional imports inside the try or catch ImportError so dependency-light test runs skip instead of failing.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in the follow-up commit: the transformers/tokenizers imports moved inside a try that returns None on ImportError, which routes into the existing pytest.skip. Verified both ways: 3 passed with deps present, 3 skipped with the imports blocked.
|
@codex review |
1 similar comment
|
@codex review |
The transformers/tokenizers imports ran before the try block, so dependency-light runs errored during setup instead of skipping.
|
@codex review |
|
Codex Review: Didn't find any major issues. 👍 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
What was broken
get_chat_template_partsfailed on Zephyr (unsloth/zephyr-sft,HuggingFaceH4/zephyr-7b-beta) with:Zephyr's role tags (
<|user|>,<|assistant|>) are ordinary text, not special tokens. SentencePiece tokenizes them context-dependently:▁< | ass istant | ></s>+ newline:▁ <0x0A> < | ass istant | >(bare<, no word-start piece)The validation step tokenized each candidate marker standalone (plus a leading-space variant) and searched that token core in the rendered 3-turn probe, so:
'<|assistant|>\n'never matched anywhere in the probe and the detector raised. The usual manual markers fail the same way at masking time (same standalone tokenization in_find_common_token_ids(force_match=True)), so Zephyr had no working path andtrain_on_responses_onlymasked every assistant token.'<|user|>\n'DID validate, but only via its single text-start occurrence (the very first user turn). Had detection gotten that far, masking would have missed every later user turn and trained on user content after the first assistant header.The fix
In the marker validation of
get_chat_template_parts:'\n' + part), which reproduces the mid-conversation tokenization (the newline carries the SentencePiece word-start piece, leaving the tag tokenized exactly as it appears after a turn boundary), andSingle-match acceptance is kept for the original two candidates, so any template that previously validated still validates, and templates whose first candidate matches repeatedly (all special-token templates) return exactly what they returned before.
Zephyr now auto-detects
('\n<|user|>\n', '\n<|assistant|>\n'), and those strings round-trip throughtrain_on_responses_only's exact-match tokenization correctly.Validation
Catalog sweep over the 35 Studio chat templates (real tokenizers, two-turn fixture with a system message, token-level label checks: user+system masked, all assistant turns trained, final EOS trained):
AUTO_FAILS-> auto works.user1_masked / user2_masked / system_masked / asst1_trained / asst2_trained / eos_trainedall pass; trained text is exactly'grape reply number one.</s>grape reply number two.</s> \n'.Qwen/QwQ-32Bre-checked separately): detected parts, labels, verdict and notes byte-identical before vs after the change.Tests: new
tests/test_zephyr_marker_context.pybuilds the verbatim Zephyr template on the cachedhf-internal-testing/llama-tokenizer(role tags NOT registered as special tokens, eos parsing matched to Zephyr's tokenizer) and asserts detection, token-level masking, both assistant EOS labels trained, and that special-token templates keep their previous detection. Both new regression tests fail on main and pass with the fix;test_special_token_template_unchangedpasses on both.