A fork of Caduceus expanded for finetuning using any dataset with binary classification
Python
0
81 commits
updated Jul 7, 2026
Note: This is a fork of the original Caduceus repository with added support for generic CSV-based binary/multiclass classification tasks. This allows you to fine-tune Caduceus (or Mamba/Hyena) on your own DNA sequence classification datasets.
See how Caduceus performs on the full LAMBDA genome-wide evaluation set, and compare it interactively against every other benchmarked model, on the LAMBDA benchmark dashboard:
https://leannmlindsey.github.io/lambda-benchmark/
This fork adds the ability to fine-tune on any classification task using simple CSV files.
Create a directory containing three CSV files with sequence and label columns:
my_dataset/
├── train.csv
├── dev.csv
└── test.csv
Each CSV should have this format:
sequence,label
ACGTACGTACGT...,0
TGCATGCATGCA...,1
GGCCAATTGGCC...,0
sequence: DNA sequence (A, C, G, T, N characters)label: Integer class label (0, 1 for binary; 0, 1, 2, ... for multiclass)Note: The HuggingFace checkpoints are not directly compatible with this fine-tuning pipeline (see this issue for details). You need to pretrain your own model using this codebase, or use an existing checkpoint that was trained with this repository.
To pretrain a Caduceus model, follow the Pretraining on Human Reference Genome instructions in the original README below.
After pretraining, locate these two files in your pretraining output directory:
outputs/pretrain/hg38/<your_pretraining_run_name>/
├── model_config.json # <-- CONFIG_PATH: Model architecture configuration
└── checkpoints/
└── last.ckpt # <-- PRETRAINED_PATH: Model weights
You will need the full paths to both model_config.json and last.ckpt for fine-tuning.
Important: You only need to edit one file:
slurm_scripts/wrapper_run_csv_binary.sh. The config files inconfigs/do not need to be modified.
Edit slurm_scripts/wrapper_run_csv_binary.sh with your paths:
# === REQUIRED: Dataset Configuration ===
export DATA_DIR="/path/to/my_dataset" # Directory containing train.csv, dev.csv, test.csv
export DATASET_NAME="my_dataset" # Name for output directory (no spaces)
# === REQUIRED: Model Configuration (from your pretraining output) ===
export CONFIG_PATH="/path/to/outputs/pretrain/hg38/<run_name>/model_config.json"
export PRETRAINED_PATH="/path/to/outputs/pretrain/hg38/<run_name>/checkpoints/last.ckpt"
# === Model Type (must match what you pretrained) ===
export MODEL="caduceus" # Options: hyena, mamba, caduceus
export MODEL_NAME="dna_embedding_caduceus" # Options: dna_embedding, dna_embedding_mamba, dna_embedding_caduceus
# === Hyperparameters (adjust as needed) ===
export LR="6e-4" # Learning rate
export BATCH_SIZE="32" # Batch size (reduce if OOM)
export MAX_LENGTH="1024" # Max sequence length (truncates longer sequences)
export MAX_EPOCHS="20" # Training epochs (early stopping may trigger sooner)
export D_OUTPUT="2" # Number of classes (2 for binary)
Learning Rate: Start with
6e-4. If training produces NaN loss, reduce to1e-4or5e-5. Some random seeds may require lower learning rates for stability.
Important: The RC parameters you use for fine-tuning must match the model variant you pretrained. Using mismatched settings will result in poor performance or errors.
Caduceus has special handling for reverse complement (RC) sequences. The parameters depend on which model variant you pretrained:
Caduceus Variants:
Parameter Definitions:
| Parameter | Description |
|---|---|
CONJOIN_TEST | If true, run both forward + RC sequences at test time and average predictions |
CONJOIN_TRAIN_DECODER | If true, decoder expects input with shape (..., 2) and combines both channels |
RC_AUG | If true, randomly apply RC augmentation during training |
Required Settings by Pretrained Model Type:
| If you pretrained... | CONJOIN_TEST | CONJOIN_TRAIN_DECODER | RC_AUG |
|---|---|---|---|
| Caduceus-Ph | true | false | false |
| Caduceus-PS | false | true | false |
| Mamba | false | false | true |
| Hyena | false | false | true |
The default settings in the SLURM scripts are configured for Caduceus-Ph. If you pretrained a different model variant, update these parameters in wrapper_run_csv_binary.sh:
# === Reverse Complement Settings (must match your pretrained model) ===
export CONJOIN_TEST="true" # true for Caduceus-Ph, false for others
export CONJOIN_TRAIN_DECODER="false" # true for Caduceus-PS, false for others
export RC_AUG="false" # true for Mamba/Hyena, false for Caduceus
Caduceus uses Weights & Biases for experiment tracking and logging.
Setup:
pip install wandbwandb login (enter your API key when prompted)Configuration:
By default, wandb logging is disabled in the SLURM scripts (wandb=null). To enable it:
# In your training command, remove or change the wandb=null line:
python -m train \
experiment=csv_binary \
... \
wandb.project="my_project" \
wandb.name="my_experiment"
Wandb options:
wandb=null - Disable wandb logging (default in SLURM scripts)wandb.mode=online - Log to wandb cloud in real-timewandb.mode=offline - Log locally, sync later with wandb syncwandb.mode=disabled - Disable loggingwandb.project="name" - Set project namewandb.name="name" - Set run nameBy default, training stops when the gap between training and validation accuracy exceeds a threshold, which helps prevent overfitting.
Expected Training Behavior:
With large datasets (e.g., 60k+ samples), training typically completes in 3-5 epochs before the generalization gap triggers early stopping. The max_epochs=20 setting is a ceiling, not a target. Training longer leads to overfitting where the model memorizes training data rather than learning generalizable patterns.
| Dataset Size | Typical Epochs | Notes |
|---|---|---|
| < 10k samples | 10-20 | May need more epochs to learn patterns |
| 10k-50k samples | 5-10 | Moderate training time |
| 50k+ samples | 3-5 | Early stopping usually triggers quickly |
How it works:
train/accuracy - val/accuracy (the generalization gap)max_gap (default: 0.03) for patience consecutive epochsmin_epochs before checking (default: 2 epochs)Configuration (in configs/callbacks/generalization_gap.yaml):
generalization_gap:
max_gap: 0.03 # Stop if train_acc - val_acc > this value
train_metric: "train/accuracy"
val_metric: "val/accuracy"
min_epochs: 2 # Don't check gap until this many epochs
patience: 2 # Consecutive epochs gap can exceed threshold
verbose: True
To adjust the gap threshold via command line:
python -m train experiment=csv_binary \
callbacks.generalization_gap.max_gap=0.05 \
callbacks.generalization_gap.patience=5 \
...
Alternative: Standard early stopping (monitors val/accuracy only)
Standard early stopping monitors only validation accuracy and stops when it plateaus, without considering the train/val gap. To use it instead, edit configs/pipeline/csv_binary.yaml:
- /callbacks: [base, checkpoint, test_results, early_stopping] # swap generalization_gap for early_stopping
Standard early stopping config (configs/callbacks/early_stopping.yaml):
early_stopping:
monitor: ${train.monitor} # val/accuracy
patience: 10 # epochs with no improvement
To disable all early stopping:
Edit configs/pipeline/csv_binary.yaml:
- /callbacks: [base, checkpoint, test_results] # remove stopping callback
To measure the contribution of pretraining, you can train with randomly initialized weights (same architecture, no pretrained weights). This provides a baseline to compare against the pretrained model.
Usage:
python -m train experiment=csv_binary \
train.random_init=true \
dataset.data_dir=/path/to/data \
+model.config_path=/path/to/model_config.json \
...
When train.random_init=true:
pretrained_model_path is specified)Comparing Results:
Run two experiments with the same seed and hyperparameters:
train.random_init=false (default)train.random_init=trueThe difference in metrics shows the "embedding power" gained from pretraining.
cd slurm_scripts
bash wrapper_run_csv_binary.sh
Or submit directly:
sbatch --export=ALL,DATA_DIR=/path/to/data,CONFIG_PATH=/path/to/config.json,PRETRAINED_PATH=/path/to/ckpt.ckpt,DATASET_NAME=mydata run_csv_binary.sh
Output Directory Structure:
outputs/downstream/csv_binary/{DATASET_NAME}/{MODEL}_lr-{LR}_batch_size-{BS}_rc_aug-{RC}/seed-{SEED}/
├── test_results.json # Comprehensive metrics computed on full test set
├── test_predictions.npz # Raw predictions for further analysis
└── checkpoints/ # Model checkpoints
How Test Metrics Are Computed:
This fork includes a custom test evaluation callback (src/callbacks/test_results.py) that addresses a common issue with batch-level metric averaging. Metrics like MCC (Matthews Correlation Coefficient) can be incorrectly calculated when averaged across batches, especially if individual batches have imbalanced class distributions.
Our solution:
test_results.json with full traceabilityOutput Files:
| File | Description |
|---|---|
test_results.json | All metrics computed on the full test set, plus paths to data and checkpoint |
test_predictions.npz | NumPy archive containing logits, probabilities, predictions, and labels arrays for custom analysis |
The test_results.json contains:
{
"eval_loss": 0.258,
"eval_accuracy": 0.931,
"eval_precision": 0.957,
"eval_recall": 0.904,
"eval_f1": 0.930,
"eval_mcc": 0.864,
"eval_sensitivity": 0.904,
"eval_specificity": 0.959,
"eval_auc": 0.983,
"eval_runtime": 30.15,
"eval_samples_per_second": 226.4,
"eval_steps_per_second": 3.55,
"epoch": 100.0,
"checkpoint_path": "/path/to/checkpoint.ckpt",
"train_data_path": "/path/to/train.csv",
"dev_data_path": "/path/to/dev.csv",
"test_data_path": "/path/to/test.csv"
}
Loading predictions for custom analysis:
import numpy as np
data = np.load('test_predictions.npz')
logits = data['logits'] # Raw model outputs
probs = data['probabilities'] # Softmax probabilities
preds = data['predictions'] # Predicted class labels
labels = data['labels'] # True labels
Two SLURM scripts are provided in slurm_scripts/ for running on HPC clusters (configured for NIH Biowulf):
File Locations:
slurm_scripts/
├── wrapper_run_csv_binary.sh # <-- EDIT THIS FILE with your paths
└── run_csv_binary.sh # Main job script (no edits needed)
wrapper_run_csv_binary.sh - Configuration wrapper (EDIT THIS)
sbatch to submit the jobrun_csv_binary.sh - Main SLURM job script (no edits needed)
caduceus_env)What to Edit in wrapper_run_csv_binary.sh:
| Variable | Description | Example |
|---|---|---|
DATA_DIR | Path to folder with train.csv, dev.csv, test.csv | /data/user/my_dataset |
DATASET_NAME | Name for output folder (no spaces) | phage_detection |
CONFIG_PATH | Path to model_config.json from pretraining | /data/user/pretrain/model_config.json |
PRETRAINED_PATH | Path to checkpoint from pretraining | /data/user/pretrain/checkpoints/last.ckpt |
MODEL | Model type | caduceus, mamba, or hyena |
MODEL_NAME | Model registry name | dna_embedding_caduceus |
LR | Learning rate | 6e-4 |
BATCH_SIZE | Batch size (reduce if OOM) | 32 |
MAX_LENGTH | Max sequence length | 1024 |
MAX_EPOCHS | Training epochs | 100 |
D_OUTPUT | Number of classes | 2 |
Modifying SLURM Resources (if needed):
If you need to change GPU type, memory, or time limit, edit run_csv_binary.sh:
#SBATCH --partition=gpu # Partition name
#SBATCH --gres=gpu:a100:1 # GPU type and count
#SBATCH --mem=64g # Memory
#SBATCH --cpus-per-task=8 # CPU cores
#SBATCH --time=24:00:00 # Time limit
Changing the Conda Environment:
If your conda environment has a different name, edit this line in run_csv_binary.sh:
source activate caduceus_env # Change 'caduceus_env' to your env name
Extract embeddings from a model and analyze their quality using linear probes, silhouette scores, PCA visualization, and a 3-layer neural network classifier.
Using the SLURM wrapper (recommended):
slurm_scripts/wrapper_run_embedding_analysis.sh:export CSV_DIR="/path/to/csv/data" # Directory with train.csv, dev.csv, test.csv
export CONFIG_PATH="/path/to/model_config.json"
export CHECKPOINT_PATH="/path/to/checkpoint.ckpt"
export INCLUDE_RANDOM_BASELINE="true" # Compare against random init baseline
cd slurm_scripts
bash wrapper_run_embedding_analysis.sh
Or run interactively (useful for debugging or in an interactive SLURM session):
bash run_embedding_analysis_interactive.sh
Direct command:
python -m src.embedding_analysis \
--csv_dir="/path/to/csv/data" \
--checkpoint_path="/path/to/checkpoint.ckpt" \
--config_path="/path/to/config.json" \
--output_dir="./outputs/embedding_analysis" \
--pooling="mean" \
--batch_size=32 \
--include_random_baseline
Outputs:
| File | Description |
|---|---|
embeddings_pretrained.npz | Extracted embeddings for train/val/test sets |
test_predictions_pretrained.csv | Test set predictions with probabilities |
pca_visualization_pretrained.png | PCA plot showing class separation |
three_layer_nn_pretrained.pt | Trained 3-layer NN classifier weights |
embedding_analysis_results.json | All metrics (see below) |
If --include_random_baseline is set, additional files with _random suffix are created for comparison.
Test predictions CSV columns:
sequence: Original DNA sequencelabel: Ground truth labellinear_probe_pred: Linear probe predicted labellinear_probe_prob: Linear probe probability (class 1)nn_pred: 3-layer NN predicted labelnn_prob: 3-layer NN probability (class 1)Metrics in embedding_analysis_results.json:
Run inference on a CSV file to get predictions with probabilities for threshold analysis.
Single file inference:
python -m src.inference \
--input_csv="/path/to/test.csv" \
--checkpoint_path="/path/to/checkpoint.ckpt" \
--config_path="/path/to/config.json" \
--output_csv="/path/to/predictions.csv" \
--threshold=0.5 \
--conjoin_test \
--save_metrics
Batch inference on multiple files:
# input_files.txt
/path/to/dataset1.csv
/path/to/dataset2.csv
/path/to/dataset3.csv
slurm_scripts/wrapper_run_batch_inference.sh:INPUT_LIST="/path/to/input_files.txt"
OUTPUT_DIR="/path/to/output_directory"
CONFIG_PATH="/path/to/model_config.json"
CHECKPOINT_PATH="/path/to/checkpoint.ckpt"
cd slurm_scripts
bash wrapper_run_batch_inference.sh
This submits a separate SLURM job for each input file. All predictions and logs are saved to the output directory.
Output CSV columns:
sequence: Original sequencelabel: Original label (if present)prob_0, prob_1: Class probabilitiespred_label: Predicted labelMetrics JSON (when --save_metrics is used):
Key options:
--threshold: Custom classification threshold for sensitivity/specificity tradeoff analysis (default: 0.5)--conjoin_test: Use post-hoc reverse complement averaging (recommended for Caduceus-Ph)--save_metrics: Calculate and save metrics to JSON if labels are present| File | Description |
|---|---|
src/dataloaders/datasets/csv_dataset.py | PyTorch Dataset for loading CSV files |
src/dataloaders/genomics.py | Added CSVDatasetLoader class |
src/callbacks/test_results.py | Callback for computing metrics with scikit-learn |
src/embedding_analysis.py | Extract embeddings, linear probe, PCA, 3-layer NN |
src/inference.py | Run inference with probability outputs |
configs/experiment/csv_binary.yaml | Experiment config for CSV classification |
configs/pipeline/csv_binary.yaml | Pipeline config |
configs/dataset/csv_dataset.yaml | Dataset config |
slurm_scripts/run_csv_binary.sh | SLURM job script for fine-tuning |
slurm_scripts/wrapper_run_csv_binary.sh | Wrapper for fine-tuning job submission |
slurm_scripts/run_csv_binary_interactive.sh | Interactive fine-tuning (no sbatch) |
slurm_scripts/run_embedding_analysis.sh | SLURM job script for embedding analysis |
slurm_scripts/wrapper_run_embedding_analysis.sh | Wrapper for embedding analysis submission |
slurm_scripts/run_embedding_analysis_interactive.sh | Interactive embedding analysis |
slurm_scripts/run_inference.sh | SLURM job script for inference |
slurm_scripts/wrapper_run_batch_inference.sh | Wrapper for batch inference configuration |
slurm_scripts/submit_batch_inference.sh | Batch submission for multiple inference jobs |
The remainder of this README is from the original Caduceus repository.
[Blog] | [arXiv] | [HuggingFace 🤗]
This repository contains code for reproducing the results in the paper "Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling," Schiff et al. (2024).
We have uploaded a pre-trained Caduceus model to the Huggingface hub. The available models are:
You can either use the pre-trained model directly within your trainer scripts or modify the config that initializes the model.
To use the pre-trained model for masked language modeling, use the following snippet:
from transformers import AutoModelForMaskedLM, AutoTokenizer
# See the `Caduceus` collection page on the hub for list of available models.
model_name = "kuleshov-group/caduceus-ph_seqlen-131k_d_model-256_n_layer-16"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForMaskedLM.from_pretrained(model_name)
Alternatively, you can instantiate a model from scratch to train on your own data as follows:
from transformers import AutoConfig, AutoModelForMaskedLM
# Add any config overrides here, see the `config.json` file on the hub for details.
config_overrides = {}
# See the `Caduceus` collection page on the hub for list of available models.
config = AutoConfig.from_pretrained(
"kuleshov-group/caduceus-ph_seqlen-131k_d_model-256_n_layer-16",
**config_overrides,
)
model = AutoModelForMaskedLM.from_config(config)
To get started, create a conda environment containing the required dependencies.
conda env create -f caduceus_env.yml
Activate the environment.
conda activate caduceus_env
Create the following directories to store saved models and slurm logs:
mkdir outputs
mkdir watch_folder
Below, we describe the steps required for reproducing the experiments in the paper.
Throughout, the main entry point for running experiments is the train.py script.
We also provide sample slurm scripts for launching pre-training and downstream fine-tuning experiments in the slurm_scripts/ directory.
(Data downloading instructions are copied from HyenaDNA repo)
First, download the Human Reference Genome data.
It's comprised of 2 files, 1 with all the sequences (the .fasta file), and with the intervals we use (.bed file).
The file structure should look like
data
|-- hg38/
|-- hg38.ml.fa
|-- human-sequences.bed
Download fasta (.fa format) file (of the entire human genome) into ./data/hg38.
~24 chromosomes in the whole genome (merged into 1 file), each chromosome is a continuous sequence, basically.
Then download the .bed file with sequence intervals (contains chromosome name, start, end, split, which then allow you to retrieve from the fasta file).
mkdir -p data/hg38/
curl https://storage.googleapis.com/basenji_barnyard2/hg38.ml.fa.gz > data/hg38/hg38.ml.fa.gz
gunzip data/hg38/hg38.ml.fa.gz # unzip the fasta file
curl https://storage.googleapis.com/basenji_barnyard2/sequences_human.bed > data/hg38/human-sequences.bed
Launch pretraining run using the command line
python -m train \
experiment=hg38/hg38 \
callbacks.model_checkpoint_every_n_steps.every_n_train_steps=500 \
dataset.max_length=1024 \
dataset.batch_size=1024 \
dataset.mlm=true \
dataset.mlm_probability=0.15 \
dataset.rc_aug=false \
model=caduceus \
model.config.d_model=128 \
model.config.n_layer=4 \
model.config.bidirectional=true \
model.config.bidirectional_strategy=add \
model.config.bidirectional_weight_tie=true \
model.config.rcps=true \
optimizer.lr="8e-3" \
train.global_batch_size=1024 \
trainer.max_steps=10000 \
+trainer.val_check_interval=10000 \
wandb=null
or alternatively, if using a cluster that has slurm installed, adapt the scripts below:
slurm_scripts
|-- run_pretrain_caduceus.sh
|-- run_pretrain_hyena.sh
|-- run_pretrain_mamba.sh
and run the training as a batch job:
cd slurm_scripts
sbatch run_pretrain_caduceus.sh
The GenomicBenchmarks presented in Grešová et al. (2023) is comprised of 8 classification tasks.
We can launch a downstream fine-tuning run on one of the tasks using the sample command below:
python -m train \
experiment=hg38/genomic_benchmark \
callbacks.model_checkpoint_every_n_steps.every_n_train_steps=5000 \
dataset.dataset_name="dummy_mouse_enhancers_ensembl" \
dataset.train_val_split_seed=1 \
dataset.batch_size=256 \
dataset.rc_aug=false \
+dataset.conjoin_train=false \
+dataset.conjoin_test=false \
loader.num_workers=2 \
model=caduceus \
model._name_=dna_embedding_caduceus \
+model.config_path="<path to model_config.json>" \
+model.conjoin_test=false \
+decoder.conjoin_train=true \
+decoder.conjoin_test=false \
optimizer.lr="1e-3" \
trainer.max_epochs=10 \
train.pretrained_model_path="<path to .ckpt file>" \
wandb=null
This sample run will fine-tune a pre-trained Caduceus-PS model on the dummy_mouse_enhancers_ensembl task.
Note some of the additional arguments present here, relative to the pre-training command from above:
model.config_path contains the path model config that was saved during pre-training.
This will be saved to the run directory of the pre-training experiment.train.pretrained_model_path contains the path to the pre-trained model checkpoint.dataset.conjoin_train determines whether the dataset will return a single sequence (dataset.conjoin_train=false) or the concatenation of a sequence and its reverse complement along dim=-1, during downstream fine-tuning training.dataset.conjoin_test is the same as above, but for inference (e.g., validation / test).decoder.conjoin_train determines whether the prediction head (a mean pooling and linear projection in the case of the Genomics Benchmark) is expecting an input tensor of shape (batch_size, seq_len, d_model) or (batch_size, seq_len, d_model, 2) during downstream fine-tuning training.
When set to true the decoder is run on input[..., 0] and input[..., 1] and the results are averaged to produce the final prediction.decoder.conjoin_test is the same as above, but for inference (e.g., validation / test).Note this benchmark only contains a training and test split for each task.
Therefore, to have a more principled evaluation, we randomly split the training data into training and validation sets (90/10) using the dataset.train_val_split_seed argument.
We perform early stopping on validation metric (accuracy) and repeat this for 5 random seeds.
As with pre-training, we can also launch the fine-tuning run as a batch job using the provided run_genomic_benchmark.sh script.
We also provide a helper shell script wrapper_run_genomics.sh that can be used to launch multiple fine-tuning runs in parallel.
Finally, the run_genomics_benchmark_cnn.sh script can be used to train the CNN baseline for this experiment from scratch on the downstream tasks.
The Nucleotide Transformer suite of tasks was proposed in Dalla-Torre et al. (2023). The data is available on HuggingFace: InstaDeepAI/nucleotide_transformer_downstream_tasks.
We can launch a downstream fine-tuning run on one of the tasks using the sample command below:
python -m train \
experiment=hg38/nucleotide_transformer \
callbacks.model_checkpoint_every_n_steps.every_n_train_steps=5000 \
dataset.dataset_name="${task}" \
dataset.train_val_split_seed=${seed} \
dataset.batch_size=${batch_size} \
dataset.rc_aug="${rc_aug}" \
+dataset.conjoin_test="${CONJOIN_TEST}" \
loader.num_workers=2 \
model._name_=dna_embedding_caduceus \
+model.config_path="<path to model_config.json>" \
+model.conjoin_test=false \
+decoder.conjoin_train=true \
+decoder.conjoin_test=false \
optimizer.lr="1e-3" \
trainer.max_epochs=10 \
train.pretrained_model_path="<path to .ckpt file>" \
trainer.max_epochs=20 \
wandb=null
We can also launch as batch jobs (see run_nucleotide_transformer.sh and wrapper_run_nucleotide_transformer.sh for details).
This task comes from the recently proposed Long Range Benchmark (LRB) in Kao et al., 2023. The data is available on HuggingFace: InstaDeepAI/genomics-long-range-benchmark. For this task we fit a model to the pre-trained and frozen embeddings of the DNA language models. Therefore, to perform the evaluation, we proceed in 2 steps:
vep_embeddings.py script to extract the embeddings from the pre-trained model.
See the example below:torchrun \
--standalone \
--nnodes=1 \
--nproc-per-node=8 \
vep_embeddings.py \
--num_workers=2 \
--seq_len=131072 \
--bp_per_token=1 \
--embed_dump_batch_size=1 \
--name="caduceus-ps_downstream-seqlen=131k" \
--model_name_or_path="kuleshov-group/caduceus-ps_seqlen-131k_d_model-256_n_layer-16" \
--rcps
The --rcps flag is used to indicate that the model is reverse-complement equivariant.
When using other models, set this flag to false with --no-rcps.
To speed this step up, this script utilizes torch distributed data parallelism.
Please refer to the slurm script provided in slurm_scripts/dump_vep_embeddings.sh
to launch this step as a batch job.
vep_svm.ipynb.If you find our work useful, please cite our paper using the following:
@article{schiff2024caduceus,
title={Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling},
author={Schiff, Yair and Kao, Chia-Hsiang and Gokaslan, Aaron and Dao, Tri and Gu, Albert and Kuleshov, Volodymyr},
journal={arXiv preprint arXiv:2403.03234},
year={2024}
}
This repository is adapted from the HyenaDNA repo and leverages much of the training, data loading, and logging infrastructure defined there. HyenaDNA was originally derived from the S4 and Safari repositories.
We would like to thank Evan Trop and the InstaDeep team for useful discussions about the Nucleotide Transformer leaderboard and the Long Range Benchmark task.
Finally, we would like to thank MosaicML for providing compute resources for some of the pre-training experiments.
A fork of Caduceus expanded for finetuning using any dataset with binary classification
Python
0
81 commits
updated Jul 7, 2026
Note: This is a fork of the original Caduceus repository with added support for generic CSV-based binary/multiclass classification tasks. This allows you to fine-tune Caduceus (or Mamba/Hyena) on your own DNA sequence classification datasets.
See how Caduceus performs on the full LAMBDA genome-wide evaluation set, and compare it interactively against every other benchmarked model, on the LAMBDA benchmark dashboard:
https://leannmlindsey.github.io/lambda-benchmark/
This fork adds the ability to fine-tune on any classification task using simple CSV files.
Create a directory containing three CSV files with sequence and label columns:
my_dataset/
├── train.csv
├── dev.csv
└── test.csv
Each CSV should have this format:
sequence,label
ACGTACGTACGT...,0
TGCATGCATGCA...,1
GGCCAATTGGCC...,0
sequence: DNA sequence (A, C, G, T, N characters)label: Integer class label (0, 1 for binary; 0, 1, 2, ... for multiclass)Note: The HuggingFace checkpoints are not directly compatible with this fine-tuning pipeline (see this issue for details). You need to pretrain your own model using this codebase, or use an existing checkpoint that was trained with this repository.
To pretrain a Caduceus model, follow the Pretraining on Human Reference Genome instructions in the original README below.
After pretraining, locate these two files in your pretraining output directory:
outputs/pretrain/hg38/<your_pretraining_run_name>/
├── model_config.json # <-- CONFIG_PATH: Model architecture configuration
└── checkpoints/
└── last.ckpt # <-- PRETRAINED_PATH: Model weights
You will need the full paths to both model_config.json and last.ckpt for fine-tuning.
Important: You only need to edit one file:
slurm_scripts/wrapper_run_csv_binary.sh. The config files inconfigs/do not need to be modified.
Edit slurm_scripts/wrapper_run_csv_binary.sh with your paths:
# === REQUIRED: Dataset Configuration ===
export DATA_DIR="/path/to/my_dataset" # Directory containing train.csv, dev.csv, test.csv
export DATASET_NAME="my_dataset" # Name for output directory (no spaces)
# === REQUIRED: Model Configuration (from your pretraining output) ===
export CONFIG_PATH="/path/to/outputs/pretrain/hg38/<run_name>/model_config.json"
export PRETRAINED_PATH="/path/to/outputs/pretrain/hg38/<run_name>/checkpoints/last.ckpt"
# === Model Type (must match what you pretrained) ===
export MODEL="caduceus" # Options: hyena, mamba, caduceus
export MODEL_NAME="dna_embedding_caduceus" # Options: dna_embedding, dna_embedding_mamba, dna_embedding_caduceus
# === Hyperparameters (adjust as needed) ===
export LR="6e-4" # Learning rate
export BATCH_SIZE="32" # Batch size (reduce if OOM)
export MAX_LENGTH="1024" # Max sequence length (truncates longer sequences)
export MAX_EPOCHS="20" # Training epochs (early stopping may trigger sooner)
export D_OUTPUT="2" # Number of classes (2 for binary)
Learning Rate: Start with
6e-4. If training produces NaN loss, reduce to1e-4or5e-5. Some random seeds may require lower learning rates for stability.
Important: The RC parameters you use for fine-tuning must match the model variant you pretrained. Using mismatched settings will result in poor performance or errors.
Caduceus has special handling for reverse complement (RC) sequences. The parameters depend on which model variant you pretrained:
Caduceus Variants:
Parameter Definitions:
| Parameter | Description |
|---|---|
CONJOIN_TEST | If true, run both forward + RC sequences at test time and average predictions |
CONJOIN_TRAIN_DECODER | If true, decoder expects input with shape (..., 2) and combines both channels |
RC_AUG | If true, randomly apply RC augmentation during training |
Required Settings by Pretrained Model Type:
| If you pretrained... | CONJOIN_TEST | CONJOIN_TRAIN_DECODER | RC_AUG |
|---|---|---|---|
| Caduceus-Ph | true | false | false |
| Caduceus-PS | false | true | false |
| Mamba | false | false | true |
| Hyena | false | false | true |
The default settings in the SLURM scripts are configured for Caduceus-Ph. If you pretrained a different model variant, update these parameters in wrapper_run_csv_binary.sh:
# === Reverse Complement Settings (must match your pretrained model) ===
export CONJOIN_TEST="true" # true for Caduceus-Ph, false for others
export CONJOIN_TRAIN_DECODER="false" # true for Caduceus-PS, false for others
export RC_AUG="false" # true for Mamba/Hyena, false for Caduceus
Caduceus uses Weights & Biases for experiment tracking and logging.
Setup:
pip install wandbwandb login (enter your API key when prompted)Configuration:
By default, wandb logging is disabled in the SLURM scripts (wandb=null). To enable it:
# In your training command, remove or change the wandb=null line:
python -m train \
experiment=csv_binary \
... \
wandb.project="my_project" \
wandb.name="my_experiment"
Wandb options:
wandb=null - Disable wandb logging (default in SLURM scripts)wandb.mode=online - Log to wandb cloud in real-timewandb.mode=offline - Log locally, sync later with wandb syncwandb.mode=disabled - Disable loggingwandb.project="name" - Set project namewandb.name="name" - Set run nameBy default, training stops when the gap between training and validation accuracy exceeds a threshold, which helps prevent overfitting.
Expected Training Behavior:
With large datasets (e.g., 60k+ samples), training typically completes in 3-5 epochs before the generalization gap triggers early stopping. The max_epochs=20 setting is a ceiling, not a target. Training longer leads to overfitting where the model memorizes training data rather than learning generalizable patterns.
| Dataset Size | Typical Epochs | Notes |
|---|---|---|
| < 10k samples | 10-20 | May need more epochs to learn patterns |
| 10k-50k samples | 5-10 | Moderate training time |
| 50k+ samples | 3-5 | Early stopping usually triggers quickly |
How it works:
train/accuracy - val/accuracy (the generalization gap)max_gap (default: 0.03) for patience consecutive epochsmin_epochs before checking (default: 2 epochs)Configuration (in configs/callbacks/generalization_gap.yaml):
generalization_gap:
max_gap: 0.03 # Stop if train_acc - val_acc > this value
train_metric: "train/accuracy"
val_metric: "val/accuracy"
min_epochs: 2 # Don't check gap until this many epochs
patience: 2 # Consecutive epochs gap can exceed threshold
verbose: True
To adjust the gap threshold via command line:
python -m train experiment=csv_binary \
callbacks.generalization_gap.max_gap=0.05 \
callbacks.generalization_gap.patience=5 \
...
Alternative: Standard early stopping (monitors val/accuracy only)
Standard early stopping monitors only validation accuracy and stops when it plateaus, without considering the train/val gap. To use it instead, edit configs/pipeline/csv_binary.yaml:
- /callbacks: [base, checkpoint, test_results, early_stopping] # swap generalization_gap for early_stopping
Standard early stopping config (configs/callbacks/early_stopping.yaml):
early_stopping:
monitor: ${train.monitor} # val/accuracy
patience: 10 # epochs with no improvement
To disable all early stopping:
Edit configs/pipeline/csv_binary.yaml:
- /callbacks: [base, checkpoint, test_results] # remove stopping callback
To measure the contribution of pretraining, you can train with randomly initialized weights (same architecture, no pretrained weights). This provides a baseline to compare against the pretrained model.
Usage:
python -m train experiment=csv_binary \
train.random_init=true \
dataset.data_dir=/path/to/data \
+model.config_path=/path/to/model_config.json \
...
When train.random_init=true:
pretrained_model_path is specified)Comparing Results:
Run two experiments with the same seed and hyperparameters:
train.random_init=false (default)train.random_init=trueThe difference in metrics shows the "embedding power" gained from pretraining.
cd slurm_scripts
bash wrapper_run_csv_binary.sh
Or submit directly:
sbatch --export=ALL,DATA_DIR=/path/to/data,CONFIG_PATH=/path/to/config.json,PRETRAINED_PATH=/path/to/ckpt.ckpt,DATASET_NAME=mydata run_csv_binary.sh
Output Directory Structure:
outputs/downstream/csv_binary/{DATASET_NAME}/{MODEL}_lr-{LR}_batch_size-{BS}_rc_aug-{RC}/seed-{SEED}/
├── test_results.json # Comprehensive metrics computed on full test set
├── test_predictions.npz # Raw predictions for further analysis
└── checkpoints/ # Model checkpoints
How Test Metrics Are Computed:
This fork includes a custom test evaluation callback (src/callbacks/test_results.py) that addresses a common issue with batch-level metric averaging. Metrics like MCC (Matthews Correlation Coefficient) can be incorrectly calculated when averaged across batches, especially if individual batches have imbalanced class distributions.
Our solution:
test_results.json with full traceabilityOutput Files:
| File | Description |
|---|---|
test_results.json | All metrics computed on the full test set, plus paths to data and checkpoint |
test_predictions.npz | NumPy archive containing logits, probabilities, predictions, and labels arrays for custom analysis |
The test_results.json contains:
{
"eval_loss": 0.258,
"eval_accuracy": 0.931,
"eval_precision": 0.957,
"eval_recall": 0.904,
"eval_f1": 0.930,
"eval_mcc": 0.864,
"eval_sensitivity": 0.904,
"eval_specificity": 0.959,
"eval_auc": 0.983,
"eval_runtime": 30.15,
"eval_samples_per_second": 226.4,
"eval_steps_per_second": 3.55,
"epoch": 100.0,
"checkpoint_path": "/path/to/checkpoint.ckpt",
"train_data_path": "/path/to/train.csv",
"dev_data_path": "/path/to/dev.csv",
"test_data_path": "/path/to/test.csv"
}
Loading predictions for custom analysis:
import numpy as np
data = np.load('test_predictions.npz')
logits = data['logits'] # Raw model outputs
probs = data['probabilities'] # Softmax probabilities
preds = data['predictions'] # Predicted class labels
labels = data['labels'] # True labels
Two SLURM scripts are provided in slurm_scripts/ for running on HPC clusters (configured for NIH Biowulf):
File Locations:
slurm_scripts/
├── wrapper_run_csv_binary.sh # <-- EDIT THIS FILE with your paths
└── run_csv_binary.sh # Main job script (no edits needed)
wrapper_run_csv_binary.sh - Configuration wrapper (EDIT THIS)
sbatch to submit the jobrun_csv_binary.sh - Main SLURM job script (no edits needed)
caduceus_env)What to Edit in wrapper_run_csv_binary.sh:
| Variable | Description | Example |
|---|---|---|
DATA_DIR | Path to folder with train.csv, dev.csv, test.csv | /data/user/my_dataset |
DATASET_NAME | Name for output folder (no spaces) | phage_detection |
CONFIG_PATH | Path to model_config.json from pretraining | /data/user/pretrain/model_config.json |
PRETRAINED_PATH | Path to checkpoint from pretraining | /data/user/pretrain/checkpoints/last.ckpt |
MODEL | Model type | caduceus, mamba, or hyena |
MODEL_NAME | Model registry name | dna_embedding_caduceus |
LR | Learning rate | 6e-4 |
BATCH_SIZE | Batch size (reduce if OOM) | 32 |
MAX_LENGTH | Max sequence length | 1024 |
MAX_EPOCHS | Training epochs | 100 |
D_OUTPUT | Number of classes | 2 |
Modifying SLURM Resources (if needed):
If you need to change GPU type, memory, or time limit, edit run_csv_binary.sh:
#SBATCH --partition=gpu # Partition name
#SBATCH --gres=gpu:a100:1 # GPU type and count
#SBATCH --mem=64g # Memory
#SBATCH --cpus-per-task=8 # CPU cores
#SBATCH --time=24:00:00 # Time limit
Changing the Conda Environment:
If your conda environment has a different name, edit this line in run_csv_binary.sh:
source activate caduceus_env # Change 'caduceus_env' to your env name
Extract embeddings from a model and analyze their quality using linear probes, silhouette scores, PCA visualization, and a 3-layer neural network classifier.
Using the SLURM wrapper (recommended):
slurm_scripts/wrapper_run_embedding_analysis.sh:export CSV_DIR="/path/to/csv/data" # Directory with train.csv, dev.csv, test.csv
export CONFIG_PATH="/path/to/model_config.json"
export CHECKPOINT_PATH="/path/to/checkpoint.ckpt"
export INCLUDE_RANDOM_BASELINE="true" # Compare against random init baseline
cd slurm_scripts
bash wrapper_run_embedding_analysis.sh
Or run interactively (useful for debugging or in an interactive SLURM session):
bash run_embedding_analysis_interactive.sh
Direct command:
python -m src.embedding_analysis \
--csv_dir="/path/to/csv/data" \
--checkpoint_path="/path/to/checkpoint.ckpt" \
--config_path="/path/to/config.json" \
--output_dir="./outputs/embedding_analysis" \
--pooling="mean" \
--batch_size=32 \
--include_random_baseline
Outputs:
| File | Description |
|---|---|
embeddings_pretrained.npz | Extracted embeddings for train/val/test sets |
test_predictions_pretrained.csv | Test set predictions with probabilities |
pca_visualization_pretrained.png | PCA plot showing class separation |
three_layer_nn_pretrained.pt | Trained 3-layer NN classifier weights |
embedding_analysis_results.json | All metrics (see below) |
If --include_random_baseline is set, additional files with _random suffix are created for comparison.
Test predictions CSV columns:
sequence: Original DNA sequencelabel: Ground truth labellinear_probe_pred: Linear probe predicted labellinear_probe_prob: Linear probe probability (class 1)nn_pred: 3-layer NN predicted labelnn_prob: 3-layer NN probability (class 1)Metrics in embedding_analysis_results.json:
Run inference on a CSV file to get predictions with probabilities for threshold analysis.
Single file inference:
python -m src.inference \
--input_csv="/path/to/test.csv" \
--checkpoint_path="/path/to/checkpoint.ckpt" \
--config_path="/path/to/config.json" \
--output_csv="/path/to/predictions.csv" \
--threshold=0.5 \
--conjoin_test \
--save_metrics
Batch inference on multiple files:
# input_files.txt
/path/to/dataset1.csv
/path/to/dataset2.csv
/path/to/dataset3.csv
slurm_scripts/wrapper_run_batch_inference.sh:INPUT_LIST="/path/to/input_files.txt"
OUTPUT_DIR="/path/to/output_directory"
CONFIG_PATH="/path/to/model_config.json"
CHECKPOINT_PATH="/path/to/checkpoint.ckpt"
cd slurm_scripts
bash wrapper_run_batch_inference.sh
This submits a separate SLURM job for each input file. All predictions and logs are saved to the output directory.
Output CSV columns:
sequence: Original sequencelabel: Original label (if present)prob_0, prob_1: Class probabilitiespred_label: Predicted labelMetrics JSON (when --save_metrics is used):
Key options:
--threshold: Custom classification threshold for sensitivity/specificity tradeoff analysis (default: 0.5)--conjoin_test: Use post-hoc reverse complement averaging (recommended for Caduceus-Ph)--save_metrics: Calculate and save metrics to JSON if labels are present| File | Description |
|---|---|
src/dataloaders/datasets/csv_dataset.py | PyTorch Dataset for loading CSV files |
src/dataloaders/genomics.py | Added CSVDatasetLoader class |
src/callbacks/test_results.py | Callback for computing metrics with scikit-learn |
src/embedding_analysis.py | Extract embeddings, linear probe, PCA, 3-layer NN |
src/inference.py | Run inference with probability outputs |
configs/experiment/csv_binary.yaml | Experiment config for CSV classification |
configs/pipeline/csv_binary.yaml | Pipeline config |
configs/dataset/csv_dataset.yaml | Dataset config |
slurm_scripts/run_csv_binary.sh | SLURM job script for fine-tuning |
slurm_scripts/wrapper_run_csv_binary.sh | Wrapper for fine-tuning job submission |
slurm_scripts/run_csv_binary_interactive.sh | Interactive fine-tuning (no sbatch) |
slurm_scripts/run_embedding_analysis.sh | SLURM job script for embedding analysis |
slurm_scripts/wrapper_run_embedding_analysis.sh | Wrapper for embedding analysis submission |
slurm_scripts/run_embedding_analysis_interactive.sh | Interactive embedding analysis |
slurm_scripts/run_inference.sh | SLURM job script for inference |
slurm_scripts/wrapper_run_batch_inference.sh | Wrapper for batch inference configuration |
slurm_scripts/submit_batch_inference.sh | Batch submission for multiple inference jobs |
The remainder of this README is from the original Caduceus repository.
[Blog] | [arXiv] | [HuggingFace 🤗]
This repository contains code for reproducing the results in the paper "Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling," Schiff et al. (2024).
We have uploaded a pre-trained Caduceus model to the Huggingface hub. The available models are:
You can either use the pre-trained model directly within your trainer scripts or modify the config that initializes the model.
To use the pre-trained model for masked language modeling, use the following snippet:
from transformers import AutoModelForMaskedLM, AutoTokenizer
# See the `Caduceus` collection page on the hub for list of available models.
model_name = "kuleshov-group/caduceus-ph_seqlen-131k_d_model-256_n_layer-16"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForMaskedLM.from_pretrained(model_name)
Alternatively, you can instantiate a model from scratch to train on your own data as follows:
from transformers import AutoConfig, AutoModelForMaskedLM
# Add any config overrides here, see the `config.json` file on the hub for details.
config_overrides = {}
# See the `Caduceus` collection page on the hub for list of available models.
config = AutoConfig.from_pretrained(
"kuleshov-group/caduceus-ph_seqlen-131k_d_model-256_n_layer-16",
**config_overrides,
)
model = AutoModelForMaskedLM.from_config(config)
To get started, create a conda environment containing the required dependencies.
conda env create -f caduceus_env.yml
Activate the environment.
conda activate caduceus_env
Create the following directories to store saved models and slurm logs:
mkdir outputs
mkdir watch_folder
Below, we describe the steps required for reproducing the experiments in the paper.
Throughout, the main entry point for running experiments is the train.py script.
We also provide sample slurm scripts for launching pre-training and downstream fine-tuning experiments in the slurm_scripts/ directory.
(Data downloading instructions are copied from HyenaDNA repo)
First, download the Human Reference Genome data.
It's comprised of 2 files, 1 with all the sequences (the .fasta file), and with the intervals we use (.bed file).
The file structure should look like
data
|-- hg38/
|-- hg38.ml.fa
|-- human-sequences.bed
Download fasta (.fa format) file (of the entire human genome) into ./data/hg38.
~24 chromosomes in the whole genome (merged into 1 file), each chromosome is a continuous sequence, basically.
Then download the .bed file with sequence intervals (contains chromosome name, start, end, split, which then allow you to retrieve from the fasta file).
mkdir -p data/hg38/
curl https://storage.googleapis.com/basenji_barnyard2/hg38.ml.fa.gz > data/hg38/hg38.ml.fa.gz
gunzip data/hg38/hg38.ml.fa.gz # unzip the fasta file
curl https://storage.googleapis.com/basenji_barnyard2/sequences_human.bed > data/hg38/human-sequences.bed
Launch pretraining run using the command line
python -m train \
experiment=hg38/hg38 \
callbacks.model_checkpoint_every_n_steps.every_n_train_steps=500 \
dataset.max_length=1024 \
dataset.batch_size=1024 \
dataset.mlm=true \
dataset.mlm_probability=0.15 \
dataset.rc_aug=false \
model=caduceus \
model.config.d_model=128 \
model.config.n_layer=4 \
model.config.bidirectional=true \
model.config.bidirectional_strategy=add \
model.config.bidirectional_weight_tie=true \
model.config.rcps=true \
optimizer.lr="8e-3" \
train.global_batch_size=1024 \
trainer.max_steps=10000 \
+trainer.val_check_interval=10000 \
wandb=null
or alternatively, if using a cluster that has slurm installed, adapt the scripts below:
slurm_scripts
|-- run_pretrain_caduceus.sh
|-- run_pretrain_hyena.sh
|-- run_pretrain_mamba.sh
and run the training as a batch job:
cd slurm_scripts
sbatch run_pretrain_caduceus.sh
The GenomicBenchmarks presented in Grešová et al. (2023) is comprised of 8 classification tasks.
We can launch a downstream fine-tuning run on one of the tasks using the sample command below:
python -m train \
experiment=hg38/genomic_benchmark \
callbacks.model_checkpoint_every_n_steps.every_n_train_steps=5000 \
dataset.dataset_name="dummy_mouse_enhancers_ensembl" \
dataset.train_val_split_seed=1 \
dataset.batch_size=256 \
dataset.rc_aug=false \
+dataset.conjoin_train=false \
+dataset.conjoin_test=false \
loader.num_workers=2 \
model=caduceus \
model._name_=dna_embedding_caduceus \
+model.config_path="<path to model_config.json>" \
+model.conjoin_test=false \
+decoder.conjoin_train=true \
+decoder.conjoin_test=false \
optimizer.lr="1e-3" \
trainer.max_epochs=10 \
train.pretrained_model_path="<path to .ckpt file>" \
wandb=null
This sample run will fine-tune a pre-trained Caduceus-PS model on the dummy_mouse_enhancers_ensembl task.
Note some of the additional arguments present here, relative to the pre-training command from above:
model.config_path contains the path model config that was saved during pre-training.
This will be saved to the run directory of the pre-training experiment.train.pretrained_model_path contains the path to the pre-trained model checkpoint.dataset.conjoin_train determines whether the dataset will return a single sequence (dataset.conjoin_train=false) or the concatenation of a sequence and its reverse complement along dim=-1, during downstream fine-tuning training.dataset.conjoin_test is the same as above, but for inference (e.g., validation / test).decoder.conjoin_train determines whether the prediction head (a mean pooling and linear projection in the case of the Genomics Benchmark) is expecting an input tensor of shape (batch_size, seq_len, d_model) or (batch_size, seq_len, d_model, 2) during downstream fine-tuning training.
When set to true the decoder is run on input[..., 0] and input[..., 1] and the results are averaged to produce the final prediction.decoder.conjoin_test is the same as above, but for inference (e.g., validation / test).Note this benchmark only contains a training and test split for each task.
Therefore, to have a more principled evaluation, we randomly split the training data into training and validation sets (90/10) using the dataset.train_val_split_seed argument.
We perform early stopping on validation metric (accuracy) and repeat this for 5 random seeds.
As with pre-training, we can also launch the fine-tuning run as a batch job using the provided run_genomic_benchmark.sh script.
We also provide a helper shell script wrapper_run_genomics.sh that can be used to launch multiple fine-tuning runs in parallel.
Finally, the run_genomics_benchmark_cnn.sh script can be used to train the CNN baseline for this experiment from scratch on the downstream tasks.
The Nucleotide Transformer suite of tasks was proposed in Dalla-Torre et al. (2023). The data is available on HuggingFace: InstaDeepAI/nucleotide_transformer_downstream_tasks.
We can launch a downstream fine-tuning run on one of the tasks using the sample command below:
python -m train \
experiment=hg38/nucleotide_transformer \
callbacks.model_checkpoint_every_n_steps.every_n_train_steps=5000 \
dataset.dataset_name="${task}" \
dataset.train_val_split_seed=${seed} \
dataset.batch_size=${batch_size} \
dataset.rc_aug="${rc_aug}" \
+dataset.conjoin_test="${CONJOIN_TEST}" \
loader.num_workers=2 \
model._name_=dna_embedding_caduceus \
+model.config_path="<path to model_config.json>" \
+model.conjoin_test=false \
+decoder.conjoin_train=true \
+decoder.conjoin_test=false \
optimizer.lr="1e-3" \
trainer.max_epochs=10 \
train.pretrained_model_path="<path to .ckpt file>" \
trainer.max_epochs=20 \
wandb=null
We can also launch as batch jobs (see run_nucleotide_transformer.sh and wrapper_run_nucleotide_transformer.sh for details).
This task comes from the recently proposed Long Range Benchmark (LRB) in Kao et al., 2023. The data is available on HuggingFace: InstaDeepAI/genomics-long-range-benchmark. For this task we fit a model to the pre-trained and frozen embeddings of the DNA language models. Therefore, to perform the evaluation, we proceed in 2 steps:
vep_embeddings.py script to extract the embeddings from the pre-trained model.
See the example below:torchrun \
--standalone \
--nnodes=1 \
--nproc-per-node=8 \
vep_embeddings.py \
--num_workers=2 \
--seq_len=131072 \
--bp_per_token=1 \
--embed_dump_batch_size=1 \
--name="caduceus-ps_downstream-seqlen=131k" \
--model_name_or_path="kuleshov-group/caduceus-ps_seqlen-131k_d_model-256_n_layer-16" \
--rcps
The --rcps flag is used to indicate that the model is reverse-complement equivariant.
When using other models, set this flag to false with --no-rcps.
To speed this step up, this script utilizes torch distributed data parallelism.
Please refer to the slurm script provided in slurm_scripts/dump_vep_embeddings.sh
to launch this step as a batch job.
vep_svm.ipynb.If you find our work useful, please cite our paper using the following:
@article{schiff2024caduceus,
title={Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling},
author={Schiff, Yair and Kao, Chia-Hsiang and Gokaslan, Aaron and Dao, Tri and Gu, Albert and Kuleshov, Volodymyr},
journal={arXiv preprint arXiv:2403.03234},
year={2024}
}
This repository is adapted from the HyenaDNA repo and leverages much of the training, data loading, and logging infrastructure defined there. HyenaDNA was originally derived from the S4 and Safari repositories.
We would like to thank Evan Trop and the InstaDeep team for useful discussions about the Nucleotide Transformer leaderboard and the Long Range Benchmark task.
Finally, we would like to thank MosaicML for providing compute resources for some of the pre-training experiments.