ferritin_plms/esmc/

pretrained.rs

//! ESMC pretrained model loading.
//!
//! Downloads weights from `biohub/ESMC-{300M,600M,6B}` on HuggingFace and
//! wraps the ESMC model for sequence embedding.
//!
//! ## Weight key layout (`biohub/ESMC-*` safetensors)
//!
//! The weights are saved from `ESMCForMaskedLM`, a HuggingFace wrapper that
//! nests the backbone under an `esmc` attribute. All backbone keys are
//! therefore prefixed with `esmc.`:
//!
//! | Python (HF safetensors key)                              | Rust VarBuilder path                          |
//! |----------------------------------------------------------|-----------------------------------------------|
//! | `esmc.embed.weight`                                      | `embed.weight`                                |
//! | `esmc.transformer.blocks.{i}.attn.layernorm_qkv.0.*`   | `transformer.blocks.{i}.attn.layernorm_qkv.0.*` |
//! | `esmc.transformer.blocks.{i}.attn.layernorm_qkv.1.*`   | `transformer.blocks.{i}.attn.layernorm_qkv.1.*` |
//! | `esmc.transformer.blocks.{i}.attn.out_proj.weight`      | `transformer.blocks.{i}.attn.out_proj.weight`   |
//! | `esmc.transformer.blocks.{i}.attn.q_ln.weight`          | `transformer.blocks.{i}.attn.q_ln.weight`       |
//! | `esmc.transformer.blocks.{i}.attn.k_ln.weight`          | `transformer.blocks.{i}.attn.k_ln.weight`       |
//! | `esmc.transformer.blocks.{i}.attn.rotary.*`             | `transformer.blocks.{i}.attn.rotary.*`          |
//! | `esmc.transformer.blocks.{i}.ffn.0.*`                   | `transformer.blocks.{i}.ffn.0.*`                |
//! | `esmc.transformer.blocks.{i}.ffn.1.weight`              | `transformer.blocks.{i}.ffn.1.weight`           |
//! | `esmc.transformer.blocks.{i}.ffn.3.weight`              | `transformer.blocks.{i}.ffn.3.weight`           |
//! | `esmc.transformer.norm.weight`                           | `transformer.norm.weight`                       |
//! | `esmc.sequence_head.0.weight/bias`                       | `sequence_head.0.weight/bias`                   |
//! | `esmc.sequence_head.2.weight`                            | `sequence_head.2.weight`                        |
//! | `esmc.sequence_head.3.weight/bias`                       | `sequence_head.3.weight/bias`                   |
//!
//! The `from_pretrained` loader auto-detects whether the `esmc.` prefix is
//! present and sets the VarBuilder root accordingly, so the same Rust model
//! code works against both the wrapped and unwrapped formats.

use crate::esmc::models::esmc::{ESMC, ESMCConfig, LogitsConfig};
use crate::loader::{LoadOptions, WeightSource, optional_prefix};
use crate::plm_runner::{
    ModelMetadata, PlmRunner, SpecialTokenLayout, pad_token_batch, zero_padded_rows,
};
use crate::registry::{self, ModelCard};
use anyhow::{Result, bail};
use candle_core::{Device, Tensor};

/// Available ESMC model variants hosted on HuggingFace.
pub enum ESMCModels {
    /// ESMC 300M — 30 layers, d_model=960 (~1.3 GB weights)
    ESMC300M,
    /// ESMC 600M — 36 layers, d_model=1152
    ESMC600M,
    /// ESMC 6B — 80 layers, d_model=2560 (backbone for ESMFold2)
    ESMC6B,
}

impl ESMCModels {
    /// This variant's registry id.
    pub const fn registry_id(&self) -> &'static str {
        match self {
            Self::ESMC300M => "esmc-300m",
            Self::ESMC600M => "esmc-600m",
            Self::ESMC6B => "esmc-6b",
        }
    }

    /// This variant's [`ModelCard`].
    pub fn card(&self) -> &'static ModelCard {
        registry::lookup(self.registry_id())
            .expect("every ESMCModels variant must have a registry entry")
    }

    /// Returns `(weight_source, filename, config)` for this variant.
    ///
    /// Repo and filename come from [`REGISTRY`][crate::registry::REGISTRY]
    /// (ferritin-goh.1). Those point at the **EvolutionaryScale** originals,
    /// not the `biohub/ESMC-*` re-exports this once used: those are
    /// TransformerEngine-FUSED, so nothing this port asks for resolved
    /// (ferritin-100.23).
    ///
    /// `ESMC6B` loads through its shard index and keeps its head at
    /// `lm_head` rather than `sequence_head` (ferritin-100.24).
    pub fn model_info(&self) -> Result<(WeightSource, &'static str, ESMCConfig)> {
        let card = self.card();
        if let Some(reason) = card.unsupported {
            bail!("{} is not supported: {reason}", card.id);
        }
        let config = match self {
            Self::ESMC300M => ESMCConfig::esmc_300m(),
            Self::ESMC600M => ESMCConfig::esmc_600m(),
            Self::ESMC6B => ESMCConfig::esmc_6b(),
        };
        Ok((card.source, card.file, config))
    }
}

