From 6b4fa873294bf5be726caa1e88d0a77d2fe373e2 Mon Sep 17 00:00:00 2001 From: Charlie Ruan <53290280+CharlieFRuan@users.noreply.github.com> Date: Sun, 22 Sep 2024 22:06:26 -0700 Subject: [PATCH] [Fix][Grammar] Detach token table to prevent disposing it (#571) Prior to this PR, when reusing the same engine but for different schema will run into error "Module has already been disposed". An example to reproduce this is included in https://github.com/mlc-ai/web-llm/issues/560. This is because `this.tokenTable` is a `tvmjs.TVMObject` and will be disposed after the scope ends. We fix this by wrapping the call with `this.tvm.detachFromCurrentScope()`, and only dispose `this.tokenTable` when we dispose the entire `LLMChatPipeline`. Co-authored-by: SMarioMan --- src/llm_chat.ts | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/llm_chat.ts b/src/llm_chat.ts index a9f5c086..2f1c0019 100644 --- a/src/llm_chat.ts +++ b/src/llm_chat.ts @@ -298,6 +298,7 @@ export class LLMChatPipeline { this.logitsOnCPU?.dispose(); this.tvm.dispose(); this.tokenizer.dispose(); + this.tokenTable?.dispose(); } /** @@ -557,9 +558,11 @@ export class LLMChatPipeline { if (this.tokenTable === undefined) { const rawTokenTable = getTokenTableFromTokenizer(this.tokenizer); // Post process entire table - this.tokenTable = this.fpostProcessTokenTable( - rawTokenTable, - this.token_postproc_method, + this.tokenTable = this.tvm.detachFromCurrentScope( + this.fpostProcessTokenTable( + rawTokenTable, + this.token_postproc_method, + ), ); } const grammar: BNFGrammar =