Extend

Add a model

Implement loading, embedding, and any supported scoring methods for a new model.

The adapter contract

EmbeddingModel handles version validation, attention-backend validation, inference mode, and the batched embed() interface. Import model-specific packages inside the constructor so other adapters remain importable.

Model adapter

from collections.abc import Callable

import numpy as np
import torch

from mrna_bench.models import (
    EmbeddingModel,
    ModelBehavior,
    mean_pool,
)


class MyModel(EmbeddingModel):
    default_version = "my-model-base"
    valid_versions = ["my-model-base"]
    default_attn_implementation = "sdpa"
    valid_attn_implementations = ["eager", "sdpa"]
    hookable_layer_patterns = []
    supported_behaviors = frozenset({
        ModelBehavior.EMBEDDING,
    })

    def __init__(
        self,
        model_version: str,
        device: torch.device,
        attn_implementation: str | None,
    ):
        super().__init__(
            model_version,
            device,
            attn_implementation,
        )
        # Import optional model dependencies here.
        self.model = ...
        self.tokenizer = ...

    def embed(
        self,
        sequences: list[str],
        cds: list[np.ndarray] | None = None,
        splice: list[np.ndarray] | None = None,
        agg_fn: Callable = mean_pool,
    ) -> list[torch.Tensor]:
        ...
        return embeddings

Implement each supported output

Behavior	Required implementation
`embedding`	Batched `embed()` with pooled and unpooled output support.
`causal_likelihood`	An LM head plus compatible `logits()` and `sequence_score()`.
`pseudo_likelihood`	Expose the masked-language-model head with `_set_logits_model()`, or implement compatible `logits()` and `sequence_score()` methods.
`tracks`	`predict_tracks()` for a sequence-to-function model.

Masked-marginal scoring is separate: it evaluates only tokenizer positions affected by a substitution.

Register the class

Add the public model name and class to MODEL_CATALOG inmrna_bench/models/model_catalog.py. Versions are collected automatically from the class's valid_versions.

Add required third-party packages to base_models or a dedicated optional dependency group in pyproject.toml. Test model loading from a clean installation.

python

import torch
import mrna_bench as mb

model = mb.load_model(
    "MyModel",
    "my-model-base",
    device=torch.device("cuda"),
)

Check behavior before large runs

Each declared version loads and returns the expected output shape.
Batch output matches single-sequence output within numerical tolerance.
Unpooled output can remain ragged across sequence lengths.
Required cds and splice tracks fail clearly when absent.
Inference mode is deterministic and training mode preserves gradients.
Every declared likelihood or track behavior has a working implementation.