MLX
mxfp4 (group_size=32, 4-bit) quantization of
LiquidAI/LFM2.5-ColBERT-350M,
a bidirectional LFM2.5 late-interaction (128-d per-token, MaxSim) for multilingual retrieval. Runs on Apple Silicon via MLX.
1# pip install mlx mlx-lm transformers
2# This repo bundles `mlx_lfm2_encoder.py` — the bidirectional LFM2 encoder
3# (CLS pooling / ColBERT MaxSim) that the stock causal LFM2 loaders do NOT provide.
4import mlx.core as mx
5from transformers import AutoTokenizer
6from mlx_lfm2_encoder import load_model
7
8tok = AutoTokenizer.from_pretrained(".", trust_remote_code=True)
9model, _ = load_model(".", head="colbert") # head: "embedding" or "colbert"
10
11# ColBERT late interaction: per-token 128-d vectors + MaxSim
12import numpy as np
13def encode(texts, prefix):
14 enc = tok([prefix + t for t in texts], return_tensors="np", padding=True,
15 truncation=True, max_length=512)
16 out = model(mx.array(enc["input_ids"]), mx.array(enc["attention_mask"]))
17 mx.eval(out); arr = np.array(out.astype(mx.float32))
18 mask = enc["attention_mask"].astype(bool)
19 return [arr[i][mask[i]] for i in range(arr.shape[0])] # list of (Li, 128)
20
21q = encode(["who wrote hamlet"], "[Q] ")[0]
22d = encode(["Hamlet is a tragedy written by William Shakespeare ..."], "[D] ")[0]
23maxsim = (q @ d.T).max(axis=1).sum() # late-interaction score
These are
bidirectional encoders (non-causal attention + non-causal short-conv +
per-token MaxSim). General-purpose causal LFM2 loaders produce
wrong embeddings here, so this repo ships
mlx_lfm2_encoder.py (validated to cosine ≥ 0.999
against the original
transformers model). For the broader MLX embedding ecosystem see
mlx-embeddings.
