DNA Language Model Benchmarking Framework

A comprehensive benchmarking framework for evaluating and comparing state-of-the-art DNA language models on prophage identification and other genomic classification tasks.

Overview

This framework provides a unified interface for benchmarking multiple DNA language models, including DNABERT, Nucleotide Transformer, HyenaDNA, and others. It's specifically designed for evaluating model performance on prophage identification but can be extended to other genomic sequence classification tasks.

Features

• 10+ Pre-trained DNA Language Models: Support for DNABERT1/2, Nucleotide Transformer, ProkBERT, GROVER, GENA-LM, INHERIT, HyenaDNA, EVO, and Caduceus
• Unified Interface: Abstract base class providing consistent API across all models
• Flexible Tokenization: Support for k-mer, BPE, character-level, and specialized tokenization strategies
• Long Sequence Handling: Automatic sequence chunking and aggregation for models with different max lengths (512 to 1M+ tokens)
• Comprehensive Evaluation: Multiple metrics including accuracy, F1, precision, recall, AUROC, and AUPRC
• Cross-validation Support: Built-in k-fold cross-validation with stratification
• SLURM Integration: Ready-to-use templates for HPC cluster deployment

Installation

Base Framework Installation

# Clone the repository
git clone https://github.com/leannmlindsey/dna_benchmark.git
cd dna_benchmark

# Create base conda environment
conda create -n dna_benchmark python=3.9 -y
conda activate dna_benchmark

# Install base dependencies
pip install -r requirements.txt

# Install the package in development mode
pip install -e .

Model-Specific Environment Setup

⚠️ Important: These installation instructions are based on the model repositories as of January 2025. Always check the original GitHub repositories for the most up-to-date installation instructions, as dependencies and requirements may change.

Based on dependency analysis, we can consolidate models into 5 main environments to minimize redundancy:

1. Standard Transformers Environment (for DNABERT1, INHERIT, GENA-LM, GROVER)

# Create environment - shared by multiple models
conda create -n dna_transformers_env python=3.9 -y
conda activate dna_transformers_env

# Install PyTorch with CUDA support
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

# Install core dependencies
pip install transformers==4.30.0 biopython pandas scikit-learn sentencepiece

# Test installation
python -c "import torch, transformers; print(f'PyTorch: {torch.__version__}, Transformers: {transformers.__version__}')"

Used by: DNABERT1, INHERIT, GENA-LM, GROVER

2. DNABERT2 Environment

# DNABERT2 has specific requirements per their GitHub
conda create -n dnabert2_env python=3.8 -y
conda activate dnabert2_env

# Install dependencies
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
pip install transformers accelerate

# Optional: Install triton for flash attention (from DNABERT2 repo)
# git clone https://github.com/openai/triton.git
# cd triton/python && pip install -e .

# Test installation
python -c "from transformers import AutoTokenizer; print('DNABERT2 ready')"

Repository: https://github.com/MAGICS-LAB/DNABERT_2

3. Nucleotide Transformer Environment

# Nucleotide Transformer requires specific JAX-based setup
conda create -n nucleotide_transformer_env python=3.9 -y
conda activate nucleotide_transformer_env

# Install PyTorch
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

# Install transformers from source (required for NT)
pip install --upgrade git+https://github.com/huggingface/transformers.git
pip install einops>=0.6.0 pandas scikit-learn

# Test installation
python -c "from transformers import AutoTokenizer; print('NT ready')"

Repository: https://github.com/instadeepai/nucleotide-transformer

4. ProkBERT Environment

# ProkBERT requires Python 3.10+ and has specific package versions
conda create -n prokbert_env python=3.10 -y
conda activate prokbert_env

# Recommended: Install directly from GitHub
pip install git+https://github.com/nbrg-ppcu/prokbert.git

# Alternative: Manual installation
# conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
# pip install transformers>=4.23 biopython h5py tables blosc2 accelerate

# Test installation
python -c "import prokbert; print('ProkBERT ready')"

Repository: https://github.com/nbrg-ppcu/prokbert

5. Advanced Models Environment (for HyenaDNA, EVO, Caduceus)

# These models require newer Python and similar advanced dependencies
conda create -n dna_advanced_env python=3.10 -y
conda activate dna_advanced_env

# Install PyTorch (check for latest compatible version)
conda install pytorch==2.1.0 torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

