Package: colbert-ai
ColBERT (v2)ColBERT is afastandaccurateretrieval model, enabling scalable BERT-based search over large text collections in tens of milliseconds.Figure 1:ColBERT's late interaction, efficiently scoring the fine-grained similarity between a queries and a passage.As Figure 1 illustrates, ColBERT relies on fine-grainedcontextual late interaction: it encodes each passage into amatrixof token-level embeddings (shown above in blue). Then at search time, it embeds every query into another matrix (shown in green) and efficiently finds passages that contextually match the query using scalable vector-similarity (MaxSim) operators.These rich interactions allow ColBERT to surpass the quality ofsingle-vectorrepresentation models, while scaling efficiently to large corpora. You can read more in our papers:ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT(SIGIR'20).Relevance-guided Supervision for OpenQA with ColBERT(TACL'21).Baleen: Robust Multi-Hop Reasoning at Scale via Condensed Retrieval(NeurIPS'21).ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction(NAACL'22).PLAID: An Efficient Engine for Late Interaction Retrieval(CIKM'22).Moving Beyond Downstream Task Accuracy for Information Retrieval Benchmarking(ACL'23 Findings).UDAPDR: Unsupervised Domain Adaptation via LLM Prompting and Distillation of Rerankers(EMNLP'23).🚨Announcements(1/28/24) One of the easiest ways to use ColBERT in applications nowadays is the semi-official, fast-growingRAGatouillelibrary.(1/29/23) We have merged a new index updater feature and support for additional Hugging Face models! These are in beta so please give us feedback as you try them out.(1/24/23) If you're looking for theDSPyframework for composing retrievers like ColBERTv2 and LLMs, it's at:https://github.com/stanfordnlp/dspyColBERTv1The ColBERTv1 code from the SIGIR'20 paper is in thecolbertv1branch. Seeherefor more information on other branches.Installation(Update: nowadays you can typically dopip install colbert-ai[torch,faiss-gpu]to get things up and running, but if you face issues conda is always more reliable forfaissandtorch.)ColBERT requires Python 3.7+ and Pytorch 1.9+ and uses theHugging Face Transformerslibrary.We strongly recommend creating a conda environment using the commands below. (If you don't have conda, follow the officialconda installation guide.)We have also included a new environment file specifically for CPU-only environments (conda_env_cpu.yml), but note that if you are testing CPU execution on a machine that includes GPUs you might need to specifyCUDA_VISIBLE_DEVICES=""as part of your command. Note that a GPU is required for training and indexing.conda env create -f conda_env[_cpu].yml
conda activate colbertIf you face any problems, pleaseopen a new issueand we'll help you promptly!OverviewUsing ColBERT on a dataset typically involves the following steps.Step 0: Preprocess your collection.At its simplest, ColBERT works with tab-separated (TSV) files: a file (e.g.,collection.tsv) will contain all passages and another (e.g.,queries.tsv) will contain a set of queries for searching the collection.Step 1: Download thepre-trained ColBERTv2 checkpoint.This checkpoint has been trained on the MS MARCO Passage Ranking task. You can alsooptionallytrain your own ColBERT model.Step 2: Index your collection.Once you have a trained ColBERT model, you need toindex your collectionto permit fast retrieval. This step encodes all passages into matrices, stores them on disk, and builds data structures for efficient search.Step 3: Search the collection with your queries.Given the model and index, you canissue queries over the collectionto retrieve the top-k passages for each query.Below, we illustrate these steps via an example run on the MS MARCO Passage Ranking task.API Usage NotebookNEW: We have an experimental notebook onGoogle Colabthat you can use with free GPUs. Indexing 10,000 on the free Colab T4 GPU takes six minutes.This Jupyter notebookdocs/intro.ipynb notebookillustrates using the key features of ColBERT with the new Python API.It includes how to download the ColBERTv2 model checkpoint trained on MS MARCO Passage Ranking and how to download our new LoTTE benchmark.DataThis repository works directly with a simpletab-separated fileformat to store queries, passages, and top-k ranked lists.Queries: each line isqid \t query text.Collection: each line ispid \t passage text.Top-k Ranking: each line isqid \t pid \t rank.This works directly with the data format of theMS MARCO Passage Rankingdataset. You will need the training triples (triples.train.small.tar.gz), the official top-1000 ranked lists for the dev set queries (top1000.dev), and the dev set relevant passages (qrels.dev.small.tsv). For indexing the full collection, you will also need the list of passages (collection.tar.gz).IndexingFor fast retrieval, indexing precomputes the ColBERT representations of passages.Example usage:fromcolbert.infraimportRun,RunConfig,ColBERTConfigfromcolbertimportIndexerif__name__=='__main__':withRun().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")RetrievalWe typically recommend that you use ColBERT forend-to-endretrieval, where it directly finds its top-k passages from the full collection:fromcolbert.dataimportQueriesfromcolbert.infraimportRun,RunConfig,ColBERTConfigfromcolbertimportSearcherif__name__=='__main__':withRun().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")You can optionally specify thencells,centroid_score_threshold, andndocssearch hyperparameters to trade off between speed and result quality. Defaults for different values ofkare listed in colbert/searcher.py.We can evaluate the MSMARCO rankings using the following command:python -m utility.evaluate.msmarco_passages --ranking "/path/to/msmarco.nbits=2.ranking.tsv" --qrels "/path/to/MSMARCO/qrels.dev.small.tsv"Basic Training (ColBERTv1-style)We provide apre-trained model checkpoint, but we also detail how to train from scratch here.
Note that this example demonstrates the ColBERTv1 style of training, but the provided checkpoint was trained with ColBERTv2.Training requires a JSONL triples file with a[qid, pid+, pid-]list per line. The query IDs and passage IDs correspond to the specifiedqueries.tsvandcollection.tsvfiles respectively.Example usage (training on 4 GPUs):fromcolbert.infraimportRun,RunConfig,ColBERTConfigfromcolbertimportTrainerif__name__=='__main__':withRun().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}...")Advanced Training (ColBERTv2-style)fromcolbert.infra.runimportRunfromcolbert.infra.configimportColBERTConfig,RunConfigfromcolbertimportTrainerdeftrain():# use 4 gpus (e.g. four A100s, but you can use fewer by changing nway,accumsteps,bsize).withRun().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()Running a lightweight ColBERTv2 serverWe provide a script to run a lightweight server which serves k (upto 100) results in ranked order for a given search query, in JSON format. This script can be used to power DSP programs.To run the server, update the environment variablesINDEX_ROOTandINDEX_NAMEin the.envfile to point to the appropriate ColBERT index. The run the following command:python server.pyA sample query:http://localhost:8893/api/search?query=Who won the 2022 FIFA world cup&k=25BranchesSupported branchesmain: Stable branch with ColBERTv2 + PLAID.colbertv1: Legacy branch for ColBERTv1.Deprecated branchesnew_api: Base ColBERTv2 implementation.cpu_inference: ColBERTv2 implementation with CPU search support.fast_search: ColBERTv2 implementation with PLAID.binarization: ColBERT with a baseline binarization-based compression strategy (as opposed to ColBERTv2's residual compression, which we found to be more robust).AcknowledgmentsColBERT logo designed by Chuyi Zhang.