Prima-Bench — Guide

Good models need great benchmarks.

This guide has three parts:

1. Why Prima-Bench — what it is, who it’s for, and the philosophy/principles that shape every design decision.
2. Consumer perspective — how to evaluate a model against a task.
3. Developer perspective — how to add tasks, metrics, datasets, and curation recipes.

This is the human guide — the narrative and walkthroughs. If you’re a coding agent working in this repo, read AGENTS.md instead (the operational agent guide). For the internal directory map and layer-by-layer separation of concerns, see ARCHITECTURE.md; each src/prima_bench/<layer>/ directory carries its own README.md.

Part 1 — Why Prima-Bench

North star

A world-class biological modeling benchmark. Prima-Bench scores predictions from biological foundation models against held-out truth, across many biological questions (perturbation prediction, cell-type annotation, batch integration, variant-effect prediction, …), in a way that is reproducible, comparable, and cheap to iterate on.

What it is (and isn’t)

Prima-Bench is a library, not a service or a framework. It does exactly one thing: it takes model predictions in as a typed value (EvalData) and returns scores out as a structured record (EvalResult).

• It never imports a model framework. It does not run, train, or schedule models. The model lives entirely outside the library, owned by the model author.
• There is no Method primitive, no container, no orchestrator, no leaderboard service. The unit of contribution is predictions at evaluation time, not executable code packaged for a runtime.

Who it’s for

The primary consumer is a foundation-model team mid-training — a job already running on a GPU cluster with its own checkpointing, scheduling, and MLOps. The defining use case is in-loop validation: scoring a checkpoint every N steps is a single in-process function call, not a Docker round-trip.

The central design decision: library-driven, not framework-driven

There are two coherent patterns for a benchmark, and they invert who calls whom:

Prima-Bench chose library-driven deliberately. The main reason: you can add framework-driven on top of a library, but not the reverse. Python metrics/contracts/datasets can later be wrapped in a CLI + container and orchestrated if a public leaderboard ever becomes necessary. Going the other way is much harder.

When to revisit this (none are v0 requirements): held-out test labels consumers must never see (needs a service that holds truth behind a boundary); a public leaderboard with numbers comparable across external teams (needs a neutral runtime); or methods authored in R/Julia/etc. (the containerized-component pattern starts paying for itself).

Principles

• Predictions in, scores out. The typed boundary is EvalData → EvalResult. The contract is validated before any metric runs; errors fail loud, not silently.

• Reproducibility via hashing, not isolation. Same predictions + same dataset version + same package version + same metric selection ⇒ the same scores, always. EvalResult records SHA-256 hashes of the canonicalized predictions and truth, so two parties can verify they evaluated the same thing.

• Solution privacy is an invariant. The solution (ground-truth) split is never reachable from consumer code. The bucket has it, the runner loads it internally, and no fetch_* call ever returns its path.

• Censorship is by NaN, not by zero or by absence. A held-out answer keeps its row but its value becomes NaN. A model that touches a censored value propagates NaN loudly — failures are local to the line that read the answer.

• Immutable, versioned bytes. Dataset.splits[...].remote_uri is fixed per (name, version). Re-curation requires a new version; the versioned path is the integrity guarantee.

• Separation of concerns (each layer changes for a different reason, and depends only on the layer above):

  Recipe        writes immutable bytes in a publication-shaped layout   (external pkg)
     ▼
  Dataset       maps  name@version → {artifact: remote_uri}
     ▼
  SplitPreparer per (task, dataset): artifacts → {train, test, solution} paths + loaders
     ▼
  Runner        contract.validate → feed metric panel → EvalResult
     ▼
  Consumer      pb.fetch_data / pb.load / pb.Evaluation

  Recipes describe a publication, not an API — so third-party buckets (e.g. OP-bio) can be wrapped without re-curation. Datasets carry only identity + URIs. The (task, dataset) wiring lives in the SplitPreparer, so one dataset can serve many tasks. The runner never sees bytes or URIs — only EvalData — and owns no domain logic; metric functions do.

• Dual-mode metrics, one code path. Metrics are torchmetrics-style update/compute. The same code is correct single-process and under data-parallel distributed evaluation, with no dist.is_initialized() branching: state holds raw inputs (not computed scores), update appends a shard, compute concatenates and does all the math globally. A one-shot call and a sharded stream of the same data produce identical results.

