Skip to content

Unify stop-token resolution across text and VLM adapters - #269

Open
stikves wants to merge 13 commits into
apple:mainfrom
stikves:sukru/vlm-additional-stop-tokens
Open

stikves wants to merge 13 commits into
apple:mainfrom
stikves:sukru/vlm-additional-stop-tokens

Conversation

@stikves

@stikves stikves commented Sep 20, 2026 •

Copy link
Copy Markdown
Contributor

Unifies stop-token resolution across the text (CoreAILanguageModel) and VLM (CoreAIVisionLanguageModel) adapters, the server (LLMServerMain), and the CLI (LLMRunnerMain), and makes stop/EOS collections Set<Int32> end-to-end, sorting to Array only at IO/ordering boundaries.

Resolution splits into two functions:

  • LanguageConfig.additionalStopTokenIds(from:tokenizer:) -> Set<Int32> scans tokenizer_config.json / tokenizer.json for turn-end special tokens. Called only when a bundle ships a tokenizer directory.
  • Tokenizer.runtimeStopTokens(additional:) -> Set<Int32> (Tokenizer+StopTokens.swift) builds the runtime terminating set: {eos-if-any} ∪ {base-vocab <|im_end| if present} ∪ additional. Both adapters call it in their generation loop.

The base-vocab <|im_end|> fold lives in runtimeStopTokens, so it applies whether or not the bundle ships a tokenizer directory (a bundle using the HF-runtime tokenizer fallback previously lost it).

The agentic <|eot|> fold is text-adapter-specific and is folded inline in CoreAILanguageModel.init.

Types: CoreAILanguageModel.additionalEosTokenIds, CoreAIVLMExecutor.Configuration.additionalStopTokenIds, StopSequences.init(additionalEosTokenIds:), ServerConfig.additionalEosTokenIds, and the CLI locals/helpers are Set<Int32>. StopSequences.sequences: [[Int32]] and the xgrammar-facing stopTokenIds: [Int32]? stay arrays (positional / C-facing). Additional IDs are sorted to Array when folded into StopSequences.sequences and in the CLI stop-token log line.

Per-site agentic-EOT behavior is unchanged: the server folds an eot explicitly, the CLI does not.

Tests: RuntimeStopTokensTests covers Tokenizer.runtimeStopTokens, including the base-vocab <|im_end|> fold. AdditionalStopTokensTests covers the config scan and asserts the resolver does not fold base-vocab <|im_end|>.

The Foundation Models VLM adapter stopped generation only on the main EOS
and a hardcoded <|im_end|> lookup, so a VLM whose turn terminator is
anything else (Gemma's <end_of_turn>, Phi's <|end|>) never stopped and ran
to the response-token cap. The text server already resolves these via
LanguageConfig.additionalStopTokenIds; this reads the same set once at load
and unions it into the stop set. Models that use <|im_end|> are unaffected.
Extract the stop-token union in CoreAIVLMExecutor.respond() into a
testable stopTokenSet(tokenizer:additionalStopTokenIds:) helper, since
the executor itself needs a real CoreAISequentialVLMEngine (compiled
.aimodel assets) to construct and can't be exercised at unit level.
Tests prove additional stop tokens resolved at load are actually
unioned into the set the executor checks per generated token, not
just stored.
@stikves
stikves force-pushed the sukru/vlm-additional-stop-tokens branch from f7edea8 to 53dd712 Compare September 21, 2026 22:20
stikves and others added 2 commits September 21, 2026 16:16
Text path now also stops on a base-vocab <|im_end|> via the shared
additionalIds im_end check, matching the VLM adapter.
@stikves stikves changed the title Honor additional stop tokens in the VLM adapter Unify stop-token resolution across text and VLM adapters Sep 21, 2026
Fold <|im_end|> into the shared LanguageConfig.additionalStopTokenIds
resolver so it is universal, and route the server and CLI through it
(they previously hand-rolled resolution and missed the fold). Dissolve
the StopTokens enum: runtime union moves to Tokenizer.runtimeStopTokens,
the agentic <|eot|> fold moves into the text adapter. Sort to Array only
at IO/ordering boundaries (StopSequences fold, CLI log).
Move the base-vocab <|im_end|> fold out of additionalStopTokenIds (only
called when a bundle ships a tokenizer directory) and into
runtimeStopTokens, which every adapter runs regardless of tokenizer dir.
Bundles using the HF-runtime tokenizer fallback now stop on <|im_end|>
again.
Server's agentic-eot fold called convertTokenToId without checking
vocabContains first, unlike the text-adapter reference pattern, so an
unk-token fallback could be folded as a false stop. Add the guard.

The CLI never folded the agentic <|eot|> into its stop set at all.
Add the same detect-thinking-format + guarded fold used by the server
and text adapter. Requires making Tokenizer.vocabContains public since
llm-runner/llm-server are separate targets from CoreAILanguageModels.
@stikves
stikves marked this pull request as ready for review September 22, 2026 13:49
@stikves stikves self-assigned this Sep 22, 2026
@stikves
stikves requested review from Lewis300, carinapeng and kevchengcodes and removed request for Lewis300 and carinapeng September 22, 2026 19:37
Comment on lines 11 to 16
extension Tokenizer {
/// Whether `token` is a genuine entry in the vocabulary, not an unk-token fallback.
func vocabContains(_ token: String) -> Bool {
public func vocabContains(_ token: String) -> Bool {
guard let id = convertTokenToId(token) else { return false }
return convertIdToToken(id) == token
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should this be moved to a tokenizer file instead of tool call parser?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tokenizer is not ours... Thinking back... this should not be public. Will make it internal and not share.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants