isorite 0.2.0 → 0.3.1

package/dist/embed-worker.js

@@ -3212,14 +3212,15 @@ var init_ort_wasm_bundle_min = __esmMin((() => {
 //#endregion
 //#region packages/server/src/embed-worker.ts
 const MAX_TOKENS = 512;
-const { dir, dim } = workerData;
+const { dir, dim, pooling } = workerData;
 const ort = await Promise.resolve().then(() => (init_ort_wasm_bundle_min(), ort_wasm_bundle_min_exports));
 const { Tokenizer } = await import("@huggingface/tokenizers");
 ort.env.wasm.numThreads = 1;
 const WASM = "ort-wasm-simd-threaded.wasm";
 const beside = join(dirname(fileURLToPath(import.meta.url)), WASM);
 ort.env.wasm.wasmBinary = readFileSync(existsSync(beside) ? beside : join(dirname(createRequire(import.meta.url).resolve("onnxruntime-web")), WASM));
-const tok = new Tokenizer(JSON.parse(readFileSync(join(dir, "tokenizer.json"), "utf8")), JSON.parse(readFileSync(join(dir, "tokenizer_config.json"), "utf8")));
+const config = join(dir, "tokenizer_config.json");
+const tok = new Tokenizer(JSON.parse(readFileSync(join(dir, "tokenizer.json"), "utf8")), existsSync(config) ? JSON.parse(readFileSync(config, "utf8")) : {});
 const externalData = readdirSync(dir).filter((f) => f.endsWith(".onnx_data")).map((path) => ({
 	path,
 	data: readFileSync(join(dir, path))
@@ -3238,8 +3239,12 @@ async function embed(text) {
 		attention_mask: big(ids.map(() => 1)),
 		...typed ? { token_type_ids: big(ids.map(() => 0)) } : {}
 	});
-	const state = out.last_hidden_state ?? Object.values(out)[0];
-	return Float32Array.from(state.data.subarray(0, dim));
+	const state = out.last_hidden_state ?? Object.values(out)[0], width = state.dims.at(-1), data = state.data;
+	if (width !== dim) throw new Error(`its vectors have ${width} numbers, not ${dim}`);
+	if (pooling === "cls" || state.dims.length < 3) return Float32Array.from(data.subarray(0, dim));
+	const n = data.length / dim, mean = new Float32Array(dim);
+	for (let t = 0; t < n; t++) for (let k = 0; k < dim; k++) mean[k] += data[t * dim + k] / n;
+	return mean;
 }
 parentPort.on("message", (m) => {
 	embed(m.text).then((vector) => parentPort.postMessage({