repository·main·Indexed 25 days ago
https://github.com/stanford-futuredata/colbertA fast and accurate retrieval model using contextual late interaction for scalable BERT-based search. It provides tools for indexing text collections via the Indexer class, performing retrieval with the Searcher class, and training models using ColBERTv1 and ColBERTv2 styles. The library supports GPU acceleration via PyTorch and FAISS, and includes a lightweight server for serving top-k search results in JSON format.
ColBERT provides a script to run a lightweight server that serves top-k search results in JSON format.
.env file with INDEX_ROOT and INDEX_NAME pointing to your index.python server.pyhttp://localhost:8893/api/search?query=<QUERY_TEXT>&k=<K_VALUE>
python server.pyColBERT requires Python 3.7+ and Pytorch 1.9+. You can install it via pip or using Conda.
Pip installation:
pip install colbert-ai[torch,faiss-gpu]Conda installation (Recommended):
Create an environment using the provided .yml files. Use conda_env_cpu.yml for CPU-only environments. Note that a GPU is required for training and indexing.
conda env create -f conda_env[_cpu].yml
conda activate colbertIf testing CPU execution on a machine with GPUs, you may need to set CUDA_VISIBLE_DEVICES="".
conda env create -f conda_env[_cpu].yml
conda activate colbertTo install the ColBERT library with GPU support for FAISS and PyTorch, use the following command:
!pip install "colbert-ir[faiss-gpu, torch]"To use ColBERT for indexing and retrieval, import the Indexer, Searcher, Run, RunConfig, ColBERTConfig, Queries, and Collection classes.
import os
import sys
sys.path.insert(0, '../')
from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert.data import Queries, Collection
from colbert import Indexer, SearcherTo enable fast retrieval, you must index your collection. This process encodes passages into matrices and builds efficient data structures. You need a trained ColBERT model checkpoint and a collection in TSV format (pid \t passage text).
from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Indexer
if __name__=='__main__':
with Run().context(RunConfig(nranks=1, experiment="msmarco")):
config = ColBERTConfig(
nbits=2,
root="/path/to/experiments",
)
indexer = Indexer(checkpoint="/path/to/checkpoint", config=config)
indexer.index(name="msmarco.nbits=2", collection="/path/to/MSMARCO/collection.tsv")The trainer.Trainer class provides the interface for training ColBERT models. It supports configuring training parameters and executing the training loop.
Key methods include:
configure(): Sets up the training configuration.train(): Starts the training process.best_checkpoint_path(): Returns the file path to the best performing checkpoint discovered during training.Perform end-to-end retrieval to find the top-k passages for a set of queries. The queries should be in TSV format (qid \t query text).
from colbert.data import Queries
from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Searcher
if __name__=='__main__':
with Run().context(RunConfig(nranks=1, experiment="msmarco")):
config = ColBERTConfig(
root="/path/to/experiments",
)
searcher = Searcher(index="msmarco.nbits=2", config=config)
queries = Queries("/path/to/MSMARCO/queries.dev.small.tsv")
ranking = searcher.search_all(queries, k=100)
ranking.save("msmarco.nbits=2.ranking.tsv")Train a model from scratch using the ColBERTv1 style. This requires a JSONL triples file containing [qid, pid+, pid-] per line, along with queries.tsv and collection.tsv files.
from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Trainer
if __name__=='__main__':
with Run().context(RunConfig(nranks=4, experiment="msmarco")):
config = ColBERTConfig(
bsize=32,
root="/path/to/experiments",
)
trainer = Trainer(
triples="/path/to/MSMARCO/triples.train.small.tsv",
queries="/path/to/MSMARCO/queries.train.small.tsv",
collection="/path/to/MSMARCO/collection.tsv",
config=config,
)
checkpoint_path = trainer.train()
print(f"Saved checkpoint to {checkpoint_path}...")searcher.Searcher class is the primary interface for performing retrieval tasks using a ColBERT index. It provides methods for configuring the search environment, encoding queries, and executing both standard and dense searches against the indexed collection.Perform advanced training using the ColBERTv2 style. This method supports parameters like nway, accumsteps, and use_ib_negatives for more efficient and effective training.
from colbert.infra.run import Run
from colbert.infra.config import ColBERTConfig, RunConfig
from colbert import Trainer
def train():
# use 4 gpus (e.g. four A100s, but you can use fewer by changing nway,accumsteps,bsize).
with Run().context(RunConfig(nranks=4)):
triples = '/path/to/examples.64.json' # `wget https://huggingface.co/colbert-ir/colbertv2.0_msmarco_64way/resolve/main/examples.json?download=true` (26GB)
queries = '/path/to/MSMARCO/queries.train.tsv'
collection = '/path/to/MSMARCO/collection.tsv'
config = ColBERTConfig(bsize=32, lr=1e-05, warmup=20_000, doc_maxlen=180, dim=128, attend_to_mask_tokens=False, nway=64, accumsteps=1, similarity='cosine', use_ib_negatives=True)
trainer = Trainer(triples=triples, queries=queries, collection=collection, config=config)
trainer.train(checkpoint='colbert-ir/colbertv1.9') # or start from scratch, like `bert-base-uncased`
if __name__ == '__main__':
train()indexer.Indexer class is the primary interface for creating and managing ColBERT indices. It provides methods for configuring the indexing process, launching the indexing job, and managing the resulting index files.You can tune the performance of the Searcher by providing a ColBERTConfig object. This allows you to balance search speed against retrieval quality.
k (e.g., ncells=1, centroid_score_threshold=0.5, ndocs=256).ncells=4, centroid_score_threshold=0.4, ndocs=4096).