Skip to content

Commit e98369f

Browse files
authored
[Grammar][Fix] Pass in stop tokens to xgrammar TokenizerInfo (#642)
Prior to this PR, using models such as SmolLM, which has `<|endoftext|>` as an unk token and `<|im_end|>` as a stop token, runs into issues with XGrammar. This is because XGrammar has a builtin set of stop tokens, which includes `<|endoftext|>` but not `<|im_end|>`. This results in, at the end of a structured generation, `<|endoftext|>` is forced to be generated (as it is the only stop token recognized), but since it is not an actual stop token, the generation of the model does not stop. This PR explicitly passes in the stop tokens (recognized from `mlc-chat-config.json`) to `createTokenizerInfo()` so we do not use the built-in set of stop tokens. In the case above, `<|im_end|>` will be the only stop token used by XGrammar, fixing the issue.
1 parent 082f04e commit e98369f

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

src/llm_chat.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -554,6 +554,7 @@ export class LLMChatPipeline {
554554
this.token_postproc_method,
555555
this.prepend_space_in_encode,
556556
this.fullVocabSize,
557+
this.stopTokens,
557558
);
558559
this.grammarCompiler =
559560
await xgr.GrammarCompiler.createGrammarCompiler(

0 commit comments

Comments
 (0)