utils

Encoders that turn text and images into vectors for semantic search, from fast static models to ONNX transformers, plus a fallback that needs no download

Encoders and embeddings

Keyword search finds rows that share words with the query. Semantic search finds rows that mean something close to it, even with no word in common: “car” can find a passage about automobiles. It works by turning each piece of text into an embedding, a fixed-length list of numbers (a vector) placed so that texts with similar meanings sit close together. The model that does this is the encoder. At indexing time the encoder embeds every chunk and the vectors are stored; at query time it embeds the query, and the search returns the chunks whose vectors are nearest.

A query and the documents it searches must be embedded by the same encoder. Vectors from two different models live in unrelated spaces, so comparing them gives meaningless distances without any error. This is why a store records which encoder wrote it (see encoder_id below).

litesearch’s default encoder is static_embedder(), which loads model2vec’s potion-multilingual-128M. A model2vec model is static: it holds one precomputed vector per token and embeds a text by averaging the vectors of its tokens. There is no neural network to run, so it needs no GPU and no ONNX runtime, and it embedded about 8,600 chunks of 500 characters a second on an Apple-silicon Mac. It returns 256-dimensional float32 vectors and covers many scripts, Devanagari included. Index uses it unless you pass another encoder. static_code_embedder() loads potion-code-16M-v2, a model2vec model trained on source code; nothing in litesearch uses it by default, so pass it yourself if you index code. Both functions cache the model per name, so every caller shares one copy. The multilingual model is about 500 MB.

evals/ found that the choice of encoder moves retrieval quality less than chunk size or reranking do. Across four encoders the spread was 0.018 to 0.046 weighted MRR, and the static model won one of the three test corpora outright.

When a transformer is worth the extra time, FastEncode runs one through ONNX Runtime. The dicts in the next cell name models that have been tried with it; each gives the Hugging Face repository, the ONNX file inside it and, where the model needs them, prompt templates. You can pass any of these, any other model2vec model name to static_embedder, or your own object with an .encode method.


source

static_embedder

def static_embedder(
    model_nm:str='minishlab/potion-multilingual-128M'
):

The model2vec model model_nm, loaded once and shared by every caller.


source

ort_path

def ort_path(
    p
):

p, or hard links to it and its external data files in one real folder when p is a symlink into the Hugging Face cache.


source

download_model

def download_model(
    repo_id:str='onnx-community/embeddinggemma-300m-ONNX', # HF repo id
    md:NoneType=None, # optional explicit local dir; else the HF cache is used
    filename:NoneType=None, # one file inside the repo; None downloads the whole repo
    token:NoneType=None, # Hugging Face token; defaults to the HF_TOKEN environment variable
):

The local path of a Hugging Face model repo, or one file in it, downloading only what is not already cached.

FastEncode: transformer models through ONNX

FastEncode runs a text embedding model exported to ONNX, a portable format for neural networks, with ONNX Runtime. It works with most ONNX models that ship a Hugging Face tokenizer. On first use it downloads the model with download_model, which reuses the local Hugging Face cache when the files are already there. The default model is EmbeddingGemma 300M, which gives 768-dimensional vectors.

Each text is split into tokens, cut off at the smaller of the model’s own limit and 512 tokens (max_seq_len replaces the 512), and run through the model. The per-token outputs are averaged into one vector, scaled to length 1 when normalize=True, and returned as float16 by default. Texts are sent in batches of batch_size, parallel runs batches on that many threads, and stream=True returns the batches one by one instead of one joined array. quantize='int8' or 'uint8' writes a smaller quantized copy of the model beside the original on first use and runs that instead.

Many embedding models are trained to see a short prefix that says whether a text is a document or a query, such as search_query: for nomic models. encode_document and encode_query add the prefix from the model’s prompt templates; encode adds none.

FastEncode needs onnxruntime, which is not installed with litesearch (pip install onnxruntime). If the model fails to download or load, the constructor prints the error instead of raising, and every later encode call prints a message and returns None. Check the result before you store it. Qwen embedding models handle their padding token differently and do not work with FastEncode yet.


source

FastEncode

