refactor: reasoning markers that are text, not one token #113

Merged

rcsheets merged 2 commits from refactor/reasoning-marker into main 2026-10-04 05:11:30 +00:00

Conversation 1 Commits 2 Files changed 10 +613 -112

rcsheets commented 2026-10-04 04:24:37 +00:00

Owner

Copy link

The reasoning split and the thinking budgets assumed a chat template's closing marker is one added token in the vocabulary. That holds for Granite 4 and Nemotron 3, but not for OLMo 3 (allenai/Olmo-3-7B-Think), whose </think> is ordinary BPE text. Such a model could be neither split nor budgeted.

This change puts the marker behind a small interface owned by the tokenizer, and adds a text implementation beside the existing id one. Behavior for templates whose marker is an added token is unchanged.

Why text, and not a token sequence

A text marker has no fixed token sequence. OLMo 3's tokenizer spells it according to the surrounding text:

text	tokens
x\n</think>\n\nAnswer	x \n </ think >\n\n Answer
x\n</think>\nAnswer	x \n </ think >\n Answer
x</think>Answer	x </ think > Answer
done.</think>\n\n	done .</ think >\n\n
x </think> y	x </ think > y

Matching any one id sequence would miss most of the closes the model actually writes, so the marker is matched in the decoded text.

What changes

• tokenizer.ReasoningEnd is an interface with the three things the engine needs from a marker:

  • Find(id, text, prev): did the token just generated complete the marker, and where do reasoning and answer split?
  • Pending(text): how many trailing bytes a stream must hold back because they may yet become the marker.
  • Close(lead, trail): the tokens that write "\n</think>\n", with either newline optional.

  ReasoningTokenizer.ReasoningEnd() returns one of these, or nil for no reasoning convention. It replaces ReasoningEndToken() int32.

• TokenMarker (added token) matches only its id, so the split stays unforgeable from prose. Its Pending is always 0, and Close splices the id between newline ids.

• TextMarker (ordinary text) finds the marker in the decoded output, starting far enough back to catch a marker split across tokens. The marker may end inside the token that completed it; the rest of that token belongs to the answer. Pending is the longest tail that is a prefix of the marker. Close encodes the whole string, so a forced close is the tokenizer's own spelling (>\n as one token).

• withChatTemplate picks the marker kind. It uses TokenMarker when the template's marker is an added token and TextMarker otherwise. Before, a missing token disabled reasoning for that template.

• The engine has one path.

  • It calls Find after each decode, holds back Pending while streaming, and forces Close for a budget cut.
  • ReasonTokens is counted back over the tokens that cover the marker; for an added token that is the same count as before.
  • The newline round-trip check moves from the engine into the markers. The hard budget's reserve for the leading newline is measured once from Close.

• AGENTS.md and README describe both marker kinds.

Not covered: under a grammar, bytes that share a token with the end of a text marker were sampled unconstrained and are not fed to the grammar. For whitespace, which is what a model writes there, this is harmless.

Verification

• go build ./..., go vet ./..., go test ./... pass; gofmt is clean.
• Existing reasoning, budget, stop-string and constrained-decoding tests pass, with their fakes building markers through TokenMarker.
• New tokenizer tests (reasoning_test.go, on a toy tokenizer that spells </think> the way OLMo 3 does):
  • TestTextMarkerFind: matches across token boundaries and inside a token, but not a marker earlier than the newest token could have completed.
  • TestTextMarkerPending: the holdback rule.
  • TestMarkerClose: the merged >\n spelling for text, the id between newlines for a token, and no newline from a tokenizer that cannot round-trip one.
  • TestTokenMarkerFind: matches on the id only, never on text.
• TestHFReasoningEndText: a template marker absent from the added tokens yields a TextMarker on a real Granite tokenizer. TestHFReasoningEndToken and the Nemotron goldens confirm both shipped templates still get TokenMarker with their ids.
• New engine tests:
  • TestGenerateChatReasoningSplitText: a text marker spanning two tokens splits an identical token stream exactly, and the streamed deltas reassemble both regions.
  • TestGenerateChatMaxThinkingTokensHardText: a hard budget forces \n </ think >\n, the reasoning is the thought plus its newline, and reasoning_tokens stays within the budget.
  • With Pending stubbed to 0, the hard-budget test failed 28 of 30 runs on streamed reasoning. It depends on when the stream goroutine flushes, so it is not deterministic; TestTextMarkerPending pins the rule itself.