/// Wraps a loaded ESMC model for sequence embedding inference.
pub struct ESMCRunner {
    model: ESMC,
    /// Retained so `PlmRunner::metadata` can report dimensions.
    config: ESMCConfig,
}

impl ESMCRunner {
    /// Download weights from HuggingFace and load the model.
    ///
    /// Handles the `ESMCForMaskedLM` wrapper prefix (`esmc.`) transparently:
    /// probes for `esmc.embed.weight` in the safetensors and sets the
    /// VarBuilder root to `vb.pp("esmc")` when found, otherwise uses the
    /// flat (unwrapped) layout.
    pub fn from_pretrained(model: ESMCModels, device: Device) -> Result<Self> {
        Self::from_pretrained_with(model, &LoadOptions::new(device))
    }

    /// Load with an explicit device and dtype.
    ///
    /// `ESMC6B` is ~24 GB at F32 but ~12 GB at BF16 — its on-disk precision —
    /// so reduced precision is the difference between loadable and not on a
    /// 32 GB machine. See `tests/test_plm_dtype_parity.rs` for the numerical
    /// cost (ferritin-100.9).
    pub fn from_pretrained_with(model: ESMCModels, opts: &LoadOptions) -> Result<Self> {
        let (source, filename, config) = model.model_info()?;
        let vb = source.var_builder(filename, opts)?;
        // Weights saved from ESMCForMaskedLM nest the backbone under "esmc".
        // ESMC-6B nests its backbone under "esmc." but keeps its head at the
        // top level as "lm_head"; 300M/600M are flat with "sequence_head".
        // So the head is resolved from the ORIGINAL root, not the descended
        // one (ferritin-100.24).
        let head = if vb.contains_tensor("lm_head.0.weight") {
            vb.pp("lm_head")
        } else {
            vb.pp("sequence_head")
        };
        let vb_root = optional_prefix(vb, "esmc", "embed.weight");
        let esmc = ESMC::load_with_head(vb_root, head, config.clone())?;
        Ok(Self {
            model: esmc,
            config,
        })
    }

    /// Tokenize `sequence` and run a forward pass.
    ///
    /// Returns per-residue embeddings with shape `(1, L, d_model)` where
    /// `L` includes the BOS and EOS tokens added by the tokenizer.
    pub fn embed_sequence(&self, sequence: &str) -> Result<Tensor> {
        let tokens = self.model.encode(sequence)?;
        let tokens = tokens.unsqueeze(0)?;
        let output = self.model.forward(&tokens, None, false)?;
        output
            .embeddings
            .ok_or_else(|| anyhow::anyhow!("ESMC forward() returned no embeddings"))
    }
}

impl PlmRunner for ESMCRunner {
    /// Delegate to `embed_sequence`, returning per-residue embeddings `(1, L, d_model)`.
    fn embed(&self, sequence: &str) -> Result<Tensor> {
        self.embed_sequence(sequence)
    }

    fn model_name(&self) -> &str {
        "esmc"
    }

    /// `ESMC::encode` calls `tokenize_sequence(.., true)`, which wraps the
    /// sequence in BOS and EOS.
    fn special_tokens(&self) -> SpecialTokenLayout {
        SpecialTokenLayout::BOS_EOS
    }

    fn metadata(&self) -> ModelMetadata {
        ModelMetadata {
            d_model: self.config.d_model,
            n_layers: self.config.n_layers,
            vocab_size: self.config.embedding_dim,
            // Rotary positions: no hard architectural cap.
            max_positions: None,
        }
    }

    fn device(&self) -> &Device {
        self.model.device()
    }

    /// Sequence logits `(1, L + 2, vocab_size)`.
    fn logits(&self, sequence: &str) -> Result<Tensor> {
        let tokens = self.model.encode(sequence)?.unsqueeze(0)?;
        let out = self.model.logits(
            &tokens,
            LogitsConfig {
                sequence: true,
                return_embeddings: false,
                return_hidden_states: false,
            },
        )?;
        out.sequence_logits
            .ok_or_else(|| anyhow::anyhow!("ESMC logits() returned no sequence logits"))
    }

    /// One batched forward pass over right-padded sequences (ferritin-100.12).
    ///
    /// ESM-C's `sequence_id` is a `where_cond` predicate — `1` keeps a key,
    /// `0` replaces its score with `-inf` — so the `U8` mask goes in as-is.
    /// It is passed explicitly rather than relying on `ESMC::forward`'s
    /// fallback (`tokens.ne(pad_id)`), which would happen to agree here but
    /// ties correctness to the padding id rather than to the batch itself.
    fn embed_batch(&self, sequences: &[&str]) -> Result<Tensor> {
        let rows: Vec<Vec<u32>> = sequences.iter().map(|s| self.model.tokenize(s)).collect();
        let batch = pad_token_batch(&rows, self.model.pad_token_id(), self.model.device())?;
        let output = self.model.forward(&batch.ids, Some(&batch.mask), false)?;
        let embeddings = output
            .embeddings
            .ok_or_else(|| anyhow::anyhow!("ESMC forward() returned no embeddings"))?;
        zero_padded_rows(&embeddings, &batch.mask)
    }
}