Ecosyste.ms: Awesome

An open API service indexing awesome lists of open source software.

Awesome Lists | Featured Topics | Projects

https://github.com/eva-kaushik/bert-ai

COL-BERT
https://github.com/eva-kaushik/bert-ai

huggingface huggingface-transformers machine-learning

Last synced: 2 days ago
JSON representation

COL-BERT

Awesome Lists containing this project

README 

        

# BERT-AI (v2)



### BERT-AI is a _fast_ and _accurate_ retrieval model, enabling scalable BERT-based search over large text collections in tens of milliseconds.





  Figure 1: late interaction, efficiently scoring the fine-grained similarity between a queries and a passage.




As 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.



These rich interactions allow to surpass the quality of _single-vector_ representation 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**](https://arxiv.org/abs/2004.12832) (SIGIR'20).

* [**Relevance-guided Supervision for OpenQA with ColBERT**](https://arxiv.org/abs/2007.00814) (TACL'21).

* [**Baleen: Robust Multi-Hop Reasoning at Scale via Condensed Retrieval**](https://arxiv.org/abs/2101.00436) (NeurIPS'21).

* [**ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction**](https://arxiv.org/abs/2112.01488) (NAACL'22).

* [**PLAID: An Efficient Engine for Late Interaction Retrieval**](https://arxiv.org/abs/2205.09707) (CIKM'22).

* [**Moving Beyond Downstream Task Accuracy for Information Retrieval Benchmarking**](https://arxiv.org/abs/2212.01340) (ACL'23 Findings).

* [**UDAPDR: Unsupervised Domain Adaptation via LLM Prompting and Distillation of Rerankers**](https://arxiv.org/abs/2303.00807) (EMNLP'23).



----



## 🚨 **Announcements** 



* (1/28/24) One of the easiest ways to use in applications nowadays is the semi-official, fast-growing [RAGatouille](https://github.com/bclavie/ragatouille) library.

* (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 the **DSPy** framework for composing retrievers like ColBERTv2 and LLMs, it's at: https://github.com/stanfordnlp/dspy



----



## Installation



(Update: nowadays you can typically do `pip install colbert-ai[torch,faiss-gpu]` to get things up and running, but if you face issues conda is always more reliable for `faiss` and `torch`.)



ColBERT requires Python 3.7+ and Pytorch 1.9+ and uses the [Hugging Face Transformers](https://github.com/huggingface/transformers) library.



We 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).)



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 specify `CUDA_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 colbert

```



## Data



This repository works directly with a simple **tab-separated file** format to store queries, passages, and top-k ranked lists.



* Queries: each line is `qid \t query text`.

* Collection: each line is `pid \t passage text`.

* Top-k Ranking: each line is `qid \t pid \t rank`.



This 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`).



## Indexing



For fast retrieval, indexing precomputes the ColBERT representations of passages.



Example usage:



```python

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")

```



## Retrieval



We typically recommend that you use ColBERT for **end-to-end** retrieval, where it directly finds its top-k passages from the full collection:



```python

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")

```



You 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.



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"

```



Training 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.



Example usage (training on 4 GPUs):



```python

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}...")

```



## Advanced Training (ColBERTv2-style)



```python

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()

```



## Running a lightweight ColBERTv2 server

We 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 variables `INDEX_ROOT` and `INDEX_NAME` in the `.env` file to point to the appropriate ColBERT index. The run the following command:

```

python server.py

```



A sample query:

```

http://localhost:8893/api/search?query=Who won the 2022 FIFA world cup&k=25

```