# Core dependencies
pip install transformers>=4.35.0 einops>=0.7.0 pandas scikit-learn

# Model-specific packages:
# For HyenaDNA (optional but recommended):
# pip install flash-attn --no-build-isolation

# For Caduceus (if using):
# pip install mamba-ssm causal-conv1d>=1.1.0

# For EVO: Ensure FlashAttention-2 compatible GPU (A100 required, V100 not supported)

# Test installation
python -c "import torch, transformers; print(f'Advanced env ready: PyTorch {torch.__version__}')"

Used by: HyenaDNA, EVO, Caduceus

Alternative: Docker Installation (Recommended for HyenaDNA)

# HyenaDNA provides a Docker image with all dependencies
git clone --recurse-submodules https://github.com/HazyResearch/hyena-dna.git
cd hyena-dna
docker pull hyenadna/hyena-dna:latest
docker run --gpus all -it -p80:3000 hyenadna/hyena-dna /bin/bash

Verifying Installations

# List created environments
conda env list | grep -E "dna_transformers|dnabert2|nucleotide|prokbert|dna_advanced"

# Quick test for all environments
for env in dna_transformers_env dnabert2_env nucleotide_transformer_env prokbert_env dna_advanced_env; do
    echo "Testing $env..."
    conda run -n $env python -c "import torch; print('$env: OK')" 2>/dev/null || echo "$env: FAILED"
done

Important Notes

1. Check Original Repositories: Installation requirements change frequently. Always verify with:

   • DNABERT1: https://github.com/jerryji1993/DNABERT
   • DNABERT2: https://github.com/MAGICS-LAB/DNABERT_2
   • Nucleotide Transformer: https://github.com/instadeepai/nucleotide-transformer
   • ProkBERT: https://github.com/nbrg-ppcu/prokbert
   • GROVER: Check HuggingFace model card
   • GENA-LM: https://github.com/AIRI-Institute/GENA_LM
   • HyenaDNA: https://github.com/HazyResearch/hyena-dna
   • EVO: https://github.com/evo-design/evo
   • Caduceus: https://github.com/kuleshov-group/caduceus

2. GPU Requirements:

   • Most models: CUDA-capable GPU with 8-16GB memory
   • HyenaDNA (1M context): A100 GPU recommended
   • EVO: A100 required (FlashAttention-2 not supported on V100)

3. Environment Mapping: Update config/models.yaml if you use different environment names than suggested above.

Requirements

• Python >= 3.8 (3.10+ for EVO and Caduceus)
• PyTorch >= 2.0.0
• Transformers >= 4.30.0
• CUDA-capable GPU (recommended, 16GB+ for large models)
• Disk space: ~50GB for all model weights

Project Structure

dna_benchmark/
├── config/                 # Configuration files
│   ├── benchmark.yaml      # Benchmark settings
│   └── models.yaml         # Model configurations
├── data/                   # Data directories
│   ├── raw/               # Raw sequence data
│   ├── processed/         # Processed datasets
│   └── splits/            # Train/val/test splits
├── src/
│   └── models/            # Model implementations
│       ├── base_model.py  # Abstract base class
│       ├── dnabert1.py    # DNABERT v1
│       ├── dnabert2.py    # DNABERT v2
│       ├── nucleotide_transformer.py
│       ├── prokbert.py
│       ├── grover.py
│       ├── gena_lm.py
│       ├── inherit.py
│       ├── hyenadna.py
│       ├── evo.py
│       └── README.md      # Model details
├── results/               # Output directories
│   ├── metrics/          # Evaluation metrics
│   ├── runs/             # Training logs
│   └── visualizations/   # Plots and figures
├── slurm/                # SLURM job templates
├── scripts/              # Utility scripts
└── tests/                # Unit tests

Supported Models

Quick Start

1. Load a Model

from src.models.dnabert2 import DNABert2Model

# Initialize model with config
config = {
    'pretrained_path': 'zhihan1996/DNABERT-2-117M',
    'num_labels': 2,
    'max_length': 512
}

model = DNABert2Model(config)
model.load_pretrained()

2. Get Embeddings

# Single sequence
sequence = "ATCGATCGATCGATCG"
embeddings = model.get_embeddings(sequence)

# Multiple sequences
sequences = ["ATCGATCG", "GCTAGCTA", "TTTTAAAA"]
embeddings = model.get_embeddings(sequences)

3. Make Predictions

