Feat/worker thread

core/indexing/LanceDbIndex.ts

@@ -39,7 +39,7 @@ export class LanceDbIndex implements CodebaseIndex {
 
   constructor(
     embeddingsProvider: EmbeddingsProvider,
-    readFile: (filepath: string) => Promise<string>,
+    readFile: (filepath: string) => Promise<string>
   ) {
     this.embeddingsProvider = embeddingsProvider;
     this.readFile = readFile;
@@ -65,18 +65,18 @@ export class LanceDbIndex implements CodebaseIndex {
   }
 
   private async *computeChunks(
-    items: PathAndCacheKey[],
+    items: PathAndCacheKey[]
   ): AsyncGenerator<
     | [
         number,
         LanceDbRow,
         { startLine: number; endLine: number; contents: string },
-        string,
+        string
       ]
     | PathAndCacheKey
   > {
     const contents = await Promise.all(
-      items.map(({ path }) => this.readFile(path)),
+      items.map(({ path }) => this.readFile(path))
     );
 
     for (let i = 0; i < items.length; i++) {
@@ -88,7 +88,7 @@ export class LanceDbIndex implements CodebaseIndex {
         items[i].path,
         content,
         LanceDbIndex.MAX_CHUNK_SIZE,
-        items[i].cacheKey,
+        items[i].cacheKey
       )) {
         chunks.push(chunk);
       }
@@ -98,11 +98,9 @@ export class LanceDbIndex implements CodebaseIndex {
         continue;
       }
 
-      // Calculate embeddings
       const embeddings = await this.embeddingsProvider.embed(
-        chunks.map((c) => c.content),
+        chunks.map((c) => c.content)
       );
-
       // Create row format
       for (let j = 0; j < chunks.length; j++) {
         const progress = (i + j / chunks.length) / items.length;
@@ -134,8 +132,8 @@ export class LanceDbIndex implements CodebaseIndex {
     results: RefreshIndexResults,
     markComplete: (
       items: PathAndCacheKey[],
-      resultType: IndexResultType,
-    ) => void,
+      resultType: IndexResultType
+    ) => void
   ): AsyncGenerator<IndexingProgressUpdate> {
     const lancedb = await import("vectordb");
     const tableName = this.tableNameForTag(tag);
@@ -164,7 +162,7 @@ export class LanceDbIndex implements CodebaseIndex {
           JSON.stringify(row.vector),
           data.startLine,
           data.endLine,
-          data.contents,
+          data.contents
         );
 
         yield { progress, desc };
@@ -197,7 +195,7 @@ export class LanceDbIndex implements CodebaseIndex {
       const stmt = await sqlite.prepare(
         "SELECT * FROM lance_db_cache WHERE cacheKey = ? AND path = ?",
         cacheKey,
-        path,
+        path
       );
       const cachedItems = await stmt.all();
 
@@ -234,7 +232,7 @@ export class LanceDbIndex implements CodebaseIndex {
       await sqlite.run(
         "DELETE FROM lance_db_cache WHERE cacheKey = ? AND path = ?",
         cacheKey,
-        path,
+        path
       );
     }
 
@@ -247,7 +245,7 @@ export class LanceDbIndex implements CodebaseIndex {
     n: number,
     directory: string | undefined,
     vector: number[],
-    db: any, /// lancedb.Connection
+    db: any /// lancedb.Connection
   ): Promise<LanceDbRow[]> {
     const tableName = this.tableNameForTag(tag);
     const tableNames = await db.tableNames();
@@ -271,7 +269,7 @@ export class LanceDbIndex implements CodebaseIndex {
     tags: IndexTag[],
     text: string,
     n: number,
-    directory: string | undefined,
+    directory: string | undefined
   ): Promise<Chunk[]> {
     const lancedb = await import("vectordb");
     if (!lancedb.connect) {
@@ -294,7 +292,7 @@ export class LanceDbIndex implements CodebaseIndex {
     const data = await sqliteDb.all(
       `SELECT * FROM lance_db_cache WHERE uuid in (${allResults
         .map((r) => `'${r.uuid}'`)
-        .join(",")})`,
+        .join(",")})`
     );
 
     return data.map((d) => {

core/indexing/embeddings/TransformersJsEmbeddingsProvider.ts

@@ -6,6 +6,7 @@ import {
 
 import path from "path";
 import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider";
+import { Worker } from "worker_threads";
 
 env.allowLocalModels = true;
 env.allowRemoteModels = false;
@@ -39,34 +40,32 @@ class TransformersJsEmbeddingsProvider extends BaseEmbeddingsProvider {
     return "sentence-transformers/all-MiniLM-L6-v2";
   }
 
-  async embed(chunks: string[]) {
-    let extractor = await EmbeddingsPipeline.getInstance();
+  async getInstance() {
+    return await EmbeddingsPipeline.getInstance();
+  }
 
-    if (!extractor) {
-      throw new Error("TransformerJS embeddings pipeline is not initialized");
-    }
+  async embed(chunks: string[]) {
+    return new Promise<number[][]>((resolve, reject) => {
+      const desiredDirectory = path.join(__dirname, "worker.js");
+      const worker = new Worker(desiredDirectory);
+      worker.on("message", (transcode_data) => {
+        resolve(transcode_data);
+        worker.terminate();
+      });
 
-    if (chunks.length === 0) {
-      return [];
-    }
+      worker.on("error", (err) => {
+        console.error(err);
+        reject(err);
+      });
 
-    let outputs = [];
-    for (
-      let i = 0;
-      i < chunks.length;
-      i += TransformersJsEmbeddingsProvider.MaxGroupSize
-    ) {
-      let chunkGroup = chunks.slice(
-        i,
-        i + TransformersJsEmbeddingsProvider.MaxGroupSize,
-      );
-      let output = await extractor(chunkGroup, {
-        pooling: "mean",
-        normalize: true,
+      worker.on("exit", (code) => {
+        if (code !== 0) {
+          reject(new Error(`Encoding stopped with exit code [ ${code} ]`));
+        }
       });
-      outputs.push(...output.tolist());
-    }
-    return outputs;
+
+      worker.postMessage(chunks);
+    });
   }
 }
 

core/indexing/embeddings/TransformersJsWorkerThread.js

@@ -1,8 +1,5 @@
 import path from "path";
-import {
-  env,
-  pipeline,
-} from "../../vendor/node_modules/@xenova/transformers/types/transformers";
+import { env, pipeline } from "core/node_modules/@xenova/transformers";
 import TransformersJsEmbeddingsProvider from "./TransformersJsEmbeddingsProvider";
 const { parentPort } = require("worker_threads");
 
@@ -43,7 +40,7 @@ parentPort.on("message", async (chunks) => {
     ) {
       let chunkGroup = chunks.slice(
         i,
-        i + TransformersJsEmbeddingsProvider.MaxGroupSize,
+        i + TransformersJsEmbeddingsProvider.MaxGroupSize
       );
       let output = await extractor(chunkGroup, {
         pooling: "mean",

extensions/vscode/esbuild.js

@@ -3,9 +3,9 @@ const esbuild = require("esbuild");
 const flags = process.argv.slice(2);
 
 const esbuildConfig = {
-  entryPoints: ["src/extension.ts"],
+  entryPoints: ["src/extension.ts", "src/worker.ts"],
   bundle: true,
-  outfile: "out/extension.js",
+  outdir: "out",
   external: ["vscode", "esbuild", "jsdom"],
   format: "cjs",
   platform: "node",

extensions/vscode/src/worker.ts

@@ -0,0 +1,31 @@
+import TransformersJsEmbeddingsProvider from "core/indexing/embeddings/TransformersJsEmbeddingsProvider";
+const { parentPort } = require("worker_threads");
+
+parentPort.on("message", async (chunks: any[]) => {
+  try {
+    console.log("extractor inside worker thread", chunks);
+    const extractorOne = new TransformersJsEmbeddingsProvider();
+    const extractor = await extractorOne.getInstance();
+
+    if (!extractor) {
+      throw new Error("TransformerJS embeddings pipeline is not initialized");
+    }
+
+    let outputs = [];
+    for (let i = 0; i < chunks.length; i += 4) {
+      let chunkGroup = chunks.slice(i, i + 4);
+      let output = await extractor(chunkGroup, {
+        pooling: "mean",
+        normalize: true,
+      });
+      outputs.push(...output.tolist());
+    }
+
+    parentPort.postMessage(outputs);
+    parentPort.close();
+  } catch (error: unknown) {
+    if (error instanceof Error) {
+      parentPort.postMessage({ error: error?.message });
+    }
+  }
+});