DOCSEARCHbyAlgolia
Submit Docs

Contriever: Unsupervised Dense Information Retrieval

repository·main·Indexed 20 days ago

https://github.com/facebookresearch/contriever

A framework for unsupervised dense information retrieval using contrastive learning. Contriever provides pre-trained monolingual and multilingual models (including mContriever) to generate sentence embeddings for retrieval tasks. It includes tools for evaluating Question Answering retrieval, the BEIR benchmark, and Cross-lingual MKQA, as well as scripts for pre-processing data and training models using contrastive learning.

Tokens
9.6K
Snippets
23
Records
31
Agent score
73%

What's inside Contriever

  1. Evaluate mContriever on Cross-lingual MKQA

    main

    To measure how well retrievers retrieve relevant English Wikipedia documents given a query in another language using the MKQA dataset, follow these steps:

    1. Download data: Get the mkqa.jsonl.gz file.
    2. Preprocess data: Use data_scripts/preprocess_xmkqa.py to prepare the data.
    3. Generate embeddings: Use generate_passage_embeddings.py to create embeddings for your passages. Alternatively, you can download pre-computed embeddings for mcontriever or mcontriever-msmarco from Facebook's public servers.
    4. Retrieve and compute accuracy: Use passage_retrieval.py to perform retrieval and calculate accuracy.

    When generating embeddings or performing retrieval, use the --lowercase and --normalize_text flags.

    # 1. Download data
wget https://raw.githubusercontent.com/apple/ml-mkqa/master/dataset/mkqa.jsonl.gz

# 2. Preprocess data
python data_scripts/preprocess_xmkqa.py mkqa.jsonl xmkqa

# 3. Generate embeddings
python generate_passage_embeddings.py \
    --model_name_or_path facebook/mcontriever \
    --output_dir mcontriever_embeddings  \
    --passages psgs_w100.tsv \
    --shard_id 0 --num_shards 1 \
    --lowercase --normalize_text

# (Optional) Download pre-computed embeddings instead
# wget https://dl.fbaipublicfiles.com/contriever/embeddings/mcontriever/wikipedia_embeddings.tar
# wget https://dl.fbaipublicfiles.com/contriever/embeddings/mcontriever-msmarco/wikipedia_embeddings.tar

# 4. Retrieve passages and compute retrieval accuracy
python passage_retrieval.py \
    --model_name_or_path facebook/mcontriever \
    --passages psgs_w100.tsv \
    --passages_embeddings "mcontriever_embeddings/*" \
    --data "xmkqa/*.jsonl" \
    --output_dir mcontriever_xmkqa \
    --lowercase --normalize_text

  2. Get started with Contriever for sentence embeddings

    main

    You can use Contriever to obtain dense embeddings for sentences using the Contriever class and the HuggingFace transformers library. Once embeddings are obtained, similarity between sentences can be calculated using a dot product.

    from src.contriever import Contriever
from transformers import AutoTokenizer

# Load model and tokenizer
contriever = Contriever.from_pretrained("facebook/contriever") 
tokenizer = AutoTokenizer.from_pretrained("facebook/contriever")

sentences = [
    "Where was Marie Curie born?",
    "Maria Sklodowska, later known as Marie Curie, was born on November 7, 1867.",
    "Born in Paris on 15 May 1859, Pierre Curie was the son of Eugène Curie, a doctor of French Catholic origin from Alsace."
]

# Generate embeddings
inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")
embeddings = contriever(**inputs)

# Calculate similarity scores using dot product
score01 = embeddings[0] @ embeddings[1] # Similarity between sentence 0 and 1
score02 = embeddings[0] @ embeddings[2] # Similarity between sentence 0 and 2

  3. Train Contriever or mContriever using train.py

    main

    The train.py script handles the contrastive training phase.

    Contriever (English Monolingual)

    Used for training on English Wikipedia and CCNet data. Key parameters include --retriever_model_id bert-base-uncased, --augmentation delete, and --moco_queue 131072.

    mContriever (Multilingual)

    Used for training on 29 languages using CCNet data. Key parameters include --retriever_model_id bert-base-multilingual-cased and --moco_queue 32768.

    Full training scripts for Slurm clusters are available in the example_scripts folder.

    # Example: Training Contriever (English)
