DNABERT-2

To use DNABERT-2, create a Python 3.8 virtual environment using Conda and install the required dependencies. If you want to use Flash Attention, you must also install triton from source.

Standard Installation

1. Create and activate a Conda environment named dna with Python 3.8.
2. Install dependencies from requirements.txt.

Optional: Install Triton for Flash Attention

If you require Flash Attention, follow these steps to build triton from source:

1. Clone the OpenAI Triton repository.
2. Navigate to the triton/python directory.
3. Install cmake as a build dependency.
4. Install Triton in editable mode.

# create and activate virtual python environment
conda create -n dna python=3.8
conda activate dna

# (optional if you would like to use flash attention)
# install triton from source
git clone https://github.com/openai/triton.git;
cd triton/python;
pip install cmake; # build-time dependency
pip install -e .

# install required packages
python3 -m pip install -r requirements.txt

To evaluate DNABERT-2 or other models on the Genome Understanding Evaluation (GUE) benchmark, download the GUE dataset and use the provided shell scripts in the finetune directory.

Note on Batch Sizes: The scripts are configured for 4 GPUs using DataParallel. If you have a different number of GPUs, adjust per_device_train_batch_size and gradient_accumulation_steps to maintain a global batch size of 32. For multi-GPU training, use torchrun instead of python.

export DATA_PATH=/path/to/GUE
cd finetune

# Evaluate DNABERT-2 on GUE
sh scripts/run_dnabert2.sh DATA_PATH

# Evaluate DNABERT (e.g., DNABERT with 3-mer) on GUE
# 3 for 3-mer, 4 for 4-mer, 5 for 5-mer, 6 for 6-mer
sh scripts/run_dnabert1.sh DATA_PATH 3

# Evaluate Nucleotide Transformers on GUE
# 0 for 500m-1000g, 1 for 500m-human-ref, 2 for 2.5b-1000g, 3 for 2.5b-multi-species
sh scripts/run_nt.sh DATA_PATH 0

DNABERT-2 is compatible with the transformers library. The loading method depends on your installed version of transformers.

• For version 4.28: Use AutoTokenizer and AutoModel with trust_remote_code=True.
• For version > 4.28: Use BertConfig to load the configuration first, then pass it to AutoModel with trust_remote_code=True.

# For transformers version 4.28
import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("zhihan1996/DNABERT-2-117M", trust_remote_code=True)
model = AutoModel.from_pretrained("zhihan1996/DNABERT-2-117M", trust_remote_code=True)

# For transformers version > 4.28
from transformers.models.bert.configuration_bert import BertConfig

config = BertConfig.from_pretrained("zhihan1996/DNABERT-2-117M")
model = AutoModel.from_pretrained("zhihan1996/DNABERT-2-117M", trust_remote_code=True, config=config)

1. Dataset Format

Prepare three CSV files: train.csv, dev.csv, and test.csv. The first row must be a header: sequence, label. Each subsequent row should contain a DNA sequence and a numerical label separated by a comma.

Example format:

sequence, label
ACGTCAGTCAGCGTACGT, 1

2. Training Execution

Navigate to the finetune directory and set the following environment variables:

• DATA_PATH: Path to your CSV files.
• MAX_LENGTH: Set this to approximately $0.25 imes$ your sequence length (since BPE tokenization reduces sequence length by about 5x).
• LR: Learning rate (e.g., 3e-5).

You can train using DataParallel (standard python command) or DistributedDataParallel (using torchrun for better efficiency).

cd finetune

export DATA_PATH=$path/to/data/folder
export MAX_LENGTH=100 
export LR=3e-5

# Training use DataParallel
python train.py \
    --model_name_or_path zhihan1996/DNABERT-2-117M \
    --data_path  ${DATA_PATH} \
    --kmer -1 \
    --run_name DNABERT2_${DATA_PATH} \
    --model_max_length ${MAX_LENGTH} \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 16 \
    --gradient_accumulation_steps 1 \
    --learning_rate ${LR} \
    --num_train_epochs 5 \
    --fp16 \
    --save_steps 200 \
    --output_dir output/dnabert2 \
    --evaluation_strategy steps \
    --eval_steps 200 \
    --warmup_steps 50 \
    --logging_steps 100 \
    --overwrite_output_dir True \
    --log_level info \
    --find_unused_parameters False

# Training use DistributedDataParallel (more efficient)
export num_gpu=4

torchrun --nproc_per_node=${num_gpu} train.py \
    --model_name_or_path zhihan1996/DNABERT-2-117M \
    --data_path  ${DATA_PATH} \
    --kmer -1 \
    --run_name DNABERT2_${DATA_PATH} \
    --model_max_length ${MAX_LENGTH} \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 16 \
    --gradient_accumulation_steps 1 \
    --learning_rate ${LR} \
    --num_train_epochs 5 \
    --fp16 \
    --save_steps 200 \
    --output_dir output/dnabert2 \
    --evaluation_strategy steps \
    --eval_steps 200 \
    --warmup_steps 50 \
    --logging_steps 100 \
    --overwrite_output_dir True \
    --log_level info \
    --find_unused_parameters False