litil-colbert-300m-v1-mlx

Litil ColBERT 300M v1 (MLX)

An MLX port of litillabs/litil-colbert-300m-v1 for Apple silicon. Same weights, same tokenizer, same 128-dimensional token vectors and MaxSim scoring. No PyTorch required.

The repository contains float32 weights at the top level and float16 weights in float16/.

Usage

import sys
from huggingface_hub import snapshot_download

path = snapshot_download(
    "litillabs/litil-colbert-300m-v1-mlx",
    allow_patterns=["*.py", "*.json", "model.safetensors"],  # float32 only
)
sys.path.insert(0, path)
from litil_colbert_mlx import LitilColBERT

model = LitilColBERT(path)

query = model.encode(["Kündigungsfrist bei befristetem Mietvertrag"], is_query=True)[0]
documents = model.encode(["<document 1>", "<document 2>"], is_query=False)

scores = [LitilColBERT.maxsim(query, d) for d in documents]

For the float16 weights, download with allow_patterns=["*.py", "*.json", "float16/*"] and load LitilColBERT(f"{path}/float16").

encode returns one array of normalized token vectors per text, shape (tokens, 128). Queries are read up to 1,024 tokens and documents up to 8,192. Requirements: mlx, numpy, tokenizers, and huggingface_hub for downloading.

Accuracy and speed

Compared against the PyTorch float32 implementation on 104 legal documents across six query sets in German, Chinese and English, with documents up to 8,192 tokens:

Single-text encoding latency on an Apple M5 Max, in milliseconds:

Use float32 to reproduce the published evaluation exactly; use float16 for faster indexing.

Precision note

MLX enables TF32 for float32 matrix multiplication on recent Apple GPUs by default. On this model TF32 shifts a small number of token vectors in long inputs by several percent. litil_colbert_mlx.py sets MLX_ENABLE_TF32=0 on import unless you have set it yourself, which restores exact float32 output. Import the module before running any other MLX GPU operation in the same process.

License

Apache-2.0, inherited from the original model.