contrastors是对比学习工具包,使研究人员和工程师能够有效地训练和评估对比模型。
该库依赖于 Flash Attention 存储库中的自定义内核。要设置你的环境,你需要按照以下步骤操作。
contrastors
确保你有 Cuda 11.8+。你可以通过运行来检查这一点,或者如果你已经安装了 PyTorch ,则可以运行
nvcc --version
python -c "import torch; print(torch.version.cuda)"
创建一个 python venv 并激活它
python3 -m venv env
source env/bin/activate
安装割炬。有关系统的具体说明,请参阅 PyTorch 文档(例如,截至 2023 年 12 月 12 日,默认 CUDA PyTorch 支持的版本为 12.1)。
pip3 install torch torchvision torchaudio
为 Flash Attention 安装轮子、包装、忍者(因此构建不会花费太长时间)
pip install wheel packaging ninja
安装 Flash Attention 和自定义内核
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
安装其余要求和软件包
pip install -e .
我们通过软件包提供对数据集的访问。要访问数据,你需要创建一个帐户并登录软件包。首先在 atlas.nomic.ai 创建一个账号,下载Python客户端,然后运行以下命令:
nomic-embed-text-v1
nomic
nomic
nomic
pip install nomic
nomic login # follow prompts to login
python -c "from nomic import atlas; print(atlas._get_datastream_credentials(name='contrastors'))"
这将打印出你的访问密钥。然后,你可以通过使用 或 设置 和 环境变量。
aws configure
AWS_ACCESS_KEY_ID
AWS_SECRET_ACCESS_KEY
如果你尚未安装 AWS CLI,可以在此处安装它。
要验证你的访问权限,你可以运行以下命令来列出存储桶的内容:
aws s3 ls --endpoint-url=https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/ s3://contrastive
aws s3 ls --endpoint-url=https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/ s3://contrastive-index-filtered
你应该能够看到存储桶的内容并下载数据。
如果你打算使用我们的数据和存储库进行训练,则需要设置对 Cloudflare R2 的支持。为此, 创建一个包含以下内容的文件:
contrastors
fsspec
~/.config/fsspec/s3.json
{
"s3": {
"client_kwargs": {
"endpoint_url": "https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/",
"aws_access_key_id": <ACCESS_KEY_ID>,
"aws_secret_access_key": <SECRET_KEY_ID>
}
}
}
我们的文本数据存储在 gzip 的 jsonl 文件中,我们还用该文件存储文件和 .
counts.json
offsets.json.gzip
该文件是一个字典,将文件名映射到文件中的示例数。该文件是一个字典,将文件名映射到字典,其中每个键是示例的索引,值是文件中示例的开始和结束字节偏移量的元组。我们这样做是为了允许从 R2 流入数据,尤其是当数据大于缓冲区大时。
counts.json
offsets.json.gz
下面是数据集配置可能看起来像的一个小示例:
datasets:
- name: "paq"
bucket: "s3://contrastive-index-filtered/paq_full/shard-{00000..00538}.jsonl.gz"
query_prefix: "search_query"
document_prefix: "search_document"
objective:
type: "paired"
columns: ["query", "document"]
objective定义它是成对目标还是三目标。在这两种情况下,该字段都定义要用于每个示例的列。
columns
nomic-embed-text-v1
要从头开始训练你自己的 BERT(包含所有优化),请运行
cd src/contrastors
deepspeed --num_gpus=8 train.py --config=configs/train/mlm.yaml --deepspeed_config_file=configs/deepspeed/ds_config.json --dtype=bf16
启动试验运行
cd src/contrastors
torchrun --nproc-per-node=8 train.py --config=configs/train/contrastive_pretrain.yaml --dtype=bf16
这将在所有 ~200M 示例上训练一个 bert 模型。要更改数据集,你可以修改 .
data_args.input_shards
要微调,请将配置更新为 。
nomic-bert-embed-v1-unsupervised
configs/train/contrastive_finetune.yaml
若要为管道的任何步骤生成自己的数据,可以使用中提供的脚本。
scripts/text
有关详细信息,请参阅自述文件。
scripts/text
我们提供以下位置的预训练模型:
nomic-embed-text-v1
本项目和模型在 Apache 2.0 许可证下获得许可。
我们感谢 Tri Dao 在 Flash Attention 和使这个项目成为可能的自定义内核方面所做的工作,感谢 OpenCLIP 团队的 伟大的存储库,大部分工作都基于它,以及 Huggingface 团队在 transformers 库上的出色工作。
如果你觉得模型、数据集或训练代码有用,请引用我们的工作
@misc{nussbaum2024nomic,
title={Nomic Embed: Training a Reproducible Long Context Text Embedder},
author={Zach Nussbaum and John X. Morris and Brandon Duderstadt and Andriy Mulyar},
year={2024},
eprint={2402.01613},
archivePrefix={arXiv},
primaryClass={cs.CL}
}