server/biEncoderService.ts

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";

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. Fast enough for batched passage encoding: a batch of 64 passages
 * takes ~30 ms on CPU, well within the 20 s client timeout even for six pages
 * with hundreds of passages each.
 */
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;

/** Batch size for passage encoding. Larger batches speed up encoding but use
 * more memory; 64 is a safe default on CPU. */
const BATCH_SIZE = 64;

let isReady = false;
let session: InferenceSession | null = null;
let tokenizer: Tokenizer | null = null;

/**
 * Encodes a single text into a normalized embedding vector.
 */
async function encode(
  activeSession: InferenceSession,
  loadedTokenizer: Tokenizer,
  text: string,
): Promise<Float32Array> {
  const { ids, attention_mask } = loadedTokenizer.encode(text);
  const truncatedIds = ids.slice(0, MAX_SEQUENCE_LENGTH);
  const truncatedMask = attention_mask.slice(0, MAX_SEQUENCE_LENGTH);

  const length = truncatedIds.length;
  const dimensions = [1, length];

  const { last_hidden_state } = await activeSession.run({
    input_ids: new Tensor(
      "int64",
      BigInt64Array.from(truncatedIds, BigInt),
      dimensions,
    ),
    attention_mask: new Tensor(
      "int64",
      BigInt64Array.from(truncatedMask, BigInt),
      dimensions,
    ),
    // The export declares `token_type_ids` and ONNX Runtime refuses to run with
    // a declared input missing. A single text is one segment, so it is zeros.
    token_type_ids: new Tensor("int64", new BigInt64Array(length), dimensions),
  });

  // Mean pooling: average the hidden states across non-padded tokens.
  const embedding = last_hidden_state.data as Float32Array;
  const dim = last_hidden_state.dims[2];
  const pooled = new Float32Array(dim);
  let count = 0;

  for (let t = 0; t < length; t++) {
    if (truncatedMask[t] === 0) continue;
    const offset = t * dim;
    for (let d = 0; d < dim; d++) {
      pooled[d] += embedding[offset + d];
    }
    count++;
  }

  if (count > 0) {
    for (let d = 0; d < dim; d++) {
      pooled[d] /= count;
    }
  }

  // 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;
    }
  }

  return pooled;
}

/**
 * Encodes a batch of texts into normalized embedding vectors.
 */
async function encodeBatch(
  activeSession: InferenceSession,
  loadedTokenizer: Tokenizer,
  texts: string[],
): Promise<Float32Array[]> {
  const results: Float32Array[] = [];

  for (let i = 0; i < texts.length; i += BATCH_SIZE) {
    const batch = texts.slice(i, i + BATCH_SIZE);
    const batchResults = await Promise.all(
      batch.map((text) => encode(activeSession, loadedTokenizer, text)),
    );
    results.push(...batchResults);
  }

  return results;
}

/**
 * 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;
  });
}

export async function startBiEncoderService() {
  printMessage("Preparing bi-encoder model...");
  const loaded = await loadOnnxModel(MODEL_HF_REPO, MODEL_HF_FILE);
  session = loaded.session;
  tokenizer = loaded.tokenizer;

  // Warm up with a test encoding.
  await encode(session, tokenizer, "test query");

  isReady = true;
  printMessage("Bi-encoder service ready!");
}

export async function stopBiEncoderService() {
  isReady = false;
  const currentSession = session;
  session = null;
  tokenizer = null;
  await currentSession?.release();
}

export async function getBiEncoderStatus() {
  return isReady;
}

/**
 * Returns dense (semantic) scores for passages given a query.
 * Falls back to empty array when the model is not loaded.
 */
export async function scorePassages(
  query: string,
  passages: string[],
): Promise<number[]> {
  if (!session || !tokenizer || passages.length === 0) {
    return [];
  }

  const activeSession = session;
  const loadedTokenizer = tokenizer;

  const queryEmbedding = await encode(activeSession, loadedTokenizer, query);
  const passageEmbeddings = await encodeBatch(
    activeSession,
    loadedTokenizer,
    passages,
  );

  return cosineSimilarities(queryEmbedding, passageEmbeddings);
}