对比器

contrastors

特征

• 建立在 Flash Attention 之上，可实现快速高效的训练
• 支持在多个 GPU 上进行训练
• GradCache 支持在受限的内存环境中进行大批量训练
• Huggingface 支持轻松加载常见模型（Pythia/GPTNeoX、BERT 等）
• 掩码语言建模 （MLM） 预训练

研究

• Nomic Embed： Training a Reproducible Long Context Text Embedder 作者：Zach Nussbaum、Jack Morris、Andrei Mulyar 和 Brandon Duderstadt

入门和要求

该库依赖于 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>
    }
  }
}

Nomic 数据格式

我们的文本数据存储在 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

Constrastive 预训练和微调

启动试验运行

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

• nomic-embed-text-v1 （中文）
• nomic-embed-text-v1-消融
• nomic-embed-text-v1-无监督
• 诺米克-BERT-2048

加入 Nomic 社区

• 名词：https://nomic.ai
• 不和谐：https://discord.gg/myY5YDR8z8
• 推特：https://twitter.com/nomic_ai

许可证

本项目和模型在 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}
}