• Fetch on demand, cache locally, stay bucket-agnostic. Datasets are downloaded when first needed and cached. All bucket I/O goes through a single S3-compatible boto3 client (AWS S3, GCS via its S3-interop API, Cloudflare R2, MinIO) — there is no gcloud dependency.

• Documentation is autogenerated from registry metadata as Markdown cards (prima-bench docs build), so the generated cards stay in sync with code (every task/metric/dataset/benchmark renders a card).

Core concepts

Part 2 — Consumer perspective

You run your model in your own process and hand predictions to the library. The library validates them against the task contract, runs the metric panel, and returns scores.

Setup

Datasets are fetched on demand from object storage. Credentials are a GCS HMAC keypair carried in prima-bench’s own environment variables; a gitignored project-local .env is loaded automatically. The canonical datasets live on GCS, so the HMAC keypair is all you need (copy .env.example → .env):

HMAC_BOTO_PRIMA_MENTE_BENCHMARK_KEY_ID=<gcs-hmac-access-id>
HMAC_BOTO_PRIMA_MENTE_BENCHMARK_SECRET=<gcs-hmac-secret>
HMAC_BOTO_PRIMA_MENTE_BENCHMARK_REGION=auto

(These are HMAC keys, not literal AWS keys. prima-bench reads them directly and passes them to boto3 as explicit credentials — no AWS_* / ~/.aws fallback.)

See the README for generating an HMAC keypair with gcloud (one-time) and for pointing at non-GCS stores.

Discover what’s available

import prima_bench as pb

pb.list_tasks()
# → ["batch_integration", "cell_type_annotation", "perturbation_prediction", ...]

# Dataset keys are "{publication}/{slug}@{version}" and evolve — treat this call,
# not a hardcoded string, as the source of truth.
pb.list_datasets(task="perturbation_prediction")
# → ["hao_2024_nature_methods/norman_2019_science@0.0.1"]

pb.list_metrics(task="perturbation_prediction")
# → ["l2_pseudobulk", "l2_pseudobulk_top_1000_control_genes"]

The CLI mirrors this (a thin shell over the same Python API):

prima-bench list tasks
prima-bench list metrics --task perturbation_prediction
prima-bench docs build

Fetch data → predict → evaluate

pb.fetch_data returns (train_path, test_path) — no per-task polymorphism. Test-only datasets expose no train split and raise KeyError; use pb.fetch_test for those. The solution bytes are never returned.

import prima_bench as pb

# Downloads on demand into the local cache; returns local Paths.
train_path, test_path = pb.fetch_data(task="perturbation_prediction")

# Optional: invoke the registered per-(task, split) loader.
# Returns EvalData (e.g. AnnData/tensor-backed) — you may also load the path yourself.
train = pb.load(train_path, task="perturbation_prediction", split="train")
test  = pb.load(test_path,  task="perturbation_prediction", split="test")

model.fit(train)
predictions = model(test)            # an EvalData (or aligned predictions)

run = pb.Evaluation(task="perturbation_prediction")  # loads censored solution; builds the metric panel
run.update(predictions=predictions)                  # one-shot == a single update
result = run.compute()                               # per-metric compute() → EvalResult

Why paths, not deserialized objects? A train fold can be multiple GB; forcing materialization burns memory you may not have. Paths give a uniform (Path, Path) return regardless of artifact format (h5ad, csv, fasta, …), and let you choose your loader (backed-mode AnnData, streaming, a PyTorch dataloader, polars, …). The library doesn’t compete with scanpy/PyTorch/polars, and paths scale to very large pretraining sets.

Streaming / sharded evaluation

pb.Evaluation is dual-mode. For large or distributed runs, feed predictions shard-by-shard instead of materializing everything at once — the result is identical to a single one-shot update of the same data:

run = pb.Evaluation(task="cell_type_annotation")
for shard in prediction_shards:     # produced incrementally / per rank
    run.update(predictions=shard)
result = run.compute()

Evaluation-only / zero-shot datasets

Some datasets (e.g. test-only variant-effect sets) expose no train split. pb.fetch_data raises KeyError for those — use pb.fetch_test to get just the eval input:

test_path = pb.fetch_test(task="variant_effect_classification", dataset="avsec_2026_nature/clinvar_missense@0.0.1")
test = pb.load(test_path, task="variant_effect_classification", split="test", dataset="avsec_2026_nature/clinvar_missense@0.0.1")

Variant-effect: building model inputs from the genome

Variant-effect tasks need reference-genome sequence windows as inputs. pb.get_sequences(...) turns variant rows into ref/alt sequences for your model. It is a consumer-side convenience — the runner never touches it.

Coordinate convention (silent-bug territory — get it right at every boundary): a variant position (pos) is 1-based (VCF / human-genetics convention); an index, interval, or array slice is 0-based half-open. get_sequences takes 1-based pos and converts once internally; the returned variant_index (where the variant lands inside the window) is 0-based. It is variant-centric: it requires both ref and alt and raises unless the FASTA matches one of them. See ARCHITECTURE.md → Coordinate conventions for the full pysam-bridging table.

What you get back

EvalResult is structured and comparable across models: per-metric scores plus SHA-256 hashes of the canonicalized predictions, truth, and the dataset ref. What you do with the numbers (log to W&B, write a CSV, publish) is yours — the library has no opinion and no leaderboard.

Part 3 — Developer perspective

The code lives under src/prima_bench/; each subdirectory owns one layer of the pipeline and carries its own README.md. Read ARCHITECTURE.md before making or planning changes — avoid changes that would alter the decisions recorded there; if a request can’t be met within the current architecture, update ARCHITECTURE.md and the relevant README.md to reflect the change.

Layout at a glance

src/prima_bench/
├── core/         EvalData, Contract, EvalResult, MetricResult — pure value types, no I/O
├── registry/     in-memory registries for metrics / tasks / datasets / benchmarks
├── io/           hashing, atomic writes, cache paths, S3-compatible downloads, JSON serialize
├── metrics/      the dual-mode Metric base class + device helpers
├── transforms/   shared numerical helpers (e.g. kNN) metrics call directly
├── genome/       reference-FASTA access + pure sequence-window math (VEP, internal)
├── datasets/     dataset descriptors + per-dataset loaders (one subdir per dataset)
├── tasks/        task definitions: contract, metric panel, split preparers (one subdir per task)
├── benchmarks/   Benchmark grouping records (a manifest, not an executor)
├── runner/       the pb.Evaluation entry point — feeds a task's metric panel
├── cli/          `prima-bench` click commands (discovery + docs)
├── docs/         Markdown card autogen from registry metadata
└── sequences.py  pb.get_sequences — consumer helper that builds VEP inputs

Registration is by import side-effect. Built-in tasks/datasets call register_* at module import time; prima_bench/__init__.py imports those modules eagerly, so everything is registered the moment import prima_bench returns. When you add a metric/preparer module, re-export/import it from the relevant __init__.py so its decorator/registration fires.

Before you write code

From AGENTS.md:

• Search before you write. Grep core/, metrics/, io/, and the register_* decorators for something that already does it. Reuse existing bases rather than re-deriving them; don’t hand-roll a metric torchmetrics already provides. (Reuse ≥~80% → reuse; close → extend without breaking callers; only write new when nothing fits.)
• Be concise and modular. Minimum code that solves the problem; match the surrounding file’s density.
• Worktrees. Create a new worktree under .worktrees/ before making code changes.
• Linear. File issues in the prima-bench project on the Tech team.
• PRs ship as drafts on this repo (gh pr create --draft).

Add a new task

A task is four things — a name, a Contract, a list of metric names, and a SplitPreparer per dataset. Build them in order under tasks/<task>/:

tasks/<task>/
├── __init__.py           registers the Task; imports the submodules below
├── contract.py           Contract dataclass + any shared projection helper
├── metrics/
│   ├── __init__.py       imports each metric module so decorators fire
│   └── <metric>.py       one file per metric (or closely-related pair)
├── split_preparers/
│   ├── __init__.py       imports each preparer module so registrations fire
│   └── <dataset>.py      one file per dataset this task can score against
└── aggregation.py        (optional) panel-level overall score