python train.py \
        --retriever_model_id bert-base-uncased --pooling average \
        --augmentation delete --prob_augmentation 0.1 \
        --train_data "data/wiki/ data/cc-net/" --loading_mode split \
        --ratio_min 0.1 --ratio_max 0.5 --chunk_length 256 \
        --momentum 0.9995 --moco_queue 131072 --temperature 0.05 \
        --warmup_steps 20000 --total_steps 500000 --lr 0.00005 \
        --scheduler linear --optim adamw --per_gpu_batch_size 64 \
        --output_dir /checkpoint/gizacard/contriever/xling/contriever

# Example: Training mContriever (Multilingual)
TDIR=encoded-data/bert-base-multilingual-cased/
TRAINDATASETS="${TDIR}fr_XX ${TDIR}en_XX ${TDIR}ar_AR ${TDIR}bn_IN ${TDIR}fi_FI ${TDIR}id_ID ${TDIR}ja_XX ${TDIR}ko_KR ${TDIR}ru_RU ${TDIR}sw_KE ${TDIR}hu_HU ${TDIR}he_IL ${TDIR}it_IT ${TDIR}km_KM ${TDIR}ms_MY ${TDIR}nl_XX ${TDIR}no_XX ${TDIR}pl_PL ${TDIR}pt_XX ${TDIR}sv_SE ${TDIR}te_IN ${TDIR}th_TH ${TDIR}tr_TR ${TDIR}vi_VN ${TDIR}zh_CN ${TDIR}zh_TW ${TDIR}es_XX ${TDIR}de_DE ${TDIR}da_DK"

python train.py \
        --retriever_model_id bert-base-multilingual-cased --pooling average \
        --train_data ${TRAINDATASETS} --loading_mode split \
        --ratio_min 0.1 --ratio_max 0.5 --chunk_length 256 \
        --momentum 0.999 --moco_queue 32768 --temperature 0.05 \
        --warmup_steps 20000 --total_steps 500000 --lr 0.00005 \
        --scheduler linear --optim adamw --per_gpu_batch_size 64 \
        --output_dir /checkpoint/gizacard/contriever/xling/mcontriever

  4. Pre-process data for Contriever training

    main

    Contriever pre-training uses data from CCNet and Wikipedia.

    1. Format: Convert data into a text file.
    2. Tokenization/Chunking: Use data_scripts/tokenization_script.sh to tokenize and chunk the text into multiple sub-files. These chunks can be loaded separately by different processes in a distributed training job.
    3. Normalization: For the multilingual model (mContriever), use the --normalize_text option to preprocess data and normalize common characters not present in the mBERT tokenizer.

  5. Evaluate Question Answering Retrieval

    main

    To evaluate Contriever on Question Answering tasks (like NaturalQuestions or TriviaQA), follow these steps:

    1. Download Passages: Download the Wikipedia passages used in DPR.

      wget https://dl.fbaipublicfiles.com/dpr/wikipedia_split/psgs_w100.tsv.gz

    2. Generate Passage Embeddings: Use generate_passage_embeddings.py to create embeddings for your passages.

      python generate_passage_embeddings.py \
    --model_name_or_path facebook/contriever \
    --output_dir contriever_embeddings  \
    --passages psgs_w100.tsv \
    --shard_id 0 --num_shards 1 \

      Alternatively, you can download pre-computed embeddings for contriever or contriever-msmarco from Facebook's public files.

    3. Perform Retrieval: Use passage_retrieval.py to retrieve the top-100 passages for a given dataset.

      python passage_retrieval.py \
    --model_name_or_path facebook/contriever \
    --passages psgs_w100.tsv \
    --passages_embeddings "contriever_embeddings/*" \
    --data nq_dir/test.json \
    --output_dir contriever_nq \

  6. Evaluate mContriever on Mr. TyDi v1.1

    main

    To reproduce multilingual evaluation results on the Mr. TyDi v1.1 dataset, follow these steps (example provided for Swahili):

    1. Download data: Fetch the dataset and extract it.
    2. Convert data: Use data_scripts/convertmrtydi2beir.py to convert the dataset into the BEIR format.
    3. Run Evaluation: Use beireval.py with the facebook/mcontriever model path and the converted dataset path.

    Note: Ensure you use the --normalize_text flag during evaluation.

    # 1. Download data