No GPU run: the change is in tokenization and step-loop bookkeeping, covered by the CPU tests above.

🤖 Generated with Claude Code

The reasoning split and the thinking budgets assumed a chat template's closing marker is one added token in the vocabulary. That holds for Granite 4 and Nemotron 3, but not for OLMo 3 (`allenai/Olmo-3-7B-Think`), whose `</think>` is ordinary BPE text. Such a model could be neither split nor budgeted. This change puts the marker behind a small interface owned by the tokenizer, and adds a text implementation beside the existing id one. Behavior for templates whose marker is an added token is unchanged. ## Why text, and not a token sequence A text marker has no fixed token sequence. OLMo 3's tokenizer spells it according to the surrounding text: | text | tokens | | --- | --- | | `x\n</think>\n\nAnswer` | `x` `\n` `</` `think` `>\n\n` `Answer` | | `x\n</think>\nAnswer` | `x` `\n` `</` `think` `>\n` `Answer` | | `x</think>Answer` | `x` `</` `think` `>` `Answer` | | `done.</think>\n\n` | `done` `.</` `think` `>\n\n` | | `x </think> y` | `x` ` </` `think` `>` ` y` | Matching any one id sequence would miss most of the closes the model actually writes, so the marker is matched in the decoded text. ## What changes - **`tokenizer.ReasoningEnd` is an interface** with the three things the engine needs from a marker: - `Find(id, text, prev)`: did the token just generated complete the marker, and where do reasoning and answer split? - `Pending(text)`: how many trailing bytes a stream must hold back because they may yet become the marker. - `Close(lead, trail)`: the tokens that write `"\n</think>\n"`, with either newline optional. `ReasoningTokenizer.ReasoningEnd()` returns one of these, or nil for no reasoning convention. It replaces `ReasoningEndToken() int32`. - **`TokenMarker`** (added token) matches only its id, so the split stays unforgeable from prose. Its `Pending` is always 0, and `Close` splices the id between newline ids. - **`TextMarker`** (ordinary text) finds the marker in the decoded output, starting far enough back to catch a marker split across tokens. The marker may end inside the token that completed it; the rest of that token belongs to the answer. `Pending` is the longest tail that is a prefix of the marker. `Close` encodes the whole string, so a forced close is the tokenizer's own spelling (`>\n` as one token). - **`withChatTemplate` picks the marker kind.** It uses `TokenMarker` when the template's marker is an added token and `TextMarker` otherwise. Before, a missing token disabled reasoning for that template. - **The engine has one path.** - It calls `Find` after each decode, holds back `Pending` while streaming, and forces `Close` for a budget cut. - `ReasonTokens` is counted back over the tokens that cover the marker; for an added token that is the same count as before. - The newline round-trip check moves from the engine into the markers. The hard budget's reserve for the leading newline is measured once from `Close`. - AGENTS.md and README describe both marker kinds. Not covered: under a grammar, bytes that share a token with the end of a text marker were sampled unconstrained and are not fed to the grammar. For whitespace, which is what a model writes there, this is harmless. ## Verification - `go build ./...`, `go vet ./...`, `go test ./...` pass; gofmt is clean. - Existing reasoning, budget, stop-string and constrained-decoding tests pass, with their fakes building markers through `TokenMarker`. - New tokenizer tests (`reasoning_test.go`, on a toy tokenizer that spells `</think>` the way OLMo 3 does): - `TestTextMarkerFind`: matches across token boundaries and inside a token, but not a marker earlier than the newest token could have completed. - `TestTextMarkerPending`: the holdback rule. - `TestMarkerClose`: the merged `>\n` spelling for text, the id between newlines for a token, and no newline from a tokenizer that cannot round-trip one. - `TestTokenMarkerFind`: matches on the id only, never on text. - `TestHFReasoningEndText`: a template marker absent from the added tokens yields a `TextMarker` on a real Granite tokenizer. `TestHFReasoningEndToken` and the Nemotron goldens confirm both shipped templates still get `TokenMarker` with their ids. - New engine tests: - `TestGenerateChatReasoningSplitText`: a text marker spanning two tokens splits an identical token stream exactly, and the streamed deltas reassemble both regions. - `TestGenerateChatMaxThinkingTokensHardText`: a hard budget forces `\n` `</` `think` `>\n`, the reasoning is the thought plus its newline, and `reasoning_tokens` stays within the budget. - With `Pending` stubbed to 0, the hard-budget test failed 28 of 30 runs on streamed reasoning. It depends on when the stream goroutine flushes, so it is not deterministic; `TestTextMarkerPending` pins the rule itself. No GPU run: the change is in tokenization and step-loop bookkeeping, covered by the CPU tests above. 🤖 Generated with [Claude Code](https://claude.com/claude-code)

