Module plm_runner 

Unified trait for protein language model runners.

PlmRunner provides a common interface for sequence embedding across ESM2, AMPLIFY, ESMC and ESM3, so downstream code can be generic over the runner type (benchmarking, ensembles, parity harnesses).

The special-token contract

This is the part that matters for correctness. Every runner’s tokenizer wraps the amino-acid sequence in special tokens, so PlmRunner::embed returns more rows than the sequence has residues. The old doc comment said (1, L, d_model) where L includes “any BOS/EOS tokens” — and “any” was doing far too much work: a caller comparing two runners’ embeddings had no way to know whether their rows were aligned, so a mismatch showed up as silently misaligned residues rather than an error.

Each runner now declares its layout via PlmRunner::special_tokens, and the provided PlmRunner::embed_residues strips it:

method	rows	use it when
embed	L + leading + trailing	you want the raw model output, specials included
embed_residues	exactly sequence.len()	you are indexing by residue, or comparing across models

All four current runners are BOS_EOS (one leading, one trailing), but that is a fact about these checkpoints, not a guarantee — declare the layout in the impl rather than assuming it at the call site.

Batching

embed_batch runs several sequences through one forward pass, and embed_residues_batch is its residue-aligned form. The default embed_batch just loops and zero-pads, so it is correct for every runner; a runner overrides it only once its architecture honours a padding mask, since without one the pad tokens change the real residues’ embeddings with no error raised.

The four ported architectures want that mask in three different shapes, so pad_token_batch builds one TokenBatch and each runner converts: ESM-2 passes the 1 = real mask straight through, AMPLIFY converts it with additive_padding_mask, and ESM-C and ESM-3 use it as a where_cond predicate.

// Row i of each corresponds to residue i, whatever the tokenizers do.
let x = a.embed_residues(seq)?;
let y = b.embed_residues(seq)?;
assert_eq!(x.dim(1)?, y.dim(1)?);

Structs

ModelMetadata

Architecture facts a caller needs to size downstream layers.

SpecialTokenLayout

How many special tokens a runner’s tokenizer wraps a sequence in.

TokenBatch

A right-padded batch of token ids and the padding mask that goes with it.

Traits

PlmRunner

Trait implemented by all PLM runner types.

Functions

additive_padding_mask

Turn a (batch, seq_len) padding mask (1 = real token, 0 = pad) into an additive (batch, seq_len) bias in dtype.

attention_mask_fill_value

The additive value used for padded key positions, chosen per dtype.

pad_token_batch

Right-pad token rows to a common length and build the matching mask.

stack_padded

Right-pad a set of (1, rows_i, d) tensors with zeros to a common row count and concatenate them into (B, rows_max, d).

zero_padded_rows

Zero every row of embeddings that sits at a padded position.