Merge pull request #22 from supabase-community/feat/embed-worker

feat: move embed to web worker

Author: gregnr

apps/postgres-new/lib/embed/index.ts

@@ -0,0 +1,31 @@
+import { FeatureExtractionPipelineOptions } from '@xenova/transformers'
+import * as Comlink from 'comlink'
+
+type EmbedFn = (typeof import('./worker.ts'))['embed']
+
+let embedFn: EmbedFn
+
+// Wrap embed function in WebWorker via comlink
+function getEmbedFn() {
+  if (embedFn) {
+    return embedFn
+  }
+
+  if (typeof window === 'undefined') {
+    throw new Error('Embed function only available in the browser')
+  }
+
+  const worker = new Worker(new URL('./worker.ts', import.meta.url), { type: 'module' })
+  embedFn = Comlink.wrap<EmbedFn>(worker)
+  return embedFn
+}
+
+/**
+ * Generates an embedding for each text in `texts`.
+ *
+ * @returns An array of vectors.
+ */
+export async function embed(texts: string[], options?: FeatureExtractionPipelineOptions) {
+  const embedFn = getEmbedFn()
+  return await embedFn(texts, options)
+}

apps/postgres-new/lib/embed/worker.ts

@@ -0,0 +1,17 @@
+import { FeatureExtractionPipelineOptions, pipeline } from '@xenova/transformers'
+import * as Comlink from 'comlink'
+
+const embedPromise = pipeline('feature-extraction', 'supabase/gte-small', {
+  quantized: true,
+})
+
+export async function embed(
+  texts: string[],
+  options?: FeatureExtractionPipelineOptions
+): Promise<number[][]> {
+  const embedFn = await embedPromise
+  const tensor = await embedFn(texts, options)
+  return tensor.tolist()
+}
+
+Comlink.expose(embed)

apps/postgres-new/lib/hooks.ts

@@ -1,6 +1,5 @@
 'use client'
 
-import { FeatureExtractionPipelineOptions, pipeline } from '@xenova/transformers'
 import { generateId } from 'ai'
 import { Chart } from 'chart.js'
 import { codeBlock } from 'common-tags'
@@ -18,6 +17,7 @@ import { createPortal } from 'react-dom'
 import { useDatabaseUpdateMutation } from '~/data/databases/database-update-mutation'
 import { useTablesQuery } from '~/data/tables/tables-query'
 import { getDb } from './db'
+import { embed } from './embed'
 import { loadFile, saveFile } from './files'
 import { SmoothScroller } from './smooth-scroller'
 import { OnToolCall } from './tools'
@@ -382,13 +382,11 @@ export function useOnToolCall(databaseId: string) {
           const { texts } = toolCall.args
 
           try {
-            const tensor = await embed(texts, {
+            const embeddings = await embed(texts, {
               normalize: true,
               pooling: 'mean',
             })
 
-            const embeddings: number[][] = tensor.tolist()
-
             const sql = codeBlock`
               insert into meta.embeddings
                 (content, embedding)
@@ -425,15 +423,6 @@ export function useOnToolCall(databaseId: string) {
   )
 }
 
-const embedPromise = pipeline('feature-extraction', 'supabase/gte-small', {
-  quantized: true,
-})
-
-export async function embed(texts: string | string[], options?: FeatureExtractionPipelineOptions) {
-  const embedFn = await embedPromise
-  return embedFn(texts, options)
-}
-
 export type UseDropZoneOptions = {
   onDrop?(files: File[]): void
   cursorElement?: ReactNode

apps/postgres-new/package.json

@@ -31,6 +31,7 @@
     "chartjs-adapter-date-fns": "^3.0.0",
     "class-variance-authority": "^0.7.0",
     "clsx": "^2.1.1",
+    "comlink": "^4.4.1",
     "common-tags": "^1.8.2",
     "date-fns": "^3.6.0",
     "framer-motion": "^11.2.10",