Start

Quickstart

Install the package, choose its storage paths, load a dataset, and run one linear probe.

Install from PyPI

The base package includes dataset loading and evaluation. Install the model extra when you also need model inference.

Datasets and evaluation

shell

conda create --name mrna_bench python=3.12
conda activate mrna_bench
python -m pip install mrna-bench

Model inference

shell

python -m pip install \
  --index-url https://download.pytorch.org/whl/cu126 \
  torch==2.7.1

python -m pip install wheel packaging ninja
python -m pip install --no-build-isolation 'mrna-bench[base_models]'

Set persistent storage

mRNABench downloads processed data and model weights outside the package. Set these paths once before loading a dataset or model. If the model path is omitted, it defaults to a model_weights directory below the data path.

python

import mrna_bench as mb

mb.update_data_path("/absolute/path/to/mrnabench-data")
mb.update_model_weights_path("/absolute/path/to/model-weights")

Load a dataset

python

import mrna_bench as mb

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

print(dataset.dataset_name)
print(dataset.data_df.shape)
print(dataset.metadata.task, dataset.metadata.default_split_type)

A successful load prints go-mf, the dataframe dimensions, and metadata containing a multilabel task with a homology split.

Run one benchmark in memory

python

import torch
import mrna_bench as mb

from mrna_bench.embedder import DatasetEmbedder
from mrna_bench.linear_probe import LinearProbeBuilder

device = torch.device("cuda")

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

embeddings = DatasetEmbedder(
    model,
    dataset,
    batch_size=8,
).embed_dataset()
embeddings = torch.stack(embeddings).cpu().numpy()

probe = (
    LinearProbeBuilder(dataset)
    .fetch_embedding_by_embedding_instance(
        model.short_name,
        embeddings,
    )
    .build()
)

metrics = probe.run_linear_probe(random_seed=2541)
print(metrics)

LinearProbeBuilder uses the first task and target in the dataset metadata, along with the dataset's default split. The benchmarking guide shows how to change those choices and save larger runs. The printed dictionary contains validation metrics with val_ prefixes.