Skip to content

Commit

Permalink
[Fix][Grammar] Detach token table to prevent disposing it (#571)
Browse files Browse the repository at this point in the history
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
#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 <[email protected]>
  • Loading branch information
CharlieFRuan and SMarioMan authored Sep 23, 2024
1 parent 2f6762d commit 6b4fa87
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions src/llm_chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -298,6 +298,7 @@ export class LLMChatPipeline {
this.logitsOnCPU?.dispose();
this.tvm.dispose();
this.tokenizer.dispose();
this.tokenTable?.dispose();
}

/**
Expand Down Expand Up @@ -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 =
Expand Down

0 comments on commit 6b4fa87

Please sign in to comment.