Sync from felladrin/MiniSearch@c75b7b9 (part 2)

server/biEncoderService.integration.test.ts

@@ -7,16 +7,22 @@
  *   npx vitest run --config vitest.integration.config.ts biEncoder
  */
 
+import { monitorEventLoopDelay } from "node:perf_hooks";
+import { Worker } from "node:worker_threads";
 import type { Tokenizer } from "@huggingface/tokenizers";
-import { InferenceSession, Tensor } from "onnxruntime-node";
-import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
+import { type InferenceSession, Tensor } from "onnxruntime-node";
+import { afterAll, beforeAll, describe, expect, it } from "vitest";
 import {
-  BATCH_ROWS,
   getBiEncoderStatus,
   scorePassages,
   startBiEncoderService,
   stopBiEncoderService,
 } from "./biEncoderService";
+import {
+  BATCH_ROWS,
+  type BiEncoderRequest,
+  type BiEncoderResponse,
+} from "./biEncoderWorkerProtocol";
 import { loadOnnxModel } from "./utils/onnxModelLoader";
 
 const MODEL_HF_REPO =
@@ -178,63 +184,6 @@ describe("biEncoderService batched scoring (#2715)", () => {
     await refSession?.release();
   });
 
-  it("runs one forward pass per length bucket, not one per passage", async () => {
-    // Zigzag lengths with the sort's guarantee removed: bucket 1's last
-    // row (190 words) is longer than bucket 2's last row (25 words), so
-    // an unsorted pass sets bucket widths from those rows and they come
-    // out decreasing — the non-decreasing-width assertion below then
-    // fails. With the sort, widths are non-decreasing by construction.
-    const wordCounts = [
-      200, 10, 150, 5, 180, 20, 120, 8, 90, 160, 30, 110, 60, 140, 190, 170, 15,
-      130, 70, 25,
-    ];
-    const passages = wordCounts.map((n, i) => makePassage(n, i));
-
-    // The typings expose `InferenceSession` as a factory interface and hide
-    // the class prototype, but at runtime `run` is a prototype method, so a
-    // prototype spy intercepts the service's already-created session.
-    const sessionClass = InferenceSession as unknown as {
-      prototype: { run: InferenceSession["run"] };
-    };
-    const runSpy = vi.spyOn(sessionClass.prototype, "run");
-    await scorePassages("how do alpha bravo passages score", passages);
-    // Capture before mockRestore: restoring also resets the call history.
-    const runDims = runSpy.mock.calls.map(
-      ([feeds]) =>
-        (feeds as unknown as { input_ids: { dims: number[] } }).input_ids.dims,
-    );
-    runSpy.mockRestore();
-
-    // The query rides in the batch, so 21 rows go through ceil(21 / B)
-    // bucketed runs instead of 21 per-passage runs.
-    const expectedBuckets = Math.ceil((passages.length + 1) / BATCH_ROWS);
-    expect(runDims.length).toBe(expectedBuckets);
-    expect(runDims.length).toBeLessThan(passages.length + 1);
-
-    let totalRows = 0;
-    for (const dims of runDims) {
-      expect(dims.length).toBe(2);
-      expect(dims[0]).toBeLessThanOrEqual(BATCH_ROWS);
-      totalRows += dims[0];
-    }
-    expect(totalRows).toBe(passages.length + 1);
-
-    // Every non-final bucket is full, so it really carries multiple rows.
-    // Only the final bucket may be short — with BATCH_ROWS of 10 or 20 a
-    // 21-row pool leaves it exactly one row, so the >1 check must not
-    // depend on the tuning knob.
-    for (let i = 0; i < runDims.length - 1; i++) {
-      expect(runDims[i][0]).toBeGreaterThan(1);
-    }
-
-    // Bucket widths are non-decreasing across runs: the sort by token
-    // length put the short rows first. Delete the `.sort()` in encodeBatch
-    // and this fails.
-    for (let i = 1; i < runDims.length; i++) {
-      expect(runDims[i][1]).toBeGreaterThanOrEqual(runDims[i - 1][1]);
-    }
-  });
-
   it("keeps batched scores within 1e-5 of the per-text path", async () => {
     const wordCounts = [
       3, 8, 15, 40, 90, 150, 250, 400, 5, 20, 60, 120, 200, 300, 12, 35, 75,
@@ -310,3 +259,156 @@ describe("biEncoderService batched scoring (#2715)", () => {
     expect(isSorted).toBe(false);
   });
 });
+
+/**
+ * Runs one scoring request against a worker of its own and returns its reply.
+ *
+ * The session lives in the worker now, so a `session.run` spy on this thread
+ * would see nothing. The worker reports the dims of each forward pass instead,
+ * which is what the bucketing assertions read.
+ */
+async function scoreInOwnWorker(
+  query: string,
+  passages: string[],
+): Promise<Extract<BiEncoderResponse, { type: "scores" }>> {
+  const worker = new Worker(new URL("./biEncoderWorker.ts", import.meta.url));
+
+  try {
+    await new Promise<void>((resolve, reject) => {
+      worker.once("message", (message: BiEncoderResponse) => {
+        if (message.type === "ready") resolve();
+        else reject(new Error(`Expected "ready", got "${message.type}"`));
+      });
+      worker.once("error", reject);
+    });
+
+    return await new Promise((resolve, reject) => {
+      worker.once("message", (message: BiEncoderResponse) => {
+        if (message.type === "scores") resolve(message);
+        else if (message.type === "failed") reject(new Error(message.message));
+      });
+      worker.once("error", reject);
+
+      const request: BiEncoderRequest = {
+        type: "score",
+        id: 1,
+        query,
+        passages,
+      };
+      worker.postMessage(request);
+    });
+  } finally {
+    await worker.terminate();
+  }
+}
+
+describe("biEncoderWorker bucketing (#2715, #2730)", () => {
+  it("runs one forward pass per length bucket, not one per passage", async () => {
+    // Zigzag lengths with the sort's guarantee removed: bucket 1's last
+    // row (190 words) is longer than bucket 2's last row (25 words), so
+    // an unsorted pass sets bucket widths from those rows and they come
+    // out decreasing — the non-decreasing-width assertion below then
+    // fails. With the sort, widths are non-decreasing by construction.
+    const wordCounts = [
+      200, 10, 150, 5, 180, 20, 120, 8, 90, 160, 30, 110, 60, 140, 190, 170, 15,
+      130, 70, 25,
+    ];
+    const passages = wordCounts.map((n, i) => makePassage(n, i));
+
+    const { scores, runDimensions } = await scoreInOwnWorker(
+      "how do alpha bravo passages score",
+      passages,
+    );
+
+    expect(scores).toHaveLength(passages.length);
+
+    // The query rides in the batch, so 21 rows go through ceil(21 / B)
+    // bucketed runs instead of 21 per-passage runs.
+    const expectedBuckets = Math.ceil((passages.length + 1) / BATCH_ROWS);
+    expect(runDimensions.length).toBe(expectedBuckets);
+    expect(runDimensions.length).toBeLessThan(passages.length + 1);
+
+    let totalRows = 0;
+    for (const dims of runDimensions) {
+      expect(dims.length).toBe(2);
+      expect(dims[0]).toBeLessThanOrEqual(BATCH_ROWS);
+      totalRows += dims[0];
+    }
+    expect(totalRows).toBe(passages.length + 1);
+
+    // Every non-final bucket is full, so it really carries multiple rows.
+    // Only the final bucket may be short — with BATCH_ROWS of 10 or 20 a
+    // 21-row pool leaves it exactly one row, so the >1 check must not
+    // depend on the tuning knob.
+    for (let i = 0; i < runDimensions.length - 1; i++) {
+      expect(runDimensions[i][0]).toBeGreaterThan(1);
+    }
+
+    // Bucket widths are non-decreasing across runs: the sort by token
+    // length put the short rows first. Delete the `.sort()` in encodeBatch
+    // and this fails.
+    for (let i = 1; i < runDimensions.length; i++) {
+      expect(runDimensions[i][1]).toBeGreaterThanOrEqual(
+        runDimensions[i - 1][1],
+      );
+    }
+  });
+});
+
+describe("biEncoderService main-thread cost (#2730)", () => {
+  beforeAll(async () => {
+    await startBiEncoderService();
+  });
+
+  afterAll(async () => {
+    await stopBiEncoderService();
+  });
+
+  /**
+   * The acceptance measurement from #2730. On the main thread the same pass
+   * blocked for ~815 ms (x86_64, 32 logical cores); with the session in the
+   * worker it measures ~1.5 ms. The bar is left at the issue's 20 ms so the
+   * test reports a regression rather than host-to-host noise.
+   */
+  it("keeps the main-thread event-loop block under 20 ms for a 200-passage pass", async () => {
+    const passages = Array.from({ length: 200 }, (_, i) =>
+      makePassage(5 + ((i * 37) % 300), i),
+    );
+    const query = "how do alpha bravo passages score";
+
+    // One untimed pass, so the worker's lazily allocated arenas are not
+    // charged to the measurement.
+    await scorePassages(query, passages.slice(0, 20));
+
+    const histogram = monitorEventLoopDelay({ resolution: 1 });
+    histogram.enable();
+    const scores = await scorePassages(query, passages);
+    histogram.disable();
+
+    expect(scores).toHaveLength(passages.length);
+    expect(histogram.max / 1e6).toBeLessThan(20);
+  });
+});
+
+describe("biEncoderService when the worker goes away (#2730)", () => {
+  it("answers in-flight scoring with empty scores and reports not ready", async () => {
+    await startBiEncoderService();
+    expect(await getBiEncoderStatus()).toBe(true);
+
+    const passages = Array.from({ length: 200 }, (_, i) =>
+      makePassage(5 + ((i * 37) % 300), i),
+    );
+    // Fired but not awaited: the worker is taken away underneath it. An
+    // unanswered request would hang the page-content read that made it, so
+    // it has to come back empty, which is the signal to rank lexically.
+    const pending = scorePassages(
+      "how do alpha bravo passages score",
+      passages,
+    );
+    await stopBiEncoderService();
+
+    await expect(pending).resolves.toEqual([]);
+    expect(await getBiEncoderStatus()).toBe(false);
+    expect(await scorePassages("query", ["one"])).toEqual([]);
+  });
+});

server/biEncoderService.ts

@@ -1,229 +1,124 @@
+/**
+ * The main thread's side of the bi-encoder. It owns no model and runs no
+ * inference: it spawns `biEncoderWorker.ts`, forwards scoring requests to it
+ * and matches the replies back to their callers.
+ *
+ * The four exported functions keep the signatures they had when the session
+ * lived here, so no caller changed. What changed is that a scoring pass no
+ * longer stops the event loop: a 200-passage pool blocked the main thread for
+ * ~815 ms before and stays under 20 ms now (see docs/page-content.md).
+ */
+
 import path from "node:path";
-import type { Tokenizer } from "@huggingface/tokenizers";
-import { type InferenceSession, Tensor } from "onnxruntime-node";
-import { createModelLogger, loadOnnxModel } from "./utils/onnxModelLoader.ts";
+import { Worker } from "node:worker_threads";
+import type {
+  BiEncoderRequest,
+  BiEncoderResponse,
+} from "./biEncoderWorkerProtocol.ts";
+import { createModelLogger } from "./utils/onnxModelLoader.ts";
 
 const printMessage = createModelLogger(path.basename(import.meta.url));
 
-const MODEL_HF_REPO =
-  "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2";
-
-/**
- * The ONNX export. ~450 MB, multilingual (50+ languages), 384-dimensional
- * embeddings. CPU cost depends on passage length and pool size; the page-content
- * selector limits dense scoring to 256 passages. See docs/page-content.md.
- */
-const MODEL_HF_FILE = "onnx/model.onnx";
+let worker: Worker | null = null;
+let isReady = false;
+let nextRequestId = 0;
 
-/**
- * Maximum tokens per encoding. The model was trained with a 256-token limit;
- * passages longer than that are truncated from the end, which is where the
- * passage content sits after the query prefix.
- */
-const MAX_SEQUENCE_LENGTH = 256;
+/** Resolvers for requests the worker has not answered yet, keyed by their id. */
+const pendingScores = new Map<number, (scores: number[]) => void>();
 
 /**
- * Rows per batched forward pass. The old chunk of 64 wrapped one-text-per-call
- * runs and bought no concurrency; now that a call really carries its rows, B
- * trades padding waste and cache pressure against per-call setup. B=64 measured
- * slower than the per-text path; the sweep cannot separate the two plausible
- * causes — a 64-row bucket spans a wider length range, so padding waste
- * grows, and at L=256 with 12 heads the attention score tensor alone is
- * ~192 MiB live, which pressures cache. B=16 (~48 MiB) measured fastest on
- * a 200-passage pool on this host (x86_64, 32 cpus): 1.88x wall time and a
- * 223ms worst event-loop block against 1560ms for the per-text path.
+ * Answers every in-flight request with empty scores. Empty is the signal
+ * `pageContentService` already reads as "dense scoring is unavailable, rank
+ * lexically", so a worker that dies degrades the ranking instead of leaving
+ * page-content reads waiting for a reply that is never coming.
  */
-export const BATCH_ROWS = 16;
-
-let isReady = false;
-let session: InferenceSession | null = null;
-let tokenizer: Tokenizer | null = null;
-let padTokenId: number | null = null;
+function resolvePendingEmpty() {
+  for (const resolve of pendingScores.values()) resolve([]);
+  pendingScores.clear();
+}
 
 /**
- * Encodes a batch of texts into normalized embedding vectors with one
- * `session.run()` per length bucket.
- *
- * Rows are sorted by token length and cut into buckets of at most
- * `BATCH_ROWS`, so each batch pads to its own near-uniform width instead of
- * the global max: attention is O(L^2), and padding a mixed-length chunk out
- * to 256 can multiply the FLOPs several-fold and come out slower than the
- * per-text path.
- *
- * Right-padding with the real pad id under an attention mask is score-safe for
- * this fp32 export: it has no per-tensor dynamic quantization scale, so a
- * padded row cannot move another row's scale, and the real tokens' hidden
- * states are bit-identical to encoding the text alone. (The cross-encoder is
- * dynamically quantized and cannot batch for that reason; see the note in
- * `rerankerService.ts`.)
- *
- * Buckets run one after another with a plain `await` — deliberately NOT a
- * `Promise.all` across buckets. Each await hands the event loop a turn
- * between buckets, which is where the block-length win comes from. Do not
- * "optimize" the buckets back into a `Promise.all`: onnxruntime-node runs
- * inference synchronously on the JS thread regardless, so it would add no
- * concurrency while taking away those between-bucket turns.
+ * Drops a worker that died. Guarded on identity because a terminated worker's
+ * `exit` still arrives after `startBiEncoderService` has installed its
+ * replacement, and that late event must not mark the new one dead.
  */
-async function encodeBatch(
-  activeSession: InferenceSession,
-  loadedTokenizer: Tokenizer,
-  activePadTokenId: number,
-  texts: string[],
-): Promise<Float32Array[]> {
-  const tokenized = texts
-    .map((text, index) => {
-      const { ids, attention_mask } = loadedTokenizer.encode(text);
-      const slicedIds = ids.slice(0, MAX_SEQUENCE_LENGTH);
-      const slicedMask = attention_mask.slice(0, MAX_SEQUENCE_LENGTH);
-      // The real-token length comes from the mask's leading 1s, not from
-      // trusting every returned id as real. At @huggingface/tokenizers
-      // 0.2.0 the mask is all 1s, so this is a no-op today; it keeps the
-      // row correct if a future tokenizer pads or marks truncation in the
-      // mask (the cached tokenizer.json already declares a padding
-      // strategy with pad_id 1).
-      const firstPad = slicedMask.indexOf(0);
-      const realLength = firstPad === -1 ? slicedIds.length : firstPad;
-      return { index, ids: slicedIds.slice(0, realLength) };
-    })
-    .sort((a, b) => a.ids.length - b.ids.length);
-
-  const embeddings: Float32Array[] = new Array(texts.length);
-
-  for (let start = 0; start < tokenized.length; start += BATCH_ROWS) {
-    const rows = tokenized.slice(start, start + BATCH_ROWS);
-    // Sorted ascending, so the last row sets the bucket width.
-    const bucketLength = rows[rows.length - 1].ids.length;
-    const dimensions = [rows.length, bucketLength];
-
-    const inputIds = new BigInt64Array(rows.length * bucketLength).fill(
-      BigInt(activePadTokenId),
-    );
-    const attentionMask = new BigInt64Array(rows.length * bucketLength);
-    // The export declares `token_type_ids` and ONNX Runtime refuses to run
-    // with a declared input missing. Every text is one segment, so zeros.
-    const tokenTypeIds = new BigInt64Array(rows.length * bucketLength);
-
-    for (let row = 0; row < rows.length; row++) {
-      const base = row * bucketLength;
-      const { ids } = rows[row];
-      for (let t = 0; t < ids.length; t++) {
-        inputIds[base + t] = BigInt(ids[t]);
-        attentionMask[base + t] = 1n;
-      }
-    }
+function forgetWorker(deadWorker: Worker) {
+  if (worker !== deadWorker) return;
+  worker = null;
+  isReady = false;
+  resolvePendingEmpty();
+}
 
-    const { last_hidden_state } = await activeSession.run({
-      input_ids: new Tensor("int64", inputIds, dimensions),
-      attention_mask: new Tensor("int64", attentionMask, dimensions),
-      token_type_ids: new Tensor("int64", tokenTypeIds, dimensions),
-    });
+export async function startBiEncoderService() {
+  if (worker) await stopBiEncoderService();
 
-    const hidden = last_hidden_state.data as Float32Array;
-    const dim = last_hidden_state.dims[2];
-    // Take the row stride from the output tensor, not from the bucket
-    // length we asked for: same value today, but the pooling can never
-    // read the wrong row if the export's layout ever changes.
-    const seqStride = last_hidden_state.dims[1];
+  const biEncoderWorker = new Worker(
+    new URL("./biEncoderWorker.ts", import.meta.url),
+  );
+  worker = biEncoderWorker;
 
-    for (let row = 0; row < rows.length; row++) {
-      // Right-padding keeps every real token inside `ids.length`, so
-      // pooling over that count never touches a pad. (Checking the mask
-      // here instead would have to compare against 0n: BigInt64Array
-      // entries are never `=== 0`.)
-      const rowLength = rows[row].ids.length;
-      const pooled = new Float32Array(dim);
-      for (let t = 0; t < rowLength; t++) {
-        const offset = (row * seqStride + t) * dim;
-        for (let d = 0; d < dim; d++) {
-          pooled[d] += hidden[offset + d];
-        }
-      }
-      if (rowLength > 0) {
-        for (let d = 0; d < dim; d++) {
-          pooled[d] /= rowLength;
-        }
-      }
+  biEncoderWorker.on("message", (message: BiEncoderResponse) => {
+    if (message.type === "ready") return;
 
-      // L2 normalize.
-      let norm = 0;
-      for (let d = 0; d < dim; d++) {
-        norm += pooled[d] * pooled[d];
-      }
-      norm = Math.sqrt(norm);
-      if (norm > 0) {
-        for (let d = 0; d < dim; d++) {
-          pooled[d] /= norm;
-        }
-      }
+    const resolve = pendingScores.get(message.id);
+    if (!resolve) return;
+    pendingScores.delete(message.id);
 
-      embeddings[rows[row].index] = pooled;
+    if (message.type === "failed") {
+      printMessage(`Bi-encoder scoring failed: ${message.message}`);
+      resolve([]);
+      return;
     }
-  }
-
-  return embeddings;
-}
-
-/**
- * Encodes a single text into a normalized embedding vector.
- */
-async function encode(
-  activeSession: InferenceSession,
-  loadedTokenizer: Tokenizer,
-  activePadTokenId: number,
-  text: string,
-): Promise<Float32Array> {
-  const [embedding] = await encodeBatch(
-    activeSession,
-    loadedTokenizer,
-    activePadTokenId,
-    [text],
-  );
-  return embedding;
-}
 
-/**
- * Computes cosine similarity between a query embedding and passage embeddings.
- * Both are assumed to be L2-normalized, so cosine similarity = dot product.
- */
-function cosineSimilarities(
-  query: Float32Array,
-  passages: Float32Array[],
-): number[] {
-  return passages.map((passage) => {
-    let sum = 0;
-    for (let d = 0; d < query.length; d++) {
-      sum += query[d] * passage[d];
-    }
-    return sum;
+    resolve(message.scores);
   });
-}
 
-export async function startBiEncoderService() {
-  printMessage("Preparing bi-encoder model...");
-  const loaded = await loadOnnxModel(MODEL_HF_REPO, MODEL_HF_FILE);
-  if (loaded.padTokenId === null) {
-    throw new Error(
-      `The bi-encoder tokenizer (${MODEL_HF_REPO}) declares no pad_token; batched scoring needs the real pad id and will not assume one`,
-    );
-  }
-  session = loaded.session;
-  tokenizer = loaded.tokenizer;
-  padTokenId = loaded.padTokenId;
+  biEncoderWorker.on("error", (error) => {
+    printMessage(`Bi-encoder worker failed: ${error.message}`);
+    forgetWorker(biEncoderWorker);
+  });
 
-  // Warm up with a test encoding.
-  await encode(session, tokenizer, padTokenId, "test query");
+  biEncoderWorker.on("exit", () => forgetWorker(biEncoderWorker));
+
+  await new Promise<void>((resolve, reject) => {
+    const stopListening = () => {
+      biEncoderWorker.off("message", onMessage);
+      biEncoderWorker.off("error", onError);
+      biEncoderWorker.off("exit", onExit);
+    };
+    const onMessage = (message: BiEncoderResponse) => {
+      if (message.type !== "ready") return;
+      stopListening();
+      resolve();
+    };
+    const onError = (error: Error) => {
+      stopListening();
+      reject(error);
+    };
+    const onExit = (code: number) => {
+      stopListening();
+      reject(
+        new Error(
+          `The bi-encoder worker exited with code ${code} before it was ready`,
+        ),
+      );
+    };
+
+    biEncoderWorker.on("message", onMessage);
+    biEncoderWorker.on("error", onError);
+    biEncoderWorker.on("exit", onExit);
+  });
 
   isReady = true;
   printMessage("Bi-encoder service ready!");
 }
 
 export async function stopBiEncoderService() {
+  const runningWorker = worker;
   isReady = false;
-  const currentSession = session;
-  session = null;
-  tokenizer = null;
-  padTokenId = null;
-  await currentSession?.release();
+  worker = null;
+  resolvePendingEmpty();
+  await runningWorker?.terminate();
 }
 
 export async function getBiEncoderStatus() {
@@ -232,31 +127,35 @@ export async function getBiEncoderStatus() {
 
 /**
  * Returns dense (semantic) scores for passages given a query.
- * Falls back to empty array when the model is not loaded.
+ * Falls back to empty array when the worker is not running.
  */
 export async function scorePassages(
   query: string,
   passages: string[],
 ): Promise<number[]> {
-  if (!session || !tokenizer || padTokenId === null || passages.length === 0) {
+  const runningWorker = worker;
+  if (!runningWorker || !isReady || passages.length === 0) {
     return [];
   }
 
-  const activeSession = session;
-  const loadedTokenizer = tokenizer;
-  const activePadTokenId = padTokenId;
-
-  // The query rides in the same batched pass as the passages instead of
-  // running alone, saving one forward pass. Its position in the result is
-  // guaranteed by the index-based restore inside encodeBatch
-  // (`embeddings[rows[row].index]`), not by where the length sort happens
-  // to place it, so it stays correct whatever the bucketing does.
-  const [queryEmbedding, ...passageEmbeddings] = await encodeBatch(
-    activeSession,
-    loadedTokenizer,
-    activePadTokenId,
-    [query, ...passages],
-  );
-
-  return cosineSimilarities(queryEmbedding, passageEmbeddings);
+  // Ids are per-process and monotonic because the service can have several
+  // scoring passes in flight at once: two concurrent page-content reads each
+  // send their own request, and the replies can come back in either order.
+  const id = nextRequestId++;
+
+  return new Promise<number[]>((resolve) => {
+    pendingScores.set(id, resolve);
+
+    const request: BiEncoderRequest = { type: "score", id, query, passages };
+    try {
+      runningWorker.postMessage(request);
+    } catch (error) {
+      // The worker was terminated between the check above and this send.
+      pendingScores.delete(id);
+      printMessage(
+        `Could not reach the bi-encoder worker: ${error instanceof Error ? error.message : String(error)}`,
+      );
+      resolve([]);
+    }
+  });
 }

server/biEncoderWorker.ts

@@ -0,0 +1,279 @@
+/**
+ * The bi-encoder's ONNX session, running off the main thread.
+ *
+ * `onnxruntime-node` runs inference synchronously on whichever JS thread calls
+ * it, so scoring a 200-passage pool used to stop the server's event loop for
+ * hundreds of milliseconds — no other request could be served while a single
+ * page-content read was ranking its passages. Here that work blocks only this
+ * thread. The model is loaded here for the same reason: the ~450 MB of weights
+ * must live in the worker's memory, not the main thread's.
+ *
+ * Node loads this file directly (type stripping), so keep it to syntax that
+ * erases: no enums, no parameter properties, no namespaces.
+ */
+
+import { availableParallelism } from "node:os";
+import path from "node:path";
+import { parentPort } from "node:worker_threads";
+import type { Tokenizer } from "@huggingface/tokenizers";
+import { type InferenceSession, Tensor } from "onnxruntime-node";
+import {
+  BATCH_ROWS,
+  type BiEncoderRequest,
+  type BiEncoderResponse,
+} from "./biEncoderWorkerProtocol.ts";
+import { createModelLogger, loadOnnxModel } from "./utils/onnxModelLoader.ts";
+
+const printMessage = createModelLogger(path.basename(import.meta.url));
+
+const MODEL_HF_REPO =
+  "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2";
+
+/**
+ * The ONNX export. ~450 MB, multilingual (50+ languages), 384-dimensional
+ * embeddings. CPU cost depends on passage length and pool size; the page-content
+ * selector limits dense scoring to 256 passages. See docs/page-content.md.
+ */
+const MODEL_HF_FILE = "onnx/model.onnx";
+
+/**
+ * Maximum tokens per encoding. The model was trained with a 256-token limit;
+ * passages longer than that are truncated from the end, which is where the
+ * passage content sits after the query prefix.
+ */
+const MAX_SEQUENCE_LENGTH = 256;
+
+/**
+ * Half the reported parallelism, which is the physical core count on the
+ * hyperthreaded x86 hosts this runs on, and at least one thread everywhere
+ * else.
+ *
+ * ONNX Runtime otherwise sizes its intra-op pool from the logical core count
+ * and its threads spin before yielding, so it oversubscribes every core and
+ * leaves nothing for the thread that has to answer HTTP. Measured on a
+ * 200-passage pool on this host (x86_64, 32 logical cores), worst main-thread
+ * event-loop delay against wall time for the pass:
+ *
+ *   ORT default (32): 31-38 ms, 7.1 s
+ *   16:                1-4 ms,  2.6 s
+ *   8:                 1-2 ms,  4.2 s
+ *   4:                 1-1 ms,  5.7 s
+ *
+ * Capping is not a trade here: 16 was both the quietest and the fastest, and
+ * the oversubscribed default was the slowest setting measured.
+ */
+const INTRA_OP_THREADS = Math.max(1, Math.floor(availableParallelism() / 2));
+
+interface EncodeResult {
+  embeddings: Float32Array[];
+  /** `input_ids` dims of each forward pass, in run order. */
+  runDimensions: number[][];
+}
+
+/**
+ * Encodes a batch of texts into normalized embedding vectors with one
+ * `session.run()` per length bucket.
+ *
+ * Rows are sorted by token length and cut into buckets of at most
+ * `BATCH_ROWS`, so each batch pads to its own near-uniform width instead of
+ * the global max: attention is O(L^2), and padding a mixed-length chunk out
+ * to 256 can multiply the FLOPs several-fold and come out slower than the
+ * per-text path.
+ *
+ * Right-padding with the real pad id under an attention mask is score-safe for
+ * this fp32 export: it has no per-tensor dynamic quantization scale, so a
+ * padded row cannot move another row's scale, and the real tokens' hidden
+ * states are bit-identical to encoding the text alone. (The cross-encoder is
+ * dynamically quantized and cannot batch for that reason; see the note in
+ * `rerankerService.ts`.)
+ *
+ * Buckets run one after another with a plain `await` — deliberately NOT a
+ * `Promise.all` across buckets. Each await hands this thread's event loop a
+ * turn between buckets, which is what lets two overlapping scoring requests
+ * interleave instead of one waiting out the other. Do not "optimize" the
+ * buckets back into a `Promise.all`: onnxruntime-node runs inference
+ * synchronously on the calling thread regardless, so it would add no
+ * concurrency while taking away those between-bucket turns.
+ */
+async function encodeBatch(
+  activeSession: InferenceSession,
+  loadedTokenizer: Tokenizer,
+  activePadTokenId: number,
+  texts: string[],
+): Promise<EncodeResult> {
+  const tokenized = texts
+    .map((text, index) => {
+      const { ids, attention_mask } = loadedTokenizer.encode(text);
+      const slicedIds = ids.slice(0, MAX_SEQUENCE_LENGTH);
+      const slicedMask = attention_mask.slice(0, MAX_SEQUENCE_LENGTH);
+      // The real-token length comes from the mask's leading 1s, not from
+      // trusting every returned id as real. At @huggingface/tokenizers
+      // 0.2.0 the mask is all 1s, so this is a no-op today; it keeps the
+      // row correct if a future tokenizer pads or marks truncation in the
+      // mask (the cached tokenizer.json already declares a padding
+      // strategy with pad_id 1).
+      const firstPad = slicedMask.indexOf(0);
+      const realLength = firstPad === -1 ? slicedIds.length : firstPad;
+      return { index, ids: slicedIds.slice(0, realLength) };
+    })
+    .sort((a, b) => a.ids.length - b.ids.length);
+
+  const embeddings: Float32Array[] = new Array(texts.length);
+  const runDimensions: number[][] = [];
+
+  for (let start = 0; start < tokenized.length; start += BATCH_ROWS) {
+    const rows = tokenized.slice(start, start + BATCH_ROWS);
+    // Sorted ascending, so the last row sets the bucket width.
+    const bucketLength = rows[rows.length - 1].ids.length;
+    const dimensions = [rows.length, bucketLength];
+
+    const inputIds = new BigInt64Array(rows.length * bucketLength).fill(
+      BigInt(activePadTokenId),
+    );
+    const attentionMask = new BigInt64Array(rows.length * bucketLength);
+    // The export declares `token_type_ids` and ONNX Runtime refuses to run
+    // with a declared input missing. Every text is one segment, so zeros.
+    const tokenTypeIds = new BigInt64Array(rows.length * bucketLength);
+
+    for (let row = 0; row < rows.length; row++) {
+      const base = row * bucketLength;
+      const { ids } = rows[row];
+      for (let t = 0; t < ids.length; t++) {
+        inputIds[base + t] = BigInt(ids[t]);
+        attentionMask[base + t] = 1n;
+      }
+    }
+
+    runDimensions.push(dimensions);
+
+    const { last_hidden_state } = await activeSession.run({
+      input_ids: new Tensor("int64", inputIds, dimensions),
+      attention_mask: new Tensor("int64", attentionMask, dimensions),
+      token_type_ids: new Tensor("int64", tokenTypeIds, dimensions),
+    });
+
+    const hidden = last_hidden_state.data as Float32Array;
+    const dim = last_hidden_state.dims[2];
+    // Take the row stride from the output tensor, not from the bucket
+    // length we asked for: same value today, but the pooling can never
+    // read the wrong row if the export's layout ever changes.
+    const seqStride = last_hidden_state.dims[1];
+
+    for (let row = 0; row < rows.length; row++) {
+      // Right-padding keeps every real token inside `ids.length`, so
+      // pooling over that count never touches a pad. (Checking the mask
+      // here instead would have to compare against 0n: BigInt64Array
+      // entries are never `=== 0`.)
+      const rowLength = rows[row].ids.length;
+      const pooled = new Float32Array(dim);
+      for (let t = 0; t < rowLength; t++) {
+        const offset = (row * seqStride + t) * dim;
+        for (let d = 0; d < dim; d++) {
+          pooled[d] += hidden[offset + d];
+        }
+      }
+      if (rowLength > 0) {
+        for (let d = 0; d < dim; d++) {
+          pooled[d] /= rowLength;
+        }
+      }
+
+      // L2 normalize.
+      let norm = 0;
+      for (let d = 0; d < dim; d++) {
+        norm += pooled[d] * pooled[d];
+      }
+      norm = Math.sqrt(norm);
+      if (norm > 0) {
+        for (let d = 0; d < dim; d++) {
+          pooled[d] /= norm;
+        }
+      }
+
+      embeddings[rows[row].index] = pooled;
+    }
+  }
+
+  return { embeddings, runDimensions };
+}
+
+/**
+ * Computes cosine similarity between a query embedding and passage embeddings.
+ * Both are assumed to be L2-normalized, so cosine similarity = dot product.
+ */
+function cosineSimilarities(
+  query: Float32Array,
+  passages: Float32Array[],
+): number[] {
+  return passages.map((passage) => {
+    let sum = 0;
+    for (let d = 0; d < query.length; d++) {
+      sum += query[d] * passage[d];
+    }
+    return sum;
+  });
+}
+
+if (!parentPort) {
+  throw new Error(
+    "biEncoderWorker.ts is a worker entry point and must be started with node:worker_threads",
+  );
+}
+
+const port = parentPort;
+
+printMessage("Preparing bi-encoder model...");
+const loaded = await loadOnnxModel(MODEL_HF_REPO, MODEL_HF_FILE, {
+  sessionOptions: { intraOpNumThreads: INTRA_OP_THREADS },
+});
+if (loaded.padTokenId === null) {
+  throw new Error(
+    `The bi-encoder tokenizer (${MODEL_HF_REPO}) declares no pad_token; batched scoring needs the real pad id and will not assume one`,
+  );
+}
+const session = loaded.session;
+const tokenizer = loaded.tokenizer;
+const padTokenId = loaded.padTokenId;
+
+// Warm up with a test encoding, so the first real request does not pay for
+// the lazily allocated arenas of the first run.
+await encodeBatch(session, tokenizer, padTokenId, ["test query"]);
+
+port.on("message", async (request: BiEncoderRequest) => {
+  if (request.type !== "score") return;
+
+  try {
+    // The query rides in the same batched pass as the passages instead of
+    // running alone, saving one forward pass. Its position in the result is
+    // guaranteed by the index-based restore inside encodeBatch
+    // (`embeddings[rows[row].index]`), not by where the length sort happens
+    // to place it, so it stays correct whatever the bucketing does.
+    const { embeddings, runDimensions } = await encodeBatch(
+      session,
+      tokenizer,
+      padTokenId,
+      [request.query, ...request.passages],
+    );
+    const [queryEmbedding, ...passageEmbeddings] = embeddings;
+
+    const response: BiEncoderResponse = {
+      type: "scores",
+      id: request.id,
+      scores: cosineSimilarities(queryEmbedding, passageEmbeddings),
+      runDimensions,
+    };
+    port.postMessage(response);
+  } catch (error) {
+    // The caller falls back to lexical ranking on empty scores, so one bad
+    // request must not take the worker down with it.
+    const response: BiEncoderResponse = {
+      type: "failed",
+      id: request.id,
+      message: error instanceof Error ? error.message : String(error),
+    };
+    port.postMessage(response);
+  }
+});
+
+const ready: BiEncoderResponse = { type: "ready" };
+port.postMessage(ready);

server/biEncoderWorkerProtocol.ts

@@ -0,0 +1,49 @@
+/**
+ * The messages `biEncoderService.ts` and `biEncoderWorker.ts` exchange, kept in
+ * their own module so the main thread can name them without importing the
+ * worker: importing that module would load the ~450 MB model into the very
+ * thread the worker exists to keep free.
+ *
+ * Only plain data crosses the boundary. Embeddings stay inside the worker and
+ * the cosine step runs there, so a scoring pass sends strings out and brings
+ * `number[]` back rather than copying 384 floats per passage.
+ */
+
+/**
+ * Rows per batched forward pass. The old chunk of 64 wrapped one-text-per-call
+ * runs and bought no concurrency; now that a call really carries its rows, B
+ * trades padding waste and cache pressure against per-call setup. B=64 measured
+ * slower than the per-text path; the sweep cannot separate the two plausible
+ * causes — a 64-row bucket spans a wider length range, so padding waste
+ * grows, and at L=256 with 12 heads the attention score tensor alone is
+ * ~192 MiB live, which pressures cache. B=16 (~48 MiB) measured fastest on
+ * a 200-passage pool on this host (x86_64, 32 cpus): 1.88x wall time and a
+ * 223ms worst event-loop block against 1560ms for the per-text path.
+ */
+export const BATCH_ROWS = 16;
+
+/** Scoring work sent to the worker. `id` correlates the reply. */
+export interface BiEncoderRequest {
+  type: "score";
+  id: number;
+  query: string;
+  passages: string[];
+}
+
+export type BiEncoderResponse =
+  /** The model is loaded and warmed up; the worker accepts scoring now. */
+  | { type: "ready" }
+  | {
+      type: "scores";
+      id: number;
+      scores: number[];
+      /**
+       * The `input_ids` dims of each forward pass this request ran, in run
+       * order. The session no longer lives on the main thread, so this is the
+       * only way a test can assert on the bucketing; it is a few pairs of
+       * numbers per request.
+       */
+      runDimensions: number[][];
+    }
+  /** Scoring threw inside the worker. The service reads this as empty scores. */
+  | { type: "failed"; id: number; message: string };

server/utils/onnxModelLoader.ts

@@ -23,6 +23,21 @@ export function createModelLogger(moduleName: string) {
   return printMessage;
 }
 
+/**
+ * Session knobs a caller may override. Left unset, ONNX Runtime picks its own
+ * defaults, which is what every caller but the bi-encoder worker wants.
+ */
+interface OnnxSessionOptions {
+  intraOpNumThreads?: number;
+}
+
+export interface LoadOnnxModelOptions {
+  /** Tokenizer file names, for a repo that does not use the standard ones. */
+  tokenizerFile?: string;
+  tokenizerConfigFile?: string;
+  sessionOptions?: OnnxSessionOptions;
+}
+
 function resolveModelPath(modelRepo: string, hfRepoFile: string) {
   return path.resolve(SERVER_DIR, "models", modelRepo, hfRepoFile);
 }
@@ -44,13 +59,22 @@ async function ensureModelFileExists(modelRepo: string, hfRepoFile: string) {
 // quiet: ONNX Runtime otherwise warns that it assigned shape operators to CPU,
 // which is expected and not actionable. The architecture is logged because it
 // selects the quantized kernel, which is the part that varies between hosts.
-function createOnnxSession(modelRepo: string, modelPath: string) {
+function createOnnxSession(
+  modelRepo: string,
+  modelPath: string,
+  sessionOptions: OnnxSessionOptions = {},
+) {
   printMessage(
-    `Creating CPU session for ${modelRepo} (arch: ${process.arch}, platform: ${process.platform})...`,
+    `Creating CPU session for ${modelRepo} (arch: ${process.arch}, platform: ${process.platform}${
+      sessionOptions.intraOpNumThreads === undefined
+        ? ""
+        : `, intra-op threads: ${sessionOptions.intraOpNumThreads}`
+    })...`,
   );
   return InferenceSession.create(modelPath, {
     executionProviders: ["cpu"],
     logSeverityLevel: 3,
+    ...sessionOptions,
   });
 }
 
@@ -105,8 +129,11 @@ function resolvePadTokenId(
 export async function loadOnnxModel(
   modelRepo: string,
   modelFile: string,
-  tokenizerFile = "tokenizer.json",
-  tokenizerConfigFile = "tokenizer_config.json",
+  {
+    tokenizerFile = "tokenizer.json",
+    tokenizerConfigFile = "tokenizer_config.json",
+    sessionOptions = {},
+  }: LoadOnnxModelOptions = {},
 ): Promise<{
   session: InferenceSession;
   tokenizer: Tokenizer;
@@ -126,7 +153,7 @@ export async function loadOnnxModel(
     tokenizerConfig,
   );
   const padTokenId = resolvePadTokenId(tokenizer, tokenizerConfig);
-  const session = await createOnnxSession(modelRepo, modelPath);
+  const session = await createOnnxSession(modelRepo, modelPath, sessionOptions);
 
   return { session, tokenizer, padTokenId };
 }