ModernBERT Documentation

ModernBERT supports two dataset classes located in src/text_data.py. You can switch between them by setting the streaming key in your configuration.

• StreamingTextDataset: Inherits from StreamingDataset. Supports MDS, CSV/TSV, or JSONL formats for both text and tokenized data. Works with local data but may have uneven memory distribution over accelerators.
• NoStreamingDataset: Requires decompressed MDS-format data. It is recommended for local data access as it enables higher training throughput.

To decompress MDS data, use src/data/mds_conversion.py with the --decompress flag.

train_loader:
  name: text
  dataset:
    streaming: false

If you are using H100 GPUs, it is recommended to clone and build Flash Attention 3 manually from the source.

git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention/hopper
python setup.py install

If your checkpoints are stored in a private or gated Hugging Face repository, you must authenticate your environment. You can do this via the CLI or by providing a hub_token directly to the evaluation scripts.

To authenticate via CLI:

huggingface-cli login

When generating an evaluation config with generate_eval_config_from_checkpoint.py, you can skip specific evaluation tasks by adding the --skip_<eval_name> flag. For example, to skip the MNLI evaluation, use --skip_mnli.

python generate_eval_config_from_checkpoint.py \
--checkpoint /path/to/checkpoint/folder \
--wandb_entity entity_name \
--wandb_project project_name \
--track_run \
--skip_mnli

To re-generate the training data used in BERT24, follow these steps:

1. Install the necessary dependencies from requirements.txt and requirements-data.txt.
2. Convert a Hugging Face (HF) dataset into the MDS format using hf_to_mds.py.
3. Sample each dataset using sample_dataset_from_config.py (Note: implementation details for this step are currently marked as TODO in the source).

For a more user-friendly and detailed interface than nvidia-smi when monitoring GPU usage during evaluations, install and use nvitop.

pip install nvitop
nvitop

Use generate_eval_config_from_checkpoint.py to create evaluation configuration files based on a model checkpoint. You can generate configs by providing a local training configuration file or by matching the checkpoint to a Weights & Biases (wandb) run.

To create a config using a local training config:

python generate_eval_config_from_checkpoint.py \
--checkpoint /path/to/checkpoint/folder \
--train_config /path/to/config.yaml

To create a config from a matching wandb run and enable wandb tracking:

python generate_eval_config_from_checkpoint.py \
--checkpoint /path/to/checkpoint/folder \
--wandb_entity entity_name \
--wandb_project project_name \
--track_run

After training, you can evaluate your models on BEIR datasets using the following scripts:

• For single vector models: python evaluate_st.py
• For multi vector models: python evaluate_pylate.py

To train a single vector retrieval model (DPR) using ModernBERT as the backbone, use the train_st.py script. This script utilizes Sentence Transformers and performs contrastive learning on the MS-MARCO dataset with mined hard negatives.

Alternatively, you can use train_st_gooaq.py to train a single vector model specifically on the GooAQ question-answer dataset.

The examples directory contains scripts for training and evaluating retrieval models:

ColBERT models (via PyLate):

• examples/train_pylate.py: Boilerplate for training.
• examples/evaluate_pylate.py: Boilerplate for evaluation.

Dense models (via Sentence Transformers):

• examples/train_st.py: Boilerplate for training.
• examples/evaluate_st.py: Boilerplate for evaluation.