def FastEncode(
    model_dict:fastcore.basics.AttrDict={'model': 'onnx-community/embeddinggemma-300m-ONNX', 'onnx_path': 'onnx/model.onnx', 'prompt': {'document': 'title: none | text: {text}', 'query': 'task: search result | query: {text}'}}, # a model dict above: repo, ONNX file path and prompt templates
    repo_id:NoneType=None, # Hugging Face repo with an ONNX file, used instead of `model_dict`
    md:NoneType=None, # local folder holding the model files
    md_nm:NoneType=None, # path of the ONNX file inside the model folder
    normalize:bool=True, # scale each vector to length 1
    dtype:type=float16, # output dtype
    tti:bool=False, # pass token type ids to the model, as BERT-style models expect; also set by `model_dict.tti`
    prompt:NoneType=None, # dict with `document` and `query` templates containing `{text}`
    hf_token:NoneType=None, # Hugging Face token; defaults to the HF_TOKEN environment variable
    batch_size:int=32, # texts per ONNX call
    parallel:int=0, # threads that encode batches at once; 0 encodes in the calling thread
    quantize:NoneType=None, # run a quantized copy of the model: None, 'int8' or 'uint8'
    max_seq_len:NoneType=None, # token limit per text in place of 512; the model's own limit still applies (e.g. 64 for SigLIP2)
):

A text encoder that runs a Hugging Face ONNX model with ONNX Runtime.

# download_model resolves to a local dir with the model files; repeated calls are network-free
p = download_model(embedding_gemma.model)
assert Path(p).exists() and (Path(p)/'config.json').exists(), p
assert Path(download_model(embedding_gemma.model)).samefile(p)  # idempotent
assert Path(_cached_snapshot(embedding_gemma.model)).samefile(p)  # scan_cache_dir hit, no network
assert _cached_snapshot('does-not/exist-xyz') is None           # missing repo -> None, never raises
print('download_model cache test passed:', p)
/Users/71293/code/litesearch/.venv/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
download_model cache test passed: /Users/71293/.cache/huggingface/hub/models--onnx-community--embeddinggemma-300m-ONNX/snapshots/5090578d9565bb06545b4552f76e6bc2c93e4a66
# a cached model's external data sits in a different blob folder; ort_path puts both in one real folder
q = ort_path(Path(p)/embedding_gemma.onnx_path)
assert q.resolve().parent == Path(f'{q}_data').resolve().parent, q
assert ort_path(Path(p)/embedding_gemma.onnx_path) == q  # reuses the links
assert ort_path(Path(p)/'config.json') == Path(p)/'config.json'  # nothing beside it to link

source

query_encoder

def query_encoder(
    embedder:FastEncode|StaticModel
):

A function that embeds a list of queries with embedder, using its query prompt when it has one.


source

doc_encoder

def doc_encoder(
    embedder:FastEncode|StaticModel
):

A function that embeds a list of documents with embedder, using its document prompt when it has one.

doc_encoder and query_encoder give every encoder the same shape: a function that takes a list of strings and returns their vectors. For a FastEncode they call encode_document or encode_query, so the model’s prompts are added. For anything else, a model2vec model or your own object, both call .encode and the two functions behave the same. Index builds both from the encoder you give it.

encoder_id names the model behind an encoder object, for example minishlab/potion-multilingual-128M, with +int8 appended for a quantized FastEncode. A string is returned as it is, and an encoder it cannot name gives ''. Index passes the name to get_store(encoder=...), which records it with the store the first time and raises StoreMismatch if the store is later opened with a different one. Without that check, adding vectors from a second model to an existing store would return wrong neighbours with no error. An empty name never conflicts, so a custom encoder with no name is not checked.


source

encoder_id

def encoder_id(
    e
)->str:

The model name of an encoder, such as its Hugging Face repo, or ’’ when it has none.

Image encoding

FastEncodeImage embeds images the way FastEncode embeds text, so you can find pictures that look alike: store a vector per image, embed a query image, and search. It takes Pillow images, file paths or image bytes. Each image is converted to RGB, resized to a square of the model’s input size (224 pixels by default) without cropping, normalised with the model’s mean and standard deviation, and run through the ONNX model. Where the model returns one output per image patch, those are averaged into one vector. The default model is nomic-embed-vision v1.5, with 768-dimensional float16 vectors. batch_size, parallel and stream work as in FastEncode, and the method is embed rather than encode. Like FastEncode, it needs onnxruntime and returns None from embed when the model failed to load.