1. Contract (contract.py): a @dataclass(frozen=True) implementing the Contract protocol. It validates (predictions, solution) before any metric runs — shapes, required metadata columns, feature alignment, pairing rules. If several metrics share a projection (e.g. joining on a key), put that helper here.
2. Metrics (metrics/<metric>.py): one registered function/class per file. description is required (doc autogen depends on it). Cross-task helpers go in transforms/; task-local primitives in metrics/utils.py.
3. Split preparers (split_preparers/<dataset>.py): a prepare(...) that fetches artifacts via dataset.fetch_artifact(...) and returns the consumer-facing dict (train/test/solution; set test = train for unsupervised tasks). How test/solution are produced is dataset-dependent — not always an eval-time derivation. prepare may censor a golden solution into a NaN-censored test (e.g. cell_type_annotation), or — when the recipe already ships an already-censored test plus a per-item solution (recipe-side censoring, e.g. sequence_classification) — just map them through and derive nothing. Provide a per-split loader (Callable[[Path], EvalData]) for every returned key — registration validates this. Call register_split_preparer(...) at module top level.
4. Register (__init__.py): register_task(Task(name=..., contract=..., metrics=(...), default_dataset="<name>@<version>", description="...")), then import metrics/ and split_preparers/ so their side-effect registrations fire.

Add a metric to an existing task: drop a file under tasks/<task>/metrics/, decorate with @register_metric(name=..., description=..., references=...), import it from metrics/__init__.py, and add the name to the task’s metrics tuple.

Add a dataset

Under datasets/<dataset>/:

• __init__.py calls register_dataset(Dataset(name=..., version=..., artifacts={...})) at import time — identity + named bucket URIs only.
• loaders/ holds the bytes → EvalData converters and any dataset-specific processing.

Per-task wiring (which artifact is train vs solution, how test is derived) does not live here — that’s the SplitPreparer’s job. The same dataset is wired to multiple tasks by registering more preparers, with no change to the descriptor.

Write a curation recipe

Recipes produce the immutable benchmark bytes. They are a separate package (recipes/), not part of installable prima_bench — they have their own dependency graph (GEARS, scanpy, boto3, …) and are run offline by curators. Consumers read the bucket via pb.fetch_data; recipes write it.

Layout — one directory per publication, one subdirectory per dataset; each (publication, dataset) is one recipe:

recipes/{first_author}_{year}_{journal}/{dataset}/
├── README.md     provenance: source/DOI/license, what it produces, preprocessing, version rationale, curator history
├── recipe.py     download + preprocess + upload all views/splits for this dataset
└── explore.py    scratch EDA / provenance for the choices recipe.py bakes in

Every new recipe PR must include both recipe.py and explore.py.

• Naming follows the company sequencing-data conventions: {first_author} (lowercase surname/consortium), 4-digit {year}, lowercase {journal}. The publication directory tracks the publication the data was downloaded from (often a derivative republication); {dataset} keeps the upstream source’s slug.
• Recipes are deliberately dumb. They write the censored test plus the per-item solution. They do not split train/test or know about “regimes” (supervised/zero-shot) or views beyond what they emit — that’s the SplitPreparer’s job at eval time, over the same bytes. One recipe can therefore serve several tasks.
• Bucket layout (canonical content lives under output/): gs://prima-mente-benchmarks/{pub}/output/{dataset}/v{version}/{view}/{split}/<artifact>, plus a manifest.json per version tying the bytes to the reviewed PR/commit and curator.
• Immutability & provenance. Versioned paths never change; any byte change requires a new v{version}/. Per-curation audit detail (who/when/which PR/tool versions) goes in the dataset README’s Curator history table — append one row per uploading execution. Run --upload against a clean post-merge main.

See recipes/README.md for the full convention and niki_2025_biorxiv/nt_benchmark/README.md for a worked example.

Testing

• Standard suite: uv run python -m pytest.

• Some integration tests are skip-by-default: they cross-check against an external reference library that isn’t a project dependency, so they importorskip and are silently skipped in CI. Run them by overlaying the dependency, e.g.:

  uv run --with kipoiseq python -m pytest tests/integration/test_indel_matches_kipoiseq.py

  Re-run the relevant differential test whenever you touch the code it validates (the covered source carries a NOTE pointing back to it).