mrna-bench 1.1.2

Benchmarking suite for mRNA property prediction.

mRNABench

This repository contains a workflow to benchmark the embedding quality of genomic foundation models on (m)RNA specific tasks. The mRNABench contains a catalogue of datasets and training split logic which can be used to evaluate the embedding quality of several catalogued models.

Jump to: Model Catalog Dataset Catalog

Setup

Several configurations of the mRNABench are available.

Datasets Only

If you are interested in the benchmark datasets only, you can run:

pip install mrna-bench

Full Version

The inference-capable version of mRNABench that can generate embeddings using Orthrus, DNA-BERT2, NucleotideTransformer, RNA-FM, and HyenaDNA can be installed as shown below. Note that this requires PyTorch version 2.2.2 with CUDA 12.1 and Triton uninstalled (due to a DNA-BERT2 issue).

conda create --name mrna_bench python=3.10
conda activate mrna_bench

pip install torch==2.2.2 --index-url https://download.pytorch.org/whl/cu121
pip install mrna-bench[base_models]
pip uninstall triton

Inference with other models will require the installation of the model's dependencies first, which are usually listed on the model's GitHub page (see below).

Post-install

After installation, please run the following in Python to set where data associated with the benchmarks will be stored.

import mrna_bench as mb

path_to_dir_to_store_data = "DESIRED_PATH"
mb.update_data_path(path_to_dir_to_store_data)

Usage

Datasets can be retrieved using:

import mrna_bench as mb

dataset = mb.load_dataset("go-mf")
data_df = dataset.data_df

The mRNABench can also be used to test out common genomic foundation models:

import torch

import mrna_bench as mb
from mrna_bench.embedder import DatasetEmbedder
from mrna_bench.linear_probe import LinearProbe

device = torch.device("cuda")

dataset = mb.load_dataset("go-mf")
model = mb.load_model("Orthrus", "orthrus-large-6-track", device)

embedder = DatasetEmbedder(model, dataset)
embeddings = embedder.embed_dataset()
embeddings = embeddings.detach().cpu().numpy()

prober = LinearProbe(
    dataset=dataset,
    embeddings=embeddings,
    task="multilabel",
    target_col="target",
    split_type="homology"
)

metrics = prober.run_linear_probe()
print(metrics)

Also see the scripts/ folder for example scripts that uses slurm to embed dataset chunks in parallel for reduce runtime, as well as an example of multi-seed linear probing.

Model Catalog

The current models catalogued are shown below. Currently, all models but AIDO.RNA are supported by the base_models optional installation.

