Extend

Add a dataset

Define metadata, produce the shared dataframe schema, choose a split, and register a reusable dataset.

The dataset contract

A registered dataset subclasses BenchmarkDataset, declares aDatasetMetadata record, and implements_get_data_from_raw(). Hosting processed data as Parquet lets users load the dataset without installing raw-data dependencies.

Dataset adapter

import pandas as pd

from mrna_bench.datasets import (
    BenchmarkDataset,
    DatasetMetadata,
)


class MyDataset(BenchmarkDataset):
    METADATA = DatasetMetadata(
        dataset_name="my-dataset",
        species="human",
        task=["regression"],
        target_col=["target"],
        default_split_type="homology",
        benchmark_set="extended",
        evaluations=("linear_probe",),
    )

    def __init__(
        self,
        force_redownload_hf: bool = False,
        force_rebuild_raw: bool = False,
    ):
        super().__init__(
            force_redownload_hf=force_redownload_hf,
            force_rebuild_raw=force_rebuild_raw,
            hf_url="https://.../my-dataset.parquet",
        )

    def _get_data_from_raw(self) -> pd.DataFrame:
        return pd.DataFrame(
            {
                "sequence": ...,
                "gene": ...,
                "chromosome": ...,
                "target": ...,
            }
        )

Required and conditional columns

Column	When it is needed
`sequence`	Always. One nucleotide sequence per row.
Target column	One column for each name in `target_col`.
`gene`	Homology splitting.
`chromosome`	Chromosome splitting.
`cds`	Models that use codon-position tracks.
`splice`	Models that use splice-site tracks.

Dataset metadata requirements

Dataset identifiers cannot contain underscores because persisted embedding filenames use underscores as separators.
Use core only for datasets intended for the central comparison.
List every supported route in evaluations.
Implement get_vep_pairs() for paired variant datasets.
Match the default split to the information available in the dataframe.

Register and check the output

Add the class to DATASET_CATALOG inmrna_bench/datasets/dataset_catalog.py. Before using it in a comparison, check metadata validation, processed schema, download and cache behavior, and the selected split with focused tests for the new adapter.

Use an ad hoc dataframe

Registration is unnecessary when you only need embeddings for a temporary table or your own downstream analysis. The dataframe needs asequence column; identifier and metadata columns remain in your dataframe but are not interpreted by the embedder.

python

import numpy as np
import pandas as pd
import torch
import mrna_bench as mb

from mrna_bench.embedder import DatasetEmbedder

dataframe = pd.DataFrame(
    {
        "sample_id": ["transcript-a", "transcript-b"],
        "sequence": [
            "ACGTTGCAACGTTGCA",
            "TTGCAACGTTGCAACG",
        ],
    }
)

model = mb.load_model(
    "Orthrus",
    "orthrus-large-4-track",
    device=torch.device("cuda"),
)
embedder = DatasetEmbedder.from_dataframe(model, dataframe)
embedding_tensors = embedder.embed_dataset()
embeddings = torch.stack(embedding_tensors).cpu().numpy()

# DatasetEmbedder preserves dataframe row order.
assert embeddings.shape[0] == len(dataframe)
np.savez_compressed(
    "custom-sequence-embeddings.npz",
    sample_id=dataframe["sample_id"].to_numpy(),
    embedding=embeddings,
)

The returned tensors follow the input row order. Save identifiers next to the embedding matrix so downstream results can be joined back to the source table.

Optional CDS and splice tracks

Add cds and splice columns only when the selected model consumes those channels. Each cell contains a one-dimensional array with the same length as its sequence. A CDS track marks the first base of each codon with 1; the all-zero example below represents a noncoding sequence.

python

# Optional model tracks must match each sequence length.
dataframe["cds"] = [
    np.zeros(len(sequence), dtype=np.int8)
    for sequence in dataframe["sequence"]
]
dataframe["splice"] = [
    np.zeros(len(sequence), dtype=np.int8)
    for sequence in dataframe["sequence"]
]

embedder = DatasetEmbedder.from_dataframe(model, dataframe)