source

FastEncodeImage

def FastEncodeImage(
    model_dict:fastcore.basics.AttrDict={'model': 'nomic-ai/nomic-embed-vision-v1.5', 'onnx_path': 'onnx/model.onnx', 'img_size': 224, 'mean': [0.485, 0.456, 0.406], 'std': [0.229, 0.224, 0.225]}, # a model dict above: repo, ONNX file path, input size, mean and std
    normalize:bool=True, # scale each vector to length 1
    dtype:type=float16, # output dtype
    batch_size:int=16, # images per ONNX call
    parallel:int=0, # threads that encode batches at once; 0 encodes in the calling thread
    hf_token:NoneType=None, # Hugging Face token; defaults to the HF_TOKEN environment variable
):

An image encoder that runs a CLIP-style Hugging Face ONNX model with ONNX Runtime.

Text and images in one space

To search images with a text query, the text and image vectors have to come from models trained together, so that a caption and its picture land close to each other. There are two ways to get that here.

Some repositories ship a text model and an image model together, such as SigLIP2. FastEncodeMultimodal loads both from one repository and gives you encode_text and encode_image. The default, siglip2_so400m, reads images at 512 pixels and cuts text at 64 tokens.

Other models come as a matched pair in separate repositories. nomic-embed-text v1.5 and nomic-embed-vision v1.5 both produce 768-dimensional vectors in the same space, so load FastEncode(nomic_text_v15) and FastEncodeImage(nomic_vision_v15) separately and compare their vectors directly.


source

FastEncodeMultimodal

def FastEncodeMultimodal(
    model_dict:fastcore.basics.AttrDict={'model': 'onnx-community/siglip2-so400m-patch16-512-ONNX', 'vision_onnx': 'onnx/vision_model.onnx', 'text_onnx': 'onnx/text_model.onnx', 'img_size': 512, 'mean': [0.5, 0.5, 0.5], 'std': [0.5, 0.5, 0.5], 'max_seq_len': 64}, # a model dict with `vision_onnx` and `text_onnx` file paths
    normalize:bool=True, # scale each vector to length 1
    dtype:type=float16, # output dtype
    batch_size_text:int=32, # texts per ONNX call
    batch_size_image:int=16, # images per ONNX call
    parallel:int=0, # threads that encode batches at once; 0 encodes in the calling thread
    hf_token:NoneType=None, # Hugging Face token; defaults to the HF_TOKEN environment variable
):

A text and image encoder pair loaded from one Hugging Face repo, such as SigLIP2.

from PIL import Image
import base64, io, os
from IPython.display import display
# Create 4 synthetic RGB images without needing any files
test_imgs = [Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)) for _ in range(4)]
img_enc = FastEncodeImage()  # default: nomic_vision_v15 (768-dim)

# Basic embed
embs = img_enc.embed(test_imgs)
assert embs.shape == (4, 768), f'Expected (4, 768), got {embs.shape}'
assert embs.dtype == np.float16
print(f'Basic image embed: shape={embs.shape}, dtype={embs.dtype}')

# Batching — batch_size=2 vs batch_size=4 must match
embs_b2 = img_enc.embed(test_imgs, batch_size=2)
embs_b4 = img_enc.embed(test_imgs, batch_size=4)
assert embs_b2.shape == embs_b4.shape == (4, 768)
assert np.allclose(embs_b2.astype(np.float32), embs_b4.astype(np.float32), atol=1e-3), 'image batching results mismatch'
print('Image batching test passed')

# Stream
stream_res = img_enc.embed(test_imgs, batch_size=2, stream=True)
assert hasattr(stream_res, '__iter__')
assert np.concatenate(list(stream_res)).shape == (4, 768)
print('Image stream test passed')

# Parallel
embs_par = img_enc.embed(test_imgs, batch_size=2, parallel=2)
assert np.allclose(embs_b2.astype(np.float32), embs_par.astype(np.float32), atol=1e-3), 'image parallel results mismatch'
print('Image parallel test passed')
Basic image embed: shape=(4, 768), dtype=float16
Image batching test passed
Image stream test passed
Image parallel test passed
# nomic paired — text (FastEncode) + vision (FastEncodeImage) share the same 768-dim space
nomic_enc     = FastEncode(nomic_text_v15, parallel=2)
nomic_img_enc = FastEncodeImage(nomic_vision_v15)