# Predict prophage vs non-prophage
predictions = model.predict(sequences)
print(predictions['labels'])  # Predicted classes
print(predictions['probabilities'])  # Class probabilities

4. Fine-tune on Custom Data

# Fine-tune model on your dataset
metrics = model.fine_tune(
    train_dataset=train_data,
    val_dataset=val_data,
    epochs=10,
    batch_size=16,
    learning_rate=2e-5
)

# Save fine-tuned model
model.save_model("path/to/save/model")

Configuration

Benchmark Configuration (config/benchmark.yaml)

datasets:
  prophage_bench_v1:
    path: "data/processed/prophage_bench_v1"
    splits:
      train: 0.7
      val: 0.15
      test: 0.15
    sequence_length: 1000
    overlap: 200
    
evaluation:
  metrics:
    - accuracy
    - f1
    - precision
    - recall
    - auroc
    - auprc
    - mcc

Model Configuration (config/models.yaml)

Each model has specific configuration parameters:

models:
  dnabert2:
    class_name: "DNABert2Model"
    pretrained_path: "zhihan1996/DNABERT-2-117M"
    tokenizer: "bpe"
    max_length: 512

Running Benchmarks

Local Execution

# Run benchmark for a specific model
python scripts/run_benchmark.py --model dnabert2 --dataset prophage_bench_v1

# Run benchmark for all models
python scripts/run_benchmark.py --model all --dataset prophage_bench_v1

SLURM Cluster Execution

# Submit job to SLURM
sbatch slurm/templates/benchmark_job.sh

Dataset Format

The framework expects datasets in the following format:

data/processed/dataset_name/
├── train.csv
├── val.csv
└── test.csv

Each CSV file should contain:

• sequence: DNA sequence (string)
• label: Class label (integer)
• id: Unique identifier (optional)

Evaluation Metrics

The framework computes the following metrics:

• Accuracy: Overall classification accuracy
• F1 Score: Harmonic mean of precision and recall
• Precision: True positives / (True positives + False positives)
• Recall: True positives / (True positives + False negatives)
• AUROC: Area under the ROC curve
• AUPRC: Area under the Precision-Recall curve
• MCC: Matthews Correlation Coefficient - a balanced measure for binary classification

Advanced Features

Sequence Chunking

For models with limited sequence length, the framework automatically chunks long sequences:

# Sequences longer than max_length are split with overlap
chunked_seqs, chunk_counts = model.split_long_sequences(
    sequences, 
    overlap=100
)

# Predictions are aggregated across chunks
predictions = model.aggregate_chunk_predictions(
    chunk_predictions,
    chunk_counts,
    method='mean'  # or 'max', 'vote'
)

Custom Model Integration

To add a new model, extend the BaseDNAModel class:

from src.models.base_model import BaseDNAModel

class CustomDNAModel(BaseDNAModel):
    def load_pretrained(self, path=None):
        # Load model weights
        pass
    
    def get_embeddings(self, sequences):
        # Generate embeddings
        pass
    
    def predict(self, sequences):
        # Make predictions
        pass
    
    def fine_tune(self, train_dataset, val_dataset=None, **kwargs):
        # Fine-tune model
        pass

Citation

If you use this framework in your research, please cite:

@software{dna_benchmark,
  title = {DNA Language Model Benchmarking Framework},
  author = {LeAnn Lindsey},
  year = {2024},
  url = {https://github.com/leannmlindsey/dna_benchmark}
}

Contributing

We welcome contributions! Please see CONTRIBUTING.md for guidelines.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

This framework builds upon the following excellent DNA language models. Please cite the original papers and repositories if you use these models in your research:

• DNABERT: GitHub Repository | HuggingFace
• DNABERT-2: GitHub Repository | HuggingFace
• Nucleotide Transformer: GitHub Repository | HuggingFace
• ProkBERT: GitHub Repository | HuggingFace
• GROVER: Zenodo | HuggingFace
• GENA-LM: GitHub Repository | HuggingFace
• INHERIT: GitHub Repository
• HyenaDNA: GitHub Repository | HuggingFace
• EVO: GitHub Repository (Evo-1) | GitHub Repository (Evo-2) | HuggingFace
• Caduceus: GitHub Repository | HuggingFace

Important: When using any of these models, please ensure you cite the original authors' work as specified in their respective repositories.

Contact

For questions and support, please open an issue on GitHub or contact [email protected].