contrastors

contrastors is an open-source contrastive learning toolkit from Nomic AI for researchers and engineers who train and evaluate contrastive models in PyTorch. It is the training code behind the Nomic Embed text and vision embedding models and is released under the Apache 2.0 license.

What It Is

The library trains embedding models contrastively. It builds on Flash Attention custom kernels for efficient training and supports multi-GPU runs. Training is configured through YAML files and launched with torchrun or DeepSpeed.

Training Workflow

The README documents a pipeline of masked language modeling pretraining for BERT-style models, followed by contrastive pretraining and contrastive finetuning. Dataset configs define paired or triplet objectives, with query and document prefixes. The README says the default pretraining config trains on roughly 200M examples. Scripts in the repo can generate your own data for any step of the pipeline.

For vision, the repo documents aligning nomic-embed-vision-v1.5 with nomic-embed-text-v1.5 using an image-text dataset.

Techniques and Model Support

• GradCache lets large batch sizes fit in constrained GPU memory.
• Matryoshka Representation Learning supports flexible embedding sizes.
• CLIP and LiT style contrastive learning are supported.
• Hugging Face models such as Pythia/GPTNeoX and BERT, and timm ViT models, can be loaded.

Setup and Data Access

Setup requires CUDA 11.8+, a Python virtual environment, PyTorch, and Flash Attention with its custom kernels. The nomic-embed-text-v1 training data is available through the nomic Python package after creating an Atlas account, and it is streamed from Cloudflare R2 storage. Pretrained Nomic Embed models are listed on Hugging Face.