Model Name	Model Versions	Description	Citation
Orthrus	orthrus-large-6-track orthrus-base-4-track	Mamba-based RNA FM pre-trained using contrastive learning on ~45M RNA transcripts to capture functional and evolutionary relationships.	[Code] [Paper]
AIDO.RNA	aido_rna_650m aido_rna_650m_cds aido_rna_1b600m aido_rna_1b600m_cds	Encoder Transformer-based RNA FM pre-trained using MLM on 42M ncRNA sequences. Version that is domain adapted to CDS is available.	[Paper]
RNA-FM	rna-fm mrna-fm	Transformer-based RNA FM pre-trained using MLM on 23M ncRNA sequences. mRNA-FM trained on mRNA CDS regions using codon tokenizer.	[Github]
DNABERT2	dnabert2	Transformer-based DNA FM pre-trained using MLM on multispecies genomic dataset. Uses BPE and other modern architectural improvements for efficiency.	[Github]
Nucleotide Transformer	2.5b-multi-species 2.5b-1000g 500m-human-ref 500m-1000g v2-50m-multi-species v2-100m-multi-species v2-250m-multi-species v2-500m-multi-species	Transformer-based DNA FM pre-trained using MLM on a variety of possible datasets at various model sizes. Sequence is tokenized using 6-mers.	[Github]
HyenaDNA	hyenadna-large-1m-seqlen-hf hyenadna-medium-450k-seqlen-hf hyenadna-medium-160k-seqlen-hf hyenadna-small-32k-seqlen-hf hyenadna-tiny-16k-seqlen-d128-hf	Hyena-based DNA FM pre-trained using NTP on the human reference genome. Available at various model sizes and pretraining sequence contexts.	[Github]
SpliceBERT	SpliceBERT.1024nt SpliceBERT-human.510nt SpliceBERT.510nt	Transformer-based RNA foundation model trained on 2M vertebrate mRNA sequences using MLM. Alternative versions trained on only human RNA, and with smaller context windows.	[Github]
RiNALMo	rinalmo	Transformer-based RNA foundation model trained on 36M ncRNA sequences using MLM and other modern architectural improvements such as RoPE, SwiGLU activations, and Flash Attention.	[Github]
UTR-LM	utrlm-te_el utrlm-mrl	Transformer-based RNA foundation model that is pre-trained on random and endogenous 5'UTR sequences from various species using MLM.	[Github]
3UTRBERT	utrbert-3mer utrbert-4mer utrbert-5mer utrbert-6mer	Transformer-based RNA foundation model that is pre-trained on the 3'UTR regions of 100K RNA sequences using MLM.	[Github]
RNA-MSM	rnamsm	Transformer-based RNA foundation model trained by using MSA from custom structure-based homology map on roughly 8M RNA sequences.	[Github]
RNAErnie	rnaernie	Transformer-based RNA foundation model trained using MLM at various mask sizes on 23M ncRNA sequences.	[Github]
RNABERT	rnabert	Transformer-based RNA foundation model trained using MLM and a structural alignment objective on 80K ncRNA sequences	[Github]
ERNIE-RNA	ernierna ernierna-ss	Transformer-based RNA foundation model trained using MLM with structural information added as attention mask biases. Pretrained on 20M ncRNA sequences.	[Github]

Many of the models wrappers (3UTRBERT, RiNALMo, UTR-LM, RNA-MSM, RNAErnie) use reimplementations from the multimolecule package. See their website for more details.

Adding a new model

All models should inherit from the template EmbeddingModel. Each model file should lazily load dependencies within its __init__ methods so each model can be used individually without install all other models. Models must implement get_model_short_name(model_version) which fetches the internal name for the model. This must be unique for every model version and must not contain underscores. Models should implement either embed_sequence or embed_sequence_sixtrack (see code for method signature). New models should be added to MODEL_CATALOG.

Dataset Catalog

The current datasets catalogued are:

Dataset Name	Catalogue Identifier	Description	Tasks	Citation
GO Molecular Function	go-mf	Classification of the molecular function of a transcript's product as defined by the GO Resource.	multilabel	website
Mean Ribosome Load (Sugimoto)	mrl‑sugimoto	Mean Ribosome Load per transcript isoform as measured in Sugimoto et al. 2022.	regression	paper
RNA Half-life (Human)	rnahl‑human	RNA half-life of human transcripts collected by Agarwal et al. 2022.	regression	paper
RNA Half-life (Mouse)	rnahl‑mouse	RNA half-life of mouse transcripts collected by Agarwal et al. 2022.	regression	paper
Protein Subcellular Localization	prot‑loc	Subcellular localization of transcript protein product defined in Protein Atlas.	multilabel	website
Protein Coding Gene Essentiality	pcg‑ess	Essentiality of PCGs as measured by CRISPR knockdown. Log-fold expression and binary essentiality available on several cell lines.	regression classification	paper

Adding a new dataset

New datasets should inherit from BenchmarkDataset. Dataset names cannot contain underscores. Each new dataset should download raw data and process it into a dataframe by overriding process_raw_data. This dataframe should store transcript as rows, using string encoding in the sequence column. If homology splitting is required, a column gene containing gene names is required. Six track embedding also requires columns cds and splice. The target column can have any name, as it is specified at time of probing. New datasets should be added to DATASET_CATALOG.