Crystalcareai/Colbertv2

1---2license: mit3language:4- en5tags:6- ColBERT7---8<p align="center">9  <img align="center" src="docs/images/colbertofficial.png" width="430px" />10</p>11<p align="left">12 13# ColBERT (v2)14 15### ColBERT is a _fast_ and _accurate_ retrieval model, enabling scalable BERT-based search over large text collections in tens of milliseconds.16 17[<img align="center" src="https://colab.research.google.com/assets/colab-badge.svg" />](https://colab.research.google.com/github/stanford-futuredata/ColBERT/blob/main/docs/intro2new.ipynb)18 19 20<p align="center">21  <img align="center" src="docs/images/ColBERT-Framework-MaxSim-W370px.png" />22</p>23<p align="center">24  <b>Figure 1:</b> ColBERT's late interaction, efficiently scoring the fine-grained similarity between a queries and a passage.25</p>26 27As Figure 1 illustrates, ColBERT relies on fine-grained **contextual late interaction**: it encodes each passage into a **matrix** of 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.28 29These rich interactions allow ColBERT to surpass the quality of _single-vector_ representation models, while scaling efficiently to large corpora. You can read more in our papers:30 31* [**ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT**](https://arxiv.org/abs/2004.12832) (SIGIR'20).32* [**Relevance-guided Supervision for OpenQA with ColBERT**](https://arxiv.org/abs/2007.00814) (TACL'21).33* [**Baleen: Robust Multi-Hop Reasoning at Scale via Condensed Retrieval**](https://arxiv.org/abs/2101.00436) (NeurIPS'21).34* [**ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction**](https://arxiv.org/abs/2112.01488) (NAACL'22).35* [**PLAID: An Efficient Engine for Late Interaction Retrieval**](https://arxiv.org/abs/2205.09707) (CIKM'22).36 37----38 39## 🚨 **Announcements** 40 41* (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.42* (1/24/23) If you're looking for the **DSP** framework for composing ColBERTv2 and LLMs, it's at: https://github.com/stanfordnlp/dsp43 44----45 46## ColBERTv147 48The ColBERTv1 code from the SIGIR'20 paper is in the [`colbertv1` branch](https://github.com/stanford-futuredata/ColBERT/tree/colbertv1). See [here](#branches) for more information on other branches.49 50 51## Installation52 53ColBERT requires Python 3.7+ and Pytorch 1.9+ and uses the [Hugging Face Transformers](https://github.com/huggingface/transformers) library.54 55We strongly recommend creating a conda environment using the commands below. (If you don't have conda, follow the official [conda installation guide](https://docs.anaconda.com/anaconda/install/linux/#installation).)56 57We 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 specify `CUDA_VISIBLE_DEVICES=""` as part of your command. Note that a GPU is required for training and indexing.58 59```60conda env create -f conda_env[_cpu].yml61conda activate colbert62```63 64If you face any problems, please [open a new issue](https://github.com/stanford-futuredata/ColBERT/issues) and we'll help you promptly!65 66 67 68## Overview69 70Using ColBERT on a dataset typically involves the following steps.71 72**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.73 74**Step 1: Download the [pre-trained ColBERTv2 checkpoint](https://downloads.cs.stanford.edu/nlp/data/colbert/colbertv2/colbertv2.0.tar.gz).** This checkpoint has been trained on the MS MARCO Passage Ranking task. You can also _optionally_ [train your own ColBERT model](#training).75 76**Step 2: Index your collection.** Once you have a trained ColBERT model, you need to [index your collection](#indexing) to permit fast retrieval. This step encodes all passages into matrices, stores them on disk, and builds data structures for efficient search.77 78**Step 3: Search the collection with your queries.** Given the model and index, you can [issue queries over the collection](#retrieval) to retrieve the top-k passages for each query.79 80Below, we illustrate these steps via an example run on the MS MARCO Passage Ranking task.81 82 83## API Usage Notebook84 85**NEW**: We have an experimental notebook on [Google Colab](https://colab.research.google.com/github/stanford-futuredata/ColBERT/blob/main/docs/intro2new.ipynb) that you can use with free GPUs. Indexing 10,000 on the free Colab T4 GPU takes six minutes.86 87This Jupyter notebook **[docs/intro.ipynb notebook](docs/intro.ipynb)** illustrates using the key features of ColBERT with the new Python API.88 89It includes how to download the ColBERTv2 model checkpoint trained on MS MARCO Passage Ranking and how to download our new LoTTE benchmark.90 91 92## Data93 94This repository works directly with a simple **tab-separated file** format to store queries, passages, and top-k ranked lists.95 96 97* Queries: each line is `qid \t query text`.98* Collection: each line is `pid \t passage text`.99* Top-k Ranking: each line is `qid \t pid \t rank`.100 101This works directly with the data format of the [MS MARCO Passage Ranking](https://github.com/microsoft/MSMARCO-Passage-Ranking) dataset. 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`).102 103 104## Indexing105 106For fast retrieval, indexing precomputes the ColBERT representations of passages.107 108Example usage:109 110```111from colbert.infra import Run, RunConfig, ColBERTConfig112from colbert import Indexer113 114if __name__=='__main__':115    with Run().context(RunConfig(nranks=1, experiment="msmarco")):116 117        config = ColBERTConfig(118            nbits=2,119            root="/path/to/experiments",120        )121        indexer = Indexer(checkpoint="/path/to/checkpoint", config=config)122        indexer.index(name="msmarco.nbits=2", collection="/path/to/MSMARCO/collection.tsv")123```124 125 126## Retrieval127 128We typically recommend that you use ColBERT for **end-to-end** retrieval, where it directly finds its top-k passages from the full collection:129 130```131from colbert.data import Queries132from colbert.infra import Run, RunConfig, ColBERTConfig133from colbert import Searcher134 135if __name__=='__main__':136    with Run().context(RunConfig(nranks=1, experiment="msmarco")):137 138        config = ColBERTConfig(139            root="/path/to/experiments",140        )141        searcher = Searcher(index="msmarco.nbits=2", config=config)142        queries = Queries("/path/to/MSMARCO/queries.dev.small.tsv")143        ranking = searcher.search_all(queries, k=100)144        ranking.save("msmarco.nbits=2.ranking.tsv")145```146 147You can optionally specify the `ncells`, `centroid_score_threshold`, and `ndocs` search hyperparameters to trade off between speed and result quality. Defaults for different values of `k` are listed in colbert/searcher.py.148 149We can evaluate the MSMARCO rankings using the following command:150 151```152python -m utility.evaluate.msmarco_passages --ranking "/path/to/msmarco.nbits=2.ranking.tsv" --qrels "/path/to/MSMARCO/qrels.dev.small.tsv"153```154 155## Training156 157We provide a [pre-trained model checkpoint](https://downloads.cs.stanford.edu/nlp/data/colbert/colbertv2/colbertv2.0.tar.gz), but we also detail how to train from scratch here.158Note that this example demonstrates the ColBERTv1 style of training, but the provided checkpoint was trained with ColBERTv2.159 160Training requires a JSONL triples file with a `[qid, pid+, pid-]` list per line. The query IDs and passage IDs correspond to the specified `queries.tsv` and `collection.tsv` files respectively.161 162Example usage (training on 4 GPUs):163 164```165from colbert.infra import Run, RunConfig, ColBERTConfig166from colbert import Trainer167 168if __name__=='__main__':169    with Run().context(RunConfig(nranks=4, experiment="msmarco")):170 171        config = ColBERTConfig(172            bsize=32,173            root="/path/to/experiments",174        )175        trainer = Trainer(176            triples="/path/to/MSMARCO/triples.train.small.tsv",177            queries="/path/to/MSMARCO/queries.train.small.tsv",178            collection="/path/to/MSMARCO/collection.tsv",179            config=config,180        )181 182        checkpoint_path = trainer.train()183 184        print(f"Saved checkpoint to {checkpoint_path}...")185```186 187## Running a lightweight ColBERTv2 server188We 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. 189 190To run the server, update the environment variables `INDEX_ROOT` and `INDEX_NAME` in the `.env` file to point to the appropriate ColBERT index. The run the following command:191```192python server.py193```194 195A sample query:196```197http://localhost:8893/api/search?query=Who won the 2022 FIFA world cup&k=25198```199 200## Branches201 202### Supported branches203 204* [`main`](https://github.com/stanford-futuredata/ColBERT/tree/main): Stable branch with ColBERTv2 + PLAID.205* [`colbertv1`](https://github.com/stanford-futuredata/ColBERT/tree/colbertv1): Legacy branch for ColBERTv1.206 207### Deprecated branches208* [`new_api`](https://github.com/stanford-futuredata/ColBERT/tree/new_api): Base ColBERTv2 implementation.209* [`cpu_inference`](https://github.com/stanford-futuredata/ColBERT/tree/cpu_inference): ColBERTv2 implementation with CPU search support.210* [`fast_search`](https://github.com/stanford-futuredata/ColBERT/tree/fast_search): ColBERTv2 implementation with PLAID.211* [`binarization`](https://github.com/stanford-futuredata/ColBERT/tree/binarization): ColBERT with a baseline binarization-based compression strategy (as opposed to ColBERTv2's residual compression, which we found to be more robust).212 213## Acknowledgments214 215ColBERT logo designed by Chuyi Zhang.