wip

sourcegraph · Dec 25, 2023 · 0b34059 · 0b34059
1 parent eff9aaa
commit 0b34059
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/provider/docs/package.json b/provider/docs/package.json
@@ -25,6 +25,7 @@
   "dependencies": {
     "@opencodegraph/provider": "workspace:*",
     "@xenova/transformers": "^2.12.1",
-    "env-paths": "^3.0.0"
+    "env-paths": "^3.0.0",
+    "onnxruntime-web": "*"
   }
 }
diff --git a/provider/docs/src/corpus/search/embeddings.ts b/provider/docs/src/corpus/search/embeddings.ts
@@ -1,7 +1,17 @@
-import { cos_sim, pipeline } from '@xenova/transformers'
+import { cos_sim, env, pipeline } from '@xenova/transformers'
+import * as onnxWeb from 'onnxruntime-web'
 import { memo, noopCache, type CorpusCache } from '../cache/cache'
 import { type CorpusSearchResult, type StoredCorpus } from '../corpus'
 
+// Force use of wasm backend for parity between Node.js and web.
+//
+// eslint-disable-next-line @typescript-eslint/ban-ts-comment
+// @ts-ignore
+// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-member-access
+env.onnx = onnxWeb.env
+// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+;(env as any).onnx.wasm.numThreads = 1
+
 export async function embeddingsSearch(
     storage: StoredCorpus,
     query: string,