Skip to main content

PlmRunner

ferritin_plms::plm_runner

Trait PlmRunner 

Source 
pub trait PlmRunner {
    // Required methods
    fn embed(&self, sequence: &str) -> Result<Tensor>;
    fn model_name(&self) -> &str;
    fn special_tokens(&self) -> SpecialTokenLayout;
    fn metadata(&self) -> ModelMetadata;
    fn device(&self) -> &Device;

    // Provided methods
    fn residue_count(&self, sequence: &str) -> usize { ... }
    fn embed_residues(&self, sequence: &str) -> Result<Tensor> { ... }
    fn logits(&self, _sequence: &str) -> Result<Tensor> { ... }
    fn embed_batch(&self, sequences: &[&str]) -> Result<Tensor> { ... }
    fn embed_residues_batch(&self, sequences: &[&str]) -> Result<Tensor> { ... }
}
Expand description

Trait implemented by all PLM runner types.

Structure and inverse-folding runners (ProteinMPNNRunner, StructureEncoderRunner) deliberately stay outside this trait — they consume or emit structure, not sequence embeddings.

Required Methods§

Source

fn embed(&self, sequence: &str) -> Result<Tensor>

Run a forward pass on sequence and return raw per-residue embeddings.

Shape: (1, L + special, d_model), where special is special_tokens().total(). Rows for special tokens are included; use embed_residues to get one row per residue.

Source

fn model_name(&self) -> &str

Model name / identifier string (e.g. “esm2”, “amplify”, “esmc”).

Source

fn special_tokens(&self) -> SpecialTokenLayout

The special tokens this runner’s tokenizer adds around a sequence.

Source

fn metadata(&self) -> ModelMetadata

Architecture dimensions, for sizing downstream layers.

Source

fn device(&self) -> &Device

Device the model’s weights live on.

Provided Methods§

Source

fn residue_count(&self, sequence: &str) -> usize

How many residues sequence encodes.

One per byte for the standard amino-acid alphabets, which is the default. SaProt is the exception: it tokenizes each residue as an (amino acid, 3Di structural state) pair over a 20x20 product alphabet, so "MdAaLp" is three residues, not six (ferritin-goh.3).

embed_residues uses this rather than sequence.len(), so a model whose alphabet is not one byte per residue still gets exactly one row per residue.

Source

fn embed_residues(&self, sequence: &str) -> Result<Tensor>

Per-residue embeddings with special-token rows removed.

Shape: (1, sequence.len(), d_model) — guaranteed, and therefore comparable row-for-row across models.

Provided in terms of embed and special_tokens, so an implementor declares its layout rather than reimplementing the strip — which callers used to hand-roll as .narrow(1, 1, l) against one model’s layout specifically.

Source

fn logits(&self, _sequence: &str) -> Result<Tensor>

Masked-LM logits over the token vocabulary.

Shape: (1, L + special, vocab_size) — the raw form, aligned with embed rather than embed_residues.

Defaults to an error: not every runner has a masked-LM head, and silently returning something else would be worse than refusing.

Source

fn embed_batch(&self, sequences: &[&str]) -> Result<Tensor>

Embed several sequences in one forward pass.

Shape: (B, T_max, d_model), where T_max is the longest residue_count + special_tokens().total() in sequences. Rows are right-padded and every padded row is exactly zero, so the result is the same whatever each architecture happens to compute at a pad position.

Row i’s real content occupies 0 .. residue_count(sequences[i]) + special_tokens().total(), which means the trailing special token sits at a different column in each row. Use embed_residues_batch when you want column j to be residue j for every row.

§Default implementation

Loops over embed and zero-pads. Correct for any runner, but it gives up the throughput this method exists for: B kernel launches instead of one. Override it with a real batched forward pass — which requires the model to honour a padding mask, because without one the padded keys leak into attention and change the real residues’ embeddings, a silent quality regression rather than an error.

Source

fn embed_residues_batch(&self, sequences: &[&str]) -> Result<Tensor>

Per-residue embeddings for a batch, special-token rows removed.

Shape: (B, R_max, d_model) where R_max is the largest residue_count. Row i, column j is residue j of sequences[i]; columns past that sequence’s length are exactly zero. This is the batched analogue of embed_residues and the form to pool over.

Provided in terms of embed_batch, and it re-checks the row count: an override that pads to the wrong length, or declares a SpecialTokenLayout its tokenizer does not follow, fails here rather than returning misaligned residues.

Dyn Compatibility§

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§