t_emb = nomic_enc.encode(['What does the Eiffel Tower look like?','What does the Eiffel Tower look like?','What does the Eiffel Tower look like?'])
i_emb = nomic_img_enc.embed([Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))])

assert t_emb.shape == (3, 768), f'Expected (3, 768), got {t_emb.shape}'
assert i_emb.shape == (1, 768), f'Expected (1, 768), got {i_emb.shape}'
assert t_emb.dtype == i_emb.dtype == np.float16
print(f'nomic paired: text={t_emb.shape}, image={i_emb.shape} — same space, ready for cross-modal search')
nomic paired: text=(3, 768), image=(1, 768) — same space, ready for cross-modal search

Embedding a PDF without a store

encode_pdf_texts and encode_pdf_images embed a PDF and hand back plain tuples, leaving the table layout to you. Use them when you keep vectors in your own schema instead of a litesearch store.

encode_pdf_texts chunks the PDF with pdf_chunks, or with your own chunk_fn, embeds the chunks with the encoder’s encode_document, and returns (page, chunk_idx, text, embedding) for each, with pages counted from 1. encode_pdf_images saves every PNG and JPEG image in the PDF to a temporary folder, embeds them with an image encoder, and returns (page, image_bytes, embedding). Its pages are counted from 0, unlike encode_pdf_texts. Both return lists.


source

encode_pdf_images

def encode_pdf_images(
    doc, # PdfDocument
    enc, # an image encoder with `embed`, such as FastEncodeImage
)->fastcore.foundation.L:

Every PNG and JPEG image in a PDF with its embedding, as (page, img_bytes, emb) tuples with pages counted from 0.


source

encode_pdf_texts

def encode_pdf_texts(
    doc, # PdfDocument
    enc, # a text encoder with `encode_document`, such as FastEncode
    chunk_fn:NoneType=None, # function from the document to [(page, chunk_idx, text)]; default doc.pdf_chunks
    **kw
)->fastcore.foundation.L: # passed to doc.pdf_chunks, which accepts only `chunker`

Every text chunk of a PDF with its embedding, as (page, chunk_idx, text, emb) tuples.

enc = FastEncode(parallel=2)
embs = enc.encode_document(['This is a test', 'Another test']*32, batch_size=4)
assert embs.shape == (64, 768), f'Expected (2, 768), got {embs.shape}'
assert embs.dtype == np.float16
print(f'Basic encode: shape={embs.shape}, dtype={embs.dtype}')
embs
Basic encode: shape=(64, 768), dtype=float16
array([[ 0.05777 ,  0.001723,  0.002573, ..., -0.0618  , -0.00662 ,
         0.03174 ],
       [ 0.02936 , -0.00818 , -0.00916 , ..., -0.02847 , -0.00226 ,
         0.02846 ],
       [ 0.05777 ,  0.001723,  0.002573, ..., -0.0618  , -0.00662 ,
         0.03174 ],
       ...,
       [ 0.02936 , -0.00818 , -0.00916 , ..., -0.02847 , -0.00226 ,
         0.02846 ],
       [ 0.05777 ,  0.001723,  0.002573, ..., -0.0618  , -0.00662 ,
         0.03174 ],
       [ 0.02936 , -0.00818 , -0.00916 , ..., -0.02847 , -0.00226 ,
         0.02846 ]], shape=(64, 768), dtype=float16)
enc_q = FastEncode(AttrDict(embedding_gemma, onnx_path='onnx/model_q4.onnx'))
assert embedding_gemma.onnx_path == 'onnx/model.onnx'  # the shared model dict is untouched
embs = enc_q.encode_document(['This is a test', 'Another test'])
assert embs.shape == (2, 768), f'Expected (2, 768), got {embs.shape}'
assert embs.dtype == np.float16
print(f'q4 encode: shape={embs.shape}, dtype={embs.dtype}')
embs
Basic encode: shape=(2, 768), dtype=float16
array([[ 0.05777 ,  0.001723,  0.002573, ..., -0.0618  , -0.00662 ,
         0.03174 ],
       [ 0.02936 , -0.00818 , -0.00916 , ..., -0.02847 , -0.00226 ,
         0.02846 ]], shape=(2, 768), dtype=float16)