wget https://git.uwaterloo.ca/jimmylin/mr.tydi/-/raw/master/data/mrtydi-v1.1-swahili.tar.gz -P mrtydi
tar -xf mrtydi/mrtydi-v1.1-swahili.tar.gz -C mrtydi
gzip -d mrtydi/mrtydi-v1.1-swahili/collection/docs.jsonl.gz

# 2. Convert data
python data_scripts/convertmrtydi2beir.py mrtydi/mrtydi-v1.1-swahili mrtydi/mrtydi-v1.1-swahili

# 3. Evaluation
python beireval.py --model_name_or_path facebook/mcontriever --dataset mrtydi/mrtydi-v1.1-swahili --normalize_text

  7. Configure pooling strategies in Contriever models

    main

    The Contriever and XLMRetriever classes support two types of pooling to generate dense embeddings from the transformer's last hidden states:

    1. "average": Computes the mean of the last hidden states across the sequence length, masked by the attention_mask to ignore padding tokens.
    2. "cls": Uses the embedding from the first token ([CLS]) of the sequence.

    Pooling is configured via the pooling argument during initialization or via the config.pooling attribute.

  8. How MultiDataset handles multiple data sources

    main

    A MultiDataset is a wrapper that combines multiple Dataset objects. It allows for sampling from different datasets during training.

    Key behaviors:

    • Sampling Probability: You can adjust the probability of selecting a specific dataset using set_prob(coeff=...). A coeff of 0.0 results in uniform sampling based on dataset size. Higher coefficients can be used to rebalance sampling.
    • Item Retrieval: When __getitem__ is called, it randomly selects a dataset based on the calculated probabilities and then returns a random sample from that specific dataset. The returned sample includes a "dataset_id" key.
    from src.data import MultiDataset

# datasets is a dict of {path: Dataset_instance}
multi_dataset = MultiDataset(datasets)

# Adjust sampling weights
multi_dataset.set_prob(coeff=0.5)

sample = multi_dataset[0]

  10. Configure data augmentation in Dataset

    main

    The Dataset class applies random cropping and augmentations to token sequences. Augmentations are controlled via the opt object.

    Supported opt.augmentation values:

    • "mask": Uses opt.mask_id and opt.prob_augmentation to mask tokens.
    • "replace": Replaces tokens with random IDs between opt.start_id and opt.vocab_size - 1 with probability opt.prob_augmentation.
    • "delete": Deletes tokens with probability opt.prob_augmentation.
    • "shuffle": Shuffles tokens with probability opt.prob_augmentation.

    Other relevant opt keys:

    • ratio_min, ratio_max: Bounds for the random crop ratio applied to chunks.

  11. Configure Indexer quantization settings

    main

    When initializing the Indexer, you can choose between exact search or compressed search using Product Quantization (PQ) to save memory.

    • Exact Search (Default): Set n_subquantizers=0. This uses faiss.IndexFlatIP, which performs exact Inner Product search. It is highly accurate but requires more memory as it stores full vectors.
    • Compressed Search (PQ): Set n_subquantizers > 0. This uses faiss.IndexPQ. This significantly reduces the memory footprint by quantizing vectors into sub-spaces, which is useful for very large datasets, though it introduces some approximation error.

    Example configuration for PQ:

    # Uses Product Quantization with 8 subquantizers and 8 bits each
indexer = Indexer(vector_sz=768, n_subquantizers=8, n_bits=8)

  12. Available Pre-trained Contriever models

    main

    Contriever provides several pre-trained model variants via HuggingFace. You can load them using Contriever.from_pretrained().

    • facebook/contriever: Unsupervised pre-training on CC-net and English Wikipedia.
    • facebook/contriever-msmarco: Contriever fine-tuned on MSMARCO.
    • facebook/mcontriever: Multilingual version pre-trained on 29 languages using CC-net.
    • facebook/mcontriever-msmarco: Multilingual version fine-tuned on MSMARCO.
    from src.contriever import Contriever

contriever = Contriever.from_pretrained("facebook/contriever") 
contriever_msmarco = Contriever.from_pretrained("facebook/contriever-msmarco")
mcontriever = Contriever.from_pretrained("facebook/mcontriever")
mcontriever_msmarco = Contriever.from_pretrained("facebook/mcontriever-msmarco")