LFM2.5-ColBERT-350M is a late interaction retriever with best-in-class multilingual performance. It stores one vector per token and matches queries to documents with MaxSim, so you can store documents in one language (for example, a product description in English) and retrieve them in many languages with high accuracy.
Find more information about LFM2.5-ColBERT-350M in our
blog post.
Make requests to embed queries and documents, and compute MaxSim similarity scores
1❯ uv run colbert-rerank.py
2
3Score: 29.04 | Q: What is panda? | D: hi
4Score: 29.57 | Q: What is panda? | D: it is a bear
5Score: 30.07 | Q: What is panda? | D: The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.
1# /// script
2# requires-python = ">=3.10"
3# dependencies = [
4# "transformers",
5# "huggingface-hub",
6# "numpy",
7# "requests",
8# "torch",
9# ]
10# ///
11
12# colbert-rerank.py
13from transformers import AutoTokenizer
14from huggingface_hub import hf_hub_download
15import numpy as np, requests, torch, torch.nn.functional as F, json
16
17
18model_id = "LiquidAI/LFM2.5-ColBERT-350M"
19tokenizer = AutoTokenizer.from_pretrained(model_id)
20config = json.load(open(hf_hub_download(model_id, "config_sentence_transformers.json")))
21skiplist = set(
22 t
23 for w in config["skiplist_words"]
24 for t in tokenizer.encode(w, add_special_tokens=False)
25)
26
27
28def maxsim(q, d):
29 return (q @ d.T).max(dim=1).values.sum().item()
30
31
32def preprocess(text, is_query):
33 prefix = config["query_prefix"] if is_query else config["document_prefix"]
34 toks = tokenizer.encode(prefix + text)
35 max_len = config["query_length"] if is_query else config["document_length"]
36 if is_query:
37 toks += [tokenizer.pad_token_id] * (max_len - len(toks))
38 else:
39 toks = toks[:max_len]
40 mask = None if is_query else [t not in skiplist for t in toks]
41 return toks, mask
42
43
44def embed(content, mask=None):
45 emb = np.array(
46 requests.post(
47 "http://localhost:8080/embedding",
48 json={"content": content},
49 ).json()[0]["embedding"]
50 )
51 if mask:
52 emb = emb[mask]
53 emb = torch.from_numpy(emb)
54 emb = F.normalize(emb, p=2, dim=-1) # L2 normalize each token embedding
55 return emb.unsqueeze(0)
56
57
58docs = [
59 "hi",
60 "it is a bear",
61 "The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.",
62]
63query = "What is panda?"
64
65q = embed(*preprocess(query, True))
66d = [embed(*preprocess(doc, False)) for doc in docs]
67s = [(query, doc, maxsim(q.squeeze(), di.squeeze())) for doc, di in zip(docs, d)]
68for q_text, d_text, score in s:
69 print(f"Score: {score:.2f} | Q: {q_text} | D: {d_text}")