rcsheets added 1 commit 2026-10-04 04:24:38 +00:00

refactor(engine): reasoning markers that are text, not one token ... 20ad827bd4

The reasoning split assumed the template's closing marker is a single
added token. OLMo 3 writes "</think>" as ordinary BPE tokens ("</",
"think", ">", with ">" merging into a following newline), so that model
could not be split or budgeted.

tokenizer.ReasoningTokenizer now returns a ReasoningEnd (the marker's
text, plus its id when it is an added token) in place of
ReasoningEndToken. The engine keeps the id path unchanged for added
tokens and, when the marker is plain text, finds it in the decoded
output, holds back a possible partial marker while streaming, and
forces the budget close by encoding "\n</think>\n" as one string.
Granite 4 and Nemotron 3 still resolve to the id path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

pr-reviewer-bot commented 2026-10-04 04:24:39 +00:00

Collaborator

Copy link

Automated review by pr-reviewer v0.54.0 | Safety Check | Nemotron 3 Nano | tracking id r-c1dde0-5131fa
This is an AI-generated review and may contain mistakes.

Status: ❌ Failed

---

Review failed. Tracking id r-c1dde0-5131fa — see logs for details.

Comment @pr-reviewer-bot retry to try again.

<!-- pr-reviewer:review --> *Automated review by [pr-reviewer](https://git.brooktrails.org/brooktrails/pr-reviewer) v0.54.0 | Safety Check | Nemotron 3 Nano | tracking id `r-c1dde0-5131fa`* *This is an AI-generated review and may contain mistakes.* **Status:** ❌ Failed --- Review failed. Tracking id `r-c1dde0-5131fa` — see logs for details. Comment `@pr-reviewer-bot retry` to try again.

rcsheets added 1 commit 2026-10-04 05:02:23 +00:00

refactor(tokenizer): reasoning marker behind Find/Pending/Close ... 002ce870bf

ReasoningEnd was a struct of the marker's text and an id that was -1 for
a text marker, and the engine branched on the id. The engine only needs
three things from a marker -- whether the token just generated completed
it, how much of the text's tail a stream must hold back, and the tokens
that write it -- so ReasoningEnd is now an interface with those methods,
implemented by TokenMarker (added token, matched by id) and TextMarker
(ordinary text, found in the decoded output). The engine has one path,
and the newline round-trip check moves with Close into the tokenizer.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

rcsheets changed title from refactor(engine): reasoning markers that are text, not one token to refactor: reasoning markers that are text, not one token 2026-10-04 05:02:42 +00:00

rcsheets merged commit 2354f683e2 into main 2026-10-04 05:11:30 +00:00

rcsheets deleted branch refactor/reasoning-marker 2026-10-04 05:11:30 +00:00

rcsheets referenced this pull request from a commit 2026-10-04 05:11:32 +00:00

Merge pull request 'refactor: reasoning markers that are text, not one token' (#113) from refactor/reasoning-marker into main

rcsheets referenced this pull request 2026-10-04 06:01:02 +00:00

feat(tokenizer): OLMo 3 Think chat template #118

Sign in to join this conversation.

Reviewers

No reviewers

Labels

Clear labels

No items

No labels

Milestone

Clear milestone

No items

No milestone

Projects

Clear projects

No items

No project

Assignees

Clear assignees

No assignees

2 participants

Notifications

Subscribe

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set

Reference

brooktrails/gllm!113

No description provided.