Enable async encoding for the GPTEncoder (continuedev#1946)

This will improve the performance of indexing when using a GPT model.
Also, improved the Async interface to also have an async decode method.

Note: I was only able to test this method with the test I added as I do
not have access to a GPT model.

Co-authored-by: Rob Leidle <[email protected]>

Author: spew

core/llm/asyncEncoder.ts

@@ -5,22 +5,23 @@ import * as path from "path";
 
 export interface AsyncEncoder {
   encode(text: string): Promise<number[]>;
-  decode(tokens: number[]): string;
+  decode(tokens: number[]): Promise<string>;
+  close(): Promise<void>;
 }
 
 export class LlamaAsyncEncoder implements AsyncEncoder {
   private workerPool: workerpool.Pool;
 
   constructor() {
-    this.workerPool = workerpool.pool(path.join(__dirname, "/llamaTokenizerWorkerPool.mjs"));
+    this.workerPool = workerpool.pool(workerCodeFilePath("llamaTokenizerWorkerPool.mjs"));
   }
 
   async encode(text: string): Promise<number[]> {
     return this.workerPool.exec("encode", [text]);
   }
 
-  decode(tokens: number[]): string {
-    return llamaTokenizer.decode(tokens);
+  async decode(tokens: number[]): Promise<string> {
+    return this.workerPool.exec("decode", [tokens]);
   }
 
   // TODO: this should be called somewhere before exit or potentially with a shutdown hook
@@ -31,17 +32,30 @@ export class LlamaAsyncEncoder implements AsyncEncoder {
 
 // this class does not yet do anything asynchronous
 export class GPTAsyncEncoder implements AsyncEncoder {
-  private tiktokenEncoding: Tiktoken;
+  private workerPool: workerpool.Pool;
 
   constructor() {
-    this.tiktokenEncoding = _encodingForModel("gpt-4");
+    this.workerPool = workerpool.pool(workerCodeFilePath("tiktokenWorkerPool.mjs"));
   }
 
   async encode(text: string): Promise<number[]> {
-    return this.tiktokenEncoding.encode(text, "all", []);
+    return this.workerPool.exec("encode", [text]);
+  }
+
+  async decode(tokens: number[]): Promise<string> {
+    return this.workerPool.exec("decode", [tokens]);
+  }
+
+  // TODO: this should be called somewhere before exit or potentially with a shutdown hook
+  public async close(): Promise<void> {
+    await this.workerPool.terminate();
   }
+}
 
-  decode(tokens: number[]): string {
-    return this.tiktokenEncoding.decode(tokens);
+function workerCodeFilePath(workerFileName: string): string {
+  if (process.env.NODE_ENV === "test") {
+    // `cross-env` seems to make it so __dirname is the root of the project and not the directory containing this file
+    return path.join(__dirname, "llm", workerFileName);
   }
+  return path.join(__dirname, workerFileName);
 }

core/llm/llamaTokenizerWorkerPool.mjs

@@ -5,6 +5,11 @@ function encode(segment) {
     return llamaTokenizer.encode(segment);
 }
 
+function decode(tokens) {
+    return llamaTokenizer.decode(tokens);
+}
+
 workerpool.worker({
+    decode,
     encode,
 });

core/llm/tiktokenWorkerPool.mjs

@@ -0,0 +1,17 @@
+import workerpool from "workerpool";
+import { encodingForModel as _encodingForModel } from "js-tiktoken";
+
+const tiktokenEncoding = _encodingForModel("gpt-4");
+
+function encode(text) {
+    return tiktokenEncoding.encode(text, "all", []);
+}
+
+function decode(tokens) {
+    return tiktokenEncoding.decode(tokens);
+}
+
+workerpool.worker({
+    decode,
+    encode,
+});

core/test/llm/asyncEncoder.test.ts

@@ -0,0 +1,41 @@
+import { GPTAsyncEncoder, LlamaAsyncEncoder } from "../../llm/asyncEncoder";
+
+describe("llama encoder", () => {
+    var tokenizer: LlamaAsyncEncoder;
+
+    beforeAll(() => {
+        tokenizer = new LlamaAsyncEncoder();
+    });
+
+    afterAll(() => {
+        tokenizer.close();
+    });
+
+    test("hello world", async () => {
+        const input = "the quick brown fox jumped over the lazy dog";
+        const output = await tokenizer.encode(input);
+        expect(output).toEqual([1, 278, 4996, 17354, 1701, 29916, 12500, 287, 975, 278, 17366, 11203]);
+        const decoded = await tokenizer.decode(output);
+        expect(decoded).toBe(input);
+    });
+});
+
+describe("tiktoken encoder", () => {
+    var tokenizer: GPTAsyncEncoder;
+
+    beforeAll(() => {
+        tokenizer = new GPTAsyncEncoder();
+    });
+
+    afterAll(() => {
+        tokenizer.close();
+    });
+
+    test("hello world", async () => {
+        const input = "the quick brown fox jumped over the lazy dog";
+        const output = await tokenizer.encode(input);
+        expect(output).toEqual([1820, 4062, 14198, 39935, 27096, 927, 279, 16053, 5679]);
+        const decoded = await tokenizer.decode(output);
+        expect(decoded).toBe(input);
+    });
+});

extensions/vscode/scripts/prepackage.js

@@ -245,6 +245,7 @@ const exe = os === "win32" ? ".exe" : "";
     "../../../core/vendor/tree-sitter.wasm",
     "../../../core/llm/llamaTokenizerWorkerPool.mjs",
     "../../../core/llm/llamaTokenizer.mjs",
+    "../../../core/llm/tiktokenWorkerPool.mjs",
   ];
   for (const f of filesToCopy) {
     fs.copyFileSync(path.join(__dirname, f), path.join(__dirname, "..", "out", path.basename(f)));