ColBERT Retrieval

A Text Embeddings model pools a whole text into one vector. A ColBERT model instead keeps one vector per token and scores a query against a document with MaxSim: for every query token, the similarity of the closest document token, summed over the query. Matching at the token level makes it more accurate on retrieval than a single pooled vector, at the cost of a larger index (one 128-dimensional vector per document token).

Quick Start

The useColbertEmbedder hook manages model downloading, tokenizer loading, and lifecycle. Embed every document once with 'document', embed each search with 'query', and rank with maxSim():

import { models, useColbertEmbedder, type ColbertTokenEmbeddings } from 'react-native-executorch';

function Search({ documents }: { documents: string[] }) {
  const colbert = useColbertEmbedder(models.colbertEmbeddings.LFM2_5_COLBERT_350M.DEFAULT);

  const search = async (query: string) => {
    const { embed, maxSim } = colbert;
    if (!embed || !maxSim) return;

    // Index once and keep the result, e.g. in state or a database.
    const index: ColbertTokenEmbeddings[] = [];
    for (const doc of documents) index.push(await embed(doc, 'document'));

    const q = await embed(query, 'query');
    return documents
      .map((text, i) => ({ text, score: maxSim(q, index[i]!) }))
      .sort((a, b) => b.score - a.score);
  };

  // ...
}

Output Format

embed() returns a ColbertTokenEmbeddings object: a row-major [numTokens, dimension] matrix of L2-normalized vectors in a single Float32Array.

• A query always yields exactly 32 vectors. Shorter queries are padded with expansion tokens whose vectors take part in scoring (ColBERT's query augmentation).
• A document yields one vector per token, up to 512 tokens, with punctuation vectors dropped.

MaxSim scores are not bounded to [0, 1]: each query vector contributes a cosine similarity, so the score grows with the query length. Compare scores against the same query, not across queries.

Imperative API

For batch indexing or manual lifecycle management, use createColbertEmbedder:

import { createColbertEmbedder, download, models } from 'react-native-executorch';

const model = await download(models.colbertEmbeddings.LFM2_5_COLBERT_350M.DEFAULT);
const colbert = await createColbertEmbedder(model);

try {
  const doc = await colbert.embed('Canberra is the capital of Australia.', 'document');
  const query = await colbert.embed('What is the capital of Australia?', 'query');
  console.log('Relevance:', colbert.maxSim(query, doc));
} finally {
  colbert.dispose();
}

A synchronous embedWorklet is available for worklet runtimes, see Worklets & Threading.

Available Models

Model Family	Variants	Vector Dim	Languages	Size Range	Supported Backends
Liquid LFM 2.5 ColBERT 350M	See	128	11	188 MB – 373 MB	XNNPACK (CPU), MLX (Apple)

Using Custom Models

Any ColBERT model exported with a [1, S] token-id and attention-mask input and a [1, S, D] per-token output can be used, as long as its inputs are [BOS, prefix, ...tokens] with no trailing special token. Pass the settings from its PyLate config_sentence_transformers.json as ColbertEmbedderOptions:

const colbert = await createColbertEmbedder({
  modelPath: 'https://example.com/my-colbert.pte',
  tokenizerPath: 'https://example.com/tokenizer.json',
  modelOpts: {
    queryPrefixToken: '[Q] ',
    documentPrefixToken: '[D] ',
    queryLength: 32,
    queryExpansionToken: '<pad>',
    bosToken: '<s>',
    skiplistTokens: ['.', ',', '!', '?'],
  },
});

API Reference

• useColbertEmbedder() — React hook for ColBERT model downloading, state, and lifecycle.
• createColbertEmbedder() — Imperative factory for ColBERT pipelines.
• ColbertEmbedder, ColbertEmbedderModel, ColbertTokenEmbeddings — Types.
• models.colbertEmbeddings — Pre-configured ColBERT models.

Source Code

View the implementation on GitHub:

• src/extensions/nlp/tasks/colbertEmbedding.ts ↗