# Batching — 30 texts in batch_size=8 vs batch_size=100 must produce identical results
texts = [f'sample text number {i}' for i in range(30)]
embs_b8   = enc.encode_document(texts, batch_size=8)
embs_b100 = enc.encode_document(texts, batch_size=100)
assert embs_b8.shape == embs_b100.shape == (30, 768)
assert np.allclose(embs_b8.astype(np.float32), embs_b100.astype(np.float32), atol=1e-3), 'batching results mismatch'
print(f'Batching test passed: {embs_b8.shape}')
Batching test passed: (30, 768)
# Stream — stream=True returns an iterable of batch arrays rather than one concatenated array
stream_results = enc.encode_document(texts, batch_size=10, stream=True)
assert hasattr(stream_results, '__iter__'), 'stream should be iterable'
concatenated = np.concatenate(list(stream_results))
assert concatenated.shape == (30, 768)
print('Stream test passed')

# Parallel — 4 thread workers should give same embeddings as single-threaded
embs_par = enc.encode_document(texts, batch_size=8, parallel=4)
assert np.allclose(embs_b8.astype(np.float32), embs_par.astype(np.float32), atol=1e-3), 'parallel results mismatch'
print('Parallel test passed')
Stream test passed
Parallel test passed
# Semantic QA — modernbert should rank the transformer/attention passage highest
modern_enc = FastEncode(modernbert)
passages = [
    'The Transformer is a model architecture using self-attention mechanisms.',
    'Recurrent neural networks process sequences step by step.',
    'The model uses encoder and decoder stacks connected by attention.',
    'Convolutional networks excel at image recognition tasks.',
]
doc_embs = modern_enc.encode_document(passages).astype(np.float32)
query = 'What architecture replaces recurrence with self-attention?'
q_emb = modern_enc.encode_query([query]).astype(np.float32)
sims = (doc_embs @ q_emb.T).flatten()
best = int(sims.argmax())
assert sims[best] > 0.3, f'Expected similarity > 0.3, got {sims[best]:.3f}'
assert any(w in passages[best].lower() for w in ['attention', 'transformer', 'encoder'])
print(f'Semantic QA test passed: best match sim={sims[best]:.3f}')
print(f'Best: {passages[best]}')
Semantic QA test passed: best match sim=0.527
Best: The Transformer is a model architecture using self-attention mechanisms.

Embedding with no model

hash_embed makes vectors without any model or download, and gives the same vectors on every run and every machine. It lowercases each text, takes every run of 3 to 5 consecutive characters (character n-grams), hashes each run with CRC32 into one of 256 slots, counts the hits per slot, and scales the result to length 1. Two texts score high when they share many character runs, so the vectors measure spelling, not meaning: “rank fusion” is close to “rank fusions”, but “cat” and “feline” score near zero.

It has two uses. One is matching names, such as people or places, where shared spelling is most of the signal anyway. The other is an index built where no real encoder can be downloaded, for example in a sandbox without network access.


source

hash_embed

def hash_embed(
    texts, # strings to embed
    ndim:int=256, # length of each vector
    ngram:tuple=(3, 5), # shortest and longest character run to hash
    dtype:type=float16
):

Unit-length vectors from hashed character n-grams, which match spelling rather than meaning and need no model.

from litesearch.utils import hash_embed
v = hash_embed(['rank fusion', 'rank fusion', 'a wholly unrelated sentence']).astype(np.float32)
assert v.shape == (3, 256)
assert (v[0] == v[1]).all()                 # same string, same vector
assert v[0] @ v[1] > v[0] @ v[2]            # and a shared spelling scores above an unshared one
assert abs(np.linalg.norm(v[0]) - 1) < 1e-2 # unit length, so a dot product is a cosine
# Devanagari is n-grammed as characters, so a codepoint is never cut in half
assert hash_embed(['सत्यमेव जयते', 'सत्यमेव जयते'])[0].tolist() == hash_embed(['सत्यमेव जयते'])[0].tolist()