{"id":24973033,"url":"https://github.com/nomic-ai/contrastors","last_synced_at":"2025-05-15T11:08:23.903Z","repository":{"id":220355919,"uuid":"750409062","full_name":"nomic-ai/contrastors","owner":"nomic-ai","description":"Train Models Contrastively in Pytorch","archived":false,"fork":false,"pushed_at":"2025-03-26T18:40:32.000Z","size":3930,"stargazers_count":689,"open_issues_count":9,"forks_count":56,"subscribers_count":14,"default_branch":"main","last_synced_at":"2025-04-08T06:34:32.021Z","etag":null,"topics":["contrastive-learning","deep-learning","dense-retrieval","embeddings","image-embeddings","multimodal","multimodal-rag","pytorch","rag","text-embeddings","transformers"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/nomic-ai.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-01-30T15:39:06.000Z","updated_at":"2025-04-07T10:27:33.000Z","dependencies_parsed_at":"2024-02-23T17:28:40.999Z","dependency_job_id":"b7bf4320-a2b0-4b84-8d19-6e438a952772","html_url":"https://github.com/nomic-ai/contrastors","commit_stats":null,"previous_names":["nomic-ai/contrastors"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/nomic-ai%2Fcontrastors","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/nomic-ai%2Fcontrastors/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/nomic-ai%2Fcontrastors/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/nomic-ai%2Fcontrastors/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/nomic-ai","download_url":"https://codeload.github.com/nomic-ai/contrastors/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":254328384,"owners_count":22052632,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["contrastive-learning","deep-learning","dense-retrieval","embeddings","image-embeddings","multimodal","multimodal-rag","pytorch","rag","text-embeddings","transformers"],"created_at":"2025-02-03T18:19:53.119Z","updated_at":"2025-05-15T11:08:23.887Z","avatar_url":"https://github.com/nomic-ai.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"# contrastors\n\n`contrastors` is contrastive learning toolkit that enables researchers and engineers to train and evaluate contrastive models efficiently.\n\n\n[![img](docs/atlas-nomic-embed.png)](https://atlas.nomic.ai/map/nomic-text-embed-v1-5m-sample)\n\n\n## Features\n\n- Built on top of [Flash Attention](https://github.com/Dao-AILab/flash-attention) for fast and efficient training\n- Support for training on multiple GPUs\n- [GradCache](https://github.com/luyug/GradCache) support for training with large batch sizes in constrained memory environments\n- Huggingface Support for easy loading of common models (Pythia/GPTNeoX, BERT, etc.)\n- Masked Language Modeling (MLM) Pretraining\n- [Matryoshka Representation Learning](https://arxiv.org/abs/2205.13147) for flexible embedding sizes\n- [CLIP](https://arxiv.org/abs/2103.00020) and [LiT](https://arxiv.org/abs/2111.07991) style contrastive learning\n- Support for loading popular ViT (e.g. [timm](https://huggingface.co/timm)) models\n\n## Research\n\n* [Nomic Embed: Training a Reproducible Long Context Text Embedder](https://arxiv.org/abs/2402.01613) by Zach Nussbaum, Jack Morris, Andriy Mulyar, and Brandon Duderstadt\n* [Nomic Embed Vision: Expanding the Latent Space](https://arxiv.org/abs/2406.18587) by Zach Nussbaum, Brandon Duderstadt, and Andriy Mulyar\n* [Training Sparse Mixture Of Experts Text Embedding Models](https://static.nomic.ai/nomic_embed_multi_preprint.pdf) by Zach Nussbaum and Brandon Duderstadt \n\n## Getting Started and Requirements\n\nThe `contrastors` library relies on custom kernels from the [Flash Attention](https://github.com/Dao-AILab/flash-attention) repository. To setup your enviornment you will need to follow the steps below.\n\nMake sure that you have Cuda 11.8+. You can check this by running `nvcc --version` or if you already have torch installed you can run `python -c \"import torch; print(torch.version.cuda)\"`\n\nCreate a python venv and activate it\n\n```bash\npython3 -m venv env\nsource env/bin/activate\n```\n\nInstall [torch](https://pytorch.org/get-started/locally/). See the torch docs for specific instructions for your system (e.g. the default CUDA torch supports is 12.1 as of 12/12/2023).\n\n```bash\npip3 install torch torchvision torchaudio\n```\n\nInstall wheel, packaging, ninja for Flash Attention (so the builds don't take too long)\n\n```bash\npip install wheel packaging ninja setuptools\n```\n\nInstall Flash Attention and the custom kernels\n\n```bash\npip 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\n```\n\nInstall the rest of the requirements and the package\n\n```bash\npip install -e . \n```\n\n## Data Access\n\nWe provide access to the `nomic-embed-text-v1` dataset via the `nomic` package. To access the data, you will need to create an account and login to the `nomic` package. First create an account at [atlas.nomic.ai](https://atlas.nomic.ai), download the `nomic` Python client, and run the following commands:\n\n```bash\npip install nomic\nnomic login # follow prompts to login\npython -c \"from nomic import atlas; print(atlas._get_datastream_credentials(name='contrastors'))\"\n```\n\nwhich will print out your access keys. You can then configure them by using `aws configure` or setting\nthe `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables.\n\nIf you do not have the AWS CLI installed, you can install it [here](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html).\n\nTo verify your access, you can run the following command to list the contents of the bucket:\n\n```bash\naws s3 ls --endpoint-url=https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/ s3://contrastive\naws s3 ls --endpoint-url=https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/ s3://contrastive-index-filtered\n```\n\nYou should be able to see the contents of the bucket and download the data.\n\nIf you intend to train using our data and the `contrastors` repo, you will need to setup `fsspec` support for Cloudflare R2. To do so,\ncreate a file `~/.config/fsspec/s3.json` with the following contents:\n\n```json\n{\n  \"s3\": {\n    \"client_kwargs\": {\n      \"endpoint_url\": \"https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/\",\n      \"aws_access_key_id\": \u003cACCESS_KEY_ID\u003e,\n      \"aws_secret_access_key\": \u003cSECRET_KEY_ID\u003e\n    }\n  }\n}\n```\n\n### Nomic Data Format\n\nOur text data is stored in gziped jsonl files with which we also store a `counts.json` file and `offsets.json.gzip`.\n\nThe `counts.json` file is a dictionary mapping the file name to the number of examples in the file. The `offsets.json.gz` file is a dictionary mapping the file name to a dictionary where each key is the index of the example and the value is a tuple of the start and end byte offset of the example in the file. We do this to allow for streaming of data in from R2, especially when the data is larger than the buffer size.\n\nHere's a small example of what a dataset configuration might look like:\n\n```yaml\ndatasets:\n  - name: \"paq\"\n    bucket: \"s3://contrastive-index-filtered/paq_full/shard-{00000..00538}.jsonl.gz\"\n    query_prefix: \"search_query\"\n    document_prefix: \"search_document\"\n    objective: \n        type: \"paired\"\n        columns: [\"query\", \"document\"]\n\n```\n\n`objective` defines if it's a paired or triplet objective. In both cases, the `columns` field defines the columns to use for each example.\n\n## Training `nomic-embed-text-v1`\n\n### Masked Language Modeling Pretraining\n\nTo train your own BERT from scratch (with all the optimizations) run\n\n```bash\ncd src/contrastors\ndeepspeed --num_gpus=8 train.py --config=configs/train/mlm.yaml --deepspeed_config=configs/deepspeed/ds_config.json --dtype=bf16\n```\n\n### Constrastive Pretraining and Finetuning\n\nTo launch an experiment run\n\n```bash\ncd src/contrastors\ntorchrun --nproc-per-node=8 train.py --config=configs/train/contrastive_pretrain.yaml --dtype=bf16\n```\n\nThis will train a bert model on all ~200M examples. To change the dataset, you can modify `data_args.input_shards`.\n\nTo finetune `nomic-bert-embed-v1-unsupervised`, update the config to `configs/train/contrastive_finetune.yaml`.\n\n### Generating Your Own Data\n\nTo generate your own data for any step of the pipeline, you can use the provided scripts in `scripts/text`. \n\nSee the [README](scripts/text/README.md) in `scripts/text` for more information.\n\n\n\n## Training `nomic-embed-vision-v1.5`\n\nTo align a vision model, you will need to curate a large image-text dataset. More details can be found [here](https://github.com/rom1504/img2dataset).\n\nTo align `nomic-embed-vision-v1.5` with `nomic-embed-text-v1.5`, you can run the following command:\n\n```bash\ndeepspeed  train.py --deepspeed_config=configs/deepspeed/image_text.json --config=configs/train/nomic_embed_vision_v1.5.yaml --dtype=bf16\n```\n\n## Pretrained Models\n\nWe provide pretrained models for `Nomic Embed` at the following locations:\n\n- [nomic-embed-text-v2-moe](https://huggingface.co/nomic-ai/nomic-embed-text-v2-moe)\n- [nomic-embed-text-v2-moe-unsupervised](https://huggingface.co/nomic-ai/nomic-embed-text-v2-moe-unsupervised)\n- [nomic-embed-text-v1](https://huggingface.co/nomic-ai/nomic-embed-text-v1)\n- [nomic-embed-vision-v1](https://huggingface.co/nomic-ai/nomic-embed-vision-v1)\n- [nomic-embed-text-v1.5](https://huggingface.co/nomic-ai/nomic-embed-text-v1.5)\n- [nomic-embed-vision-v1.5](https://huggingface.co/nomic-ai/nomic-embed-vision-v1.5)\n- [nomic-embed-text-v1-ablated](https://huggingface.co/nomic-ai/nomic-embed-text-v1-ablated)\n- [nomic-embed-text-v1-unsupervised](https://huggingface.co/nomic-ai/nomic-embed-text-v1-unsupervised)\n- [nomic-bert-2048](https://huggingface.co/nomic-ai/nomic-bert-2048)\n- [nomic-xlm-2048](https://huggingface.co/nomic-ai/nomic-xlm-2048)\n\n## Join the Nomic Community\n\n- Nomic: [https://nomic.ai](https://nomic.ai)\n- Discord: [https://discord.gg/myY5YDR8z8](https://discord.gg/myY5YDR8z8)\n- Twitter: [https://twitter.com/nomic_ai](https://twitter.com/nomic_ai)\n\n## License\n\nThis code is licensed under the [Apache 2.0 License](LICENSE). See the model cards for the individual license for each model. \n\n## Acknowledgements\n\nWe thank Tri Dao for his work on Flash Attention and the custom kernels that make this project possible, the [OpenCLIP](https://github.com/mlfoundations/open_clip) team for their\ngreat repository with which much of this work is based on, and the Huggingface team for their great work on the transformers library.\n\n\n## Citation\n\nIf you find the model, dataset, or training code useful, please cite our work\n\n```bibtex\n@misc{nussbaum2024nomic,\n      title={Nomic Embed: Training a Reproducible Long Context Text Embedder}, \n      author={Zach Nussbaum and John X. Morris and Brandon Duderstadt and Andriy Mulyar},\n      year={2024},\n      eprint={2402.01613},\n      archivePrefix={arXiv},\n      primaryClass={cs.CL}\n}\n@misc{nussbaum2024nomicembedvisionexpanding,\n      title={Nomic Embed Vision: Expanding the Latent Space}, \n      author={Zach Nussbaum and Brandon Duderstadt and Andriy Mulyar},\n      year={2024},\n      eprint={2406.18587},\n      archivePrefix={arXiv},\n      primaryClass={cs.CV},\n      url={https://arxiv.org/abs/2406.18587}, \n}\n@misc{nussbaum2025trainingsparsemixtureexperts,\n      title={Training Sparse Mixture Of Experts Text Embedding Models}, \n      author={Zach Nussbaum and Brandon Duderstadt},\n      year={2025},\n      eprint={2502.07972},\n      archivePrefix={arXiv},\n      primaryClass={cs.CL},\n      url={https://arxiv.org/abs/2502.07972}, \n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fnomic-ai%2Fcontrastors","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fnomic-ai%2Fcontrastors","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fnomic-ai%2Fcontrastors/lists"}