contrastors

The contrastors library provides different entry points for training text models using train.py.

Masked Language Modeling (MLM) Pretraining

To train a model (like BERT) from scratch:

cd src/contrastors
deepspeed --num_gpus=8 train.py --config=configs/train/mlm.yaml --deepspeed_config=configs/deepspeed/ds_config.json --dtype=bf16

Contrastive Pretraining

To launch contrastive pretraining (e.g., on ~200M examples):

cd src/contrastors
torchrun --nproc-per-node=8 train.py --config=configs/train/contrastive_pretrain.yaml --dtype=bf16

Note: Modify data_args.input_shards in the config to change the dataset.

Contrastive Finetuning

To finetune nomic-bert-embed-v1-unsupervised:

# Use the finetune config instead of pretrain
torchrun --nproc-per-node=8 train.py --config=configs/train/contrastive_finetune.yaml --dtype=bf16

# MLM Pretraining
cd src/contrastors
deepspeed --num_gpus=8 train.py --config=configs/train/mlm.yaml --deepspeed_config=configs/deepspeed/ds_config.json --dtype=bf16

# Contrastive Pretraining
torchrun --nproc-per-node=8 train.py --config=configs/train/contrastive_pretrain.yaml --dtype=bf16

To set up the contrastors environment, ensure you have CUDA 11.8+ installed. Follow these steps to create a virtual environment, install PyTorch, and compile the necessary Flash Attention kernels.

1. Create and activate a virtual environment:

python3 -m venv env
source env/bin/activate

2. Install PyTorch:

pip3 install torch torchvision torchaudio

3. Install build dependencies:

pip install wheel packaging ninja setuptools

4. Install Flash Attention and custom kernels:

pip install --no-cache-dir flash-attn --no-build-isolation git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/layer_norm git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/fused_dense_lib git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/xentropy

5. Install the package in editable mode:

pip install -e .

To improve training quality for nomic-embed-text-v1, you can filter datasets to ensure queries are actually similar to their corresponding documents. This process embeds queries and documents, builds a nearest-neighbor index, and discards any (query_i, document_i) pair where document_i is not among the top-k most similar documents to query_i.

Input Format: The input dataset must be in .jsonl format with the following structure:

{"query": "Who won the World Series in 2016?", "document": "The Chicago Cubs won the World Series against the Cleveland Guardians."}

Requirements: It is highly recommended to install faiss-gpu for efficient indexing. For CUDA 12+, refer to INSTALL_FAISS.md in the repository.

Execution: Use index_filtering.py via torchrun to process the dataset.

torchrun --nproc-per-node=<num_gpus> --dataset=<path_to_dataset_files_or_directory> --output_dir=<path_where_to_save_filtered_dataset> --query_key=<query_key_of_jsonl_file> --document_key=<document_of_key_jsonl_file> index_filtering.py

MTEB evaluation is implemented with parallel execution. You can run the evaluation by specifying the model name. If you are using a Huggingface or OpenAI model that is not natively supported by contrastors, use the --hf_model or --openai_model flags respectively.

Results are saved as .json files in results/<model_name>.

Post-processing MTEB results:

1. Merge CQADupstack: The CQADupstack results must be merged into a single result using merge_cqadupstack.py.
2. Generate Metadata: Use mteb_meta.py to generate metadata.
3. Score Results: Use score_mteb.py on the generated metadata file to calculate final scores.

# Run evaluation
OPENBLAS_NUM_THREADS=16 CUDA_VISIBLE_DEVICES=<GPUS> python eval/mteb_eval/eval_mteb.py --model_name=<model_name>

# Merge CQADupstack results
python eval/mteb_eval/merge_cqadupstack.py results/<model_name>

# Generate metadata and scores
python eval/mteb_eval/mteb_meta.py results/<model_name>
python eval/mteb_eval/score_mteb.py <last folder prefix>/mteb_metadata.md

If you are using CUDA 12+, you may need to build FAISS from source. Follow these steps to prepare your environment, install dependencies, build CMake from source, and finally compile FAISS with Python bindings.

Prerequisites:

• Ensure you have activated your environment.
• You will need sudo privileges for package installation and system-wide installs.

Steps:

1. Activate your environment.
2. Clean up existing packages and install system dependencies (libssl-dev, libblas-dev, liblapack-dev).
3. Download and build cmake (v3.28.0-rc4) from source.
4. Build faiss using the generated build directory.
5. Build swigfaiss for Python bindings.
6. Install the Python package using setup.py from the build/faiss/python directory.

# 1. Activate environment
source env/bin/activate

# 2. Install dependencies
sudo apt purge --autoremove
sudo apt-get install libssl-dev
sudo apt-get install libblas-dev liblapack-dev

# 3. Build CMake from source
wget https://github.com/Kitware/CMake/releases/download/v3.28.0-rc4/cmake-3.28.0-rc4.tar.gz
tar -xvf cmake-3.28.0-rc4.tar.gz
cd cmake-3.28.0-rc4/
./bootstrap
make -j$(nproc)
sudo make install

# 4. Build FAISS
cmake -B build .
make -C build -j faiss

# 5. Build SWIG bindings
make -C build -j swigfaiss

# 6. Install Python package
(cd build/faiss/python && python setup.py install)

To access the nomic-embed-text-v1 dataset, you must use the nomic Python client to retrieve credentials for the Cloudflare R2 storage.

1. Install the client and login:

pip install nomic
nomic login

2. Retrieve credentials: Run the following command to print your access keys:

python -c "from nomic import atlas; print(atlas._get_datastream_credentials(name='contrastors'))"

3. Configure credentials: Set the AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables, or use aws configure.

4. Configure fsspec for streaming: To enable streaming from R2, create ~/.config/fsspec/s3.json with your credentials:

{
  "s3": {
    "client_kwargs": {
      "endpoint_url": "https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/",
      "aws_access_key_id": <ACCESS_KEY_ID>,
      "aws_secret_access_key": <SECRET_KEY_ID>
    }
  }
}