aaronrene / knowtation public
inference-worker.mjs
94 lines 3.7 KB
Raw
sha256:bd1eb24434d281fe65142038c1bf217493421f5a86880c22dff7e6037d2d4700 feat(companion): add production macOS distribution path Human minor ⚠ breaking 24 days ago
1 import './inference-egress-guard.mjs';
2 import fs from 'node:fs/promises';
3 import path from 'node:path';
4 import { pathToFileURL } from 'node:url';
5
6 const runtimeRoot = path.resolve(process.argv[2] ?? '');
7 if (!path.isAbsolute(runtimeRoot)) throw new Error('PRIVATE_MODEL_UNAVAILABLE');
8
9 let model;
10 let tokenizer;
11 let ort;
12 let loadMilliseconds;
13
14 async function load() {
15 if (model) return;
16 const start = performance.now();
17 const ortRoot = path.join(runtimeRoot, 'node_modules/onnxruntime-web/dist');
18 ort = await import(pathToFileURL(path.join(ortRoot, 'ort.wasm.min.mjs')));
19 globalThis[Symbol.for('onnxruntime')] = ort;
20 const transformersPath = path.join(
21 runtimeRoot,
22 'node_modules/@huggingface/transformers/dist/transformers.web.js'
23 );
24 const { BertTokenizer, env } = await import(pathToFileURL(transformersPath));
25 env.allowRemoteModels = false;
26 env.useFSCache = false;
27 env.useBrowserCache = false;
28 ort.env.wasm.numThreads = 1;
29 ort.env.wasm.proxy = false;
30 ort.env.wasm.wasmPaths = `${ortRoot}/`;
31 ort.env.wasm.wasmBinary = await fs.readFile(path.join(ortRoot, 'ort-wasm-simd-threaded.wasm'));
32 const modelRoot = path.join(runtimeRoot, 'models/all-MiniLM-L6-v2');
33 tokenizer = new BertTokenizer(
34 JSON.parse(await fs.readFile(path.join(modelRoot, 'tokenizer.json'), 'utf8')),
35 JSON.parse(await fs.readFile(path.join(modelRoot, 'tokenizer_config.json'), 'utf8'))
36 );
37 model = await ort.InferenceSession.create(
38 await fs.readFile(path.join(modelRoot, 'onnx/model_quantized.onnx')),
39 { executionProviders: ['wasm'] }
40 );
41 loadMilliseconds = performance.now() - start;
42 }
43
44 async function embed(texts) {
45 if (!Array.isArray(texts) || texts.length === 0 || texts.length > 96 ||
46 texts.some((text) => typeof text !== 'string' || Buffer.byteLength(text) > 4096)) {
47 throw new Error('PRIVATE_INFERENCE_INPUT');
48 }
49 await load();
50 const inputs = tokenizer(texts, { padding: true, truncation: false });
51 if (inputs.input_ids.dims[1] > 256) throw new Error('PRIVATE_CONTEXT_LIMIT');
52 const feeds = Object.fromEntries(model.inputNames.map((name) => [
53 name,
54 new ort.Tensor('int64', inputs[name].data, inputs[name].dims),
55 ]));
56 const outputs = await model.run(feeds);
57 const hidden = outputs.last_hidden_state;
58 const [batch, length, dimensions] = hidden.dims;
59 const vectors = [];
60 for (let row = 0; row < batch; row += 1) {
61 const vector = Array(dimensions).fill(0);
62 let count = 0;
63 for (let token = 0; token < length; token += 1) {
64 if (inputs.attention_mask.data[row * length + token] === 0n) continue;
65 count += 1;
66 for (let column = 0; column < dimensions; column += 1) {
67 vector[column] += hidden.data[(row * length + token) * dimensions + column];
68 }
69 }
70 for (let column = 0; column < dimensions; column += 1) vector[column] /= count;
71 const norm = Math.sqrt(vector.reduce((sum, value) => sum + value * value, 0));
72 vectors.push(vector.map((value) => value / norm));
73 }
74 return vectors;
75 }
76
77 let queue = Promise.resolve();
78 process.on('message', (message) => {
79 queue = queue.then(async () => {
80 try {
81 if (message?.command !== 'embed' || Object.keys(message.arguments ?? {}).join(',') !== 'texts') {
82 throw new Error('PRIVATE_INFERENCE_INPUT');
83 }
84 const vectors = await embed(message.arguments.texts);
85 process.send?.({ id: message.id, ok: true, result: { vectors, loadMilliseconds } });
86 } catch (error) {
87 const safe = String(error?.message ?? '').startsWith('PRIVATE_')
88 ? error.message
89 : 'PRIVATE_MODEL_UNAVAILABLE';
90 process.send?.({ id: message?.id ?? 'invalid', ok: false, error: safe });
91 }
92 });
93 });
94 process.send?.({ ready: true, pid: process.pid });
File History 1 commit
sha256:bd1eb24434d281fe65142038c1bf217493421f5a86880c22dff7e6037d2d4700 feat(companion): add production macOS distribution path Human minor ⚠ 24 days ago