LFM2-ColBERT-350M

pip install -U pylate

from pylate import indexes, models, retrieve

# Load the ColBERT model
model = models.ColBERT(
    model_name_or_path="LiquidAI/LFM2-ColBERT-350M",
)
model.tokenizer.pad_token = model.tokenizer.eos_token

# Initialize the PLAID index
index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
    override=True,
)

# Encode and index documents
documents_ids = ["1", "2", "3"]
documents = ["document 1 text", "document 2 text", "document 3 text"]

documents_embeddings = model.encode(
    documents,
    batch_size=32,
    is_query=False,
    show_progress_bar=True,
)

index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

# Initialize retriever
retriever = retrieve.ColBERT(index=index)

# Encode queries
queries_embeddings = model.encode(
    ["query for document 3", "query for document 1"],
    batch_size=32,
    is_query=True,
    show_progress_bar=True,
)

# Retrieve top-k documents
scores = retriever.retrieve(
    queries_embeddings=queries_embeddings,
    k=10,
)

from pylate import rank, models

queries = ["query A", "query B"]
documents = [
    ["document A", "document B"],
    ["document 1", "document C", "document B"],
]
documents_ids = [[1, 2], [1, 3, 2]]

model = models.ColBERT(
    model_name_or_path="LiquidAI/LFM2-ColBERT-350M",
)

queries_embeddings = model.encode(queries, is_query=True)
documents_embeddings = model.encode(documents, is_query=False)

reranked_documents = rank.rerank(
    documents_ids=documents_ids,
    queries_embeddings=queries_embeddings,
    documents_embeddings=documents_embeddings,
)