ferritin_plms/ligandmpnn/
pmpnn_runner.rs

1//! ProteinMPNN Runner
2//!
3//! Loads and runs ProteinMPNN models without requiring callers to import candle directly.
4use super::configs::ProteinMPNNConfig;
5use super::model::ProteinMPNN;
6use super::proteinfeatures::ProteinFeatures;
7use crate::loader::{Format, LoadOptions, WeightSource, var_builder_from_path};
8use crate::registry::{self, ModelCard};
9use crate::types::PseudoProbability;
10use anyhow::{Result, anyhow};
11use candle_core::{Device, Tensor};
12use std::path::Path;
13
14/// ProteinMPNN `.pt` checkpoints nest their tensors under `model_state_dict`.
15const PMPNN_FORMAT: Format = Format::Pth {
16    root_key: Some("model_state_dict"),
17};
18
19pub enum ProteinMPNNModels {
20    /// proteinmpnn_v_48_020: k_neighbors=48, dropout=0.2
21    V48_020,
22    /// ligandmpnn_v_32_020_25: k_neighbors=32, 25 ligand context atoms.
23    ///
24    /// Same `ProteinMPNN` struct, a different featurizer and an extra context
25    /// path — see `ModelTypes::LigandMPNN` (ferritin-100.11).
26    LigandV32_020_25,
27}
28
29impl ProteinMPNNModels {
30    /// This variant's registry id.
31    pub const fn registry_id(&self) -> &'static str {
32        match self {
33            Self::V48_020 => "proteinmpnn-v48-020",
34            Self::LigandV32_020_25 => "ligandmpnn-v32-020-25",
35        }
36    }
37
38    /// The architecture config this checkpoint expects.
39    ///
40    /// `k_neighbors` and `atom_context_num` differ between the two and are
41    /// NOT recoverable from the tensor shapes — they live in the `.pt` as
42    /// plain Python ints, which candle's pickle reader does not surface — so
43    /// picking the wrong one loads cleanly and computes the wrong thing.
44    pub fn config(&self) -> ProteinMPNNConfig {
45        match self {
46            Self::V48_020 => ProteinMPNNConfig::proteinmpnn(),
47            Self::LigandV32_020_25 => ProteinMPNNConfig::ligandmpnn(),
48        }
49    }
50
51    /// This variant's [`ModelCard`].
52    pub fn card(&self) -> &'static ModelCard {
53        registry::lookup(self.registry_id())
54            .expect("every ProteinMPNNModels variant must have a registry entry")
55    }
56
57    /// Weight source and file for this variant.
58    ///
59    /// Delegates to [`REGISTRY`][crate::registry::REGISTRY] (ferritin-goh.1).
60    pub fn model_info(&self) -> (WeightSource, &'static str) {
61        let card = self.card();
62        (card.source, card.file)
63    }
64
65    /// Renamed to [`model_info`][Self::model_info] (ferritin-100.8).
66    #[deprecated(since = "0.4.0", note = "renamed to `model_info`")]
67    pub fn hf_info(&self) -> (WeightSource, &'static str) {
68        self.model_info()
69    }
70}
71
72pub struct ProteinMPNNRunner {
73    model: ProteinMPNN,
74}
75
76impl ProteinMPNNRunner {
77    /// Load a ProteinMPNN model from HuggingFace hub.
78    pub fn from_pretrained(modeltype: ProteinMPNNModels, device: Device) -> Result<Self> {
79        Self::from_pretrained_with(modeltype, &LoadOptions::new(device))
80    }
81
82    /// Renamed to [`from_pretrained`][Self::from_pretrained] (ferritin-100.8).
83    #[deprecated(since = "0.4.0", note = "renamed to `from_pretrained`")]
84    pub fn load_model(modeltype: ProteinMPNNModels, device: Device) -> Result<Self> {
85        Self::from_pretrained(modeltype, device)
86    }
87
88    /// Renamed to [`from_pretrained_with`][Self::from_pretrained_with] (ferritin-100.8).
89    #[deprecated(since = "0.4.0", note = "renamed to `from_pretrained_with`")]
90    pub fn load_model_with(modeltype: ProteinMPNNModels, opts: &LoadOptions) -> Result<Self> {
91        Self::from_pretrained_with(modeltype, opts)
92    }
93
94    /// Load with an explicit device and dtype (ferritin-100.9).
95    pub fn from_pretrained_with(modeltype: ProteinMPNNModels, opts: &LoadOptions) -> Result<Self> {
96        let (source, filename) = modeltype.model_info();
97        let weights_path = source.fetch(filename)?;
98        Self::from_path_as(&weights_path, &modeltype.config(), opts)
99    }
100
101    /// Load from a local .pt file (e.g. from ferritin-test-data or a cached download).
102    ///
103    /// Assumes a ProteinMPNN checkpoint. For LigandMPNN use
104    /// [`from_path_as`][Self::from_path_as] — the two need different configs
105    /// and a LigandMPNN checkpoint loaded as ProteinMPNN silently ignores its
106    /// nine ligand tensors.
107    pub fn from_path(path: impl AsRef<Path>, device: Device) -> Result<Self> {
108        Self::from_path_with(path, &LoadOptions::new(device))
109    }
110
111    /// Load from a local file with an explicit device and dtype.
112    pub fn from_path_with(path: impl AsRef<Path>, opts: &LoadOptions) -> Result<Self> {
113        Self::from_path_as(path, &ProteinMPNNConfig::proteinmpnn(), opts)
114    }
115
116    /// Load from a local file against an explicit architecture config.
117    pub fn from_path_as(
118        path: impl AsRef<Path>,
119        config: &ProteinMPNNConfig,
120        opts: &LoadOptions,
121    ) -> Result<Self> {
122        let path = path.as_ref();
123        let vb = var_builder_from_path(path, PMPNN_FORMAT, opts)?;
124        let model = ProteinMPNN::load(vb, config)
125            .map_err(|e| anyhow!("Failed to load MPNN weights from {}: {e}", path.display()))?;
126        Ok(Self { model })
127    }
128
129    /// Consume the runner and yield the loaded model.
130    ///
131    /// Lets callers that need the bare [`ProteinMPNN`] — such as
132    /// `MPNNExecConfig::load_model` — go through this tested loading path
133    /// instead of duplicating the download (ferritin-100.10).
134    pub fn into_model(self) -> ProteinMPNN {
135        self.model
136    }
137
138    /// Run ProteinMPNN and return a (L, 21) log-probability tensor for all positions.
139    ///
140    /// Useful for numerical parity tests against a Python reference.  Values are
141    /// log-softmax of the raw logits from a single structure-conditioned forward pass
142    /// (the same computation as `simple_decode`).
143    pub fn get_log_probs(&self, features: &ProteinFeatures) -> Result<Tensor> {
144        let output = self
145            .model
146            .simple_decode(features)
147            .map_err(|e| anyhow!("ProteinMPNN forward pass failed: {e}"))?;
148        // log_probs shape: (1, L, 21) — squeeze the batch dimension
149        output
150            .get_log_probs()
151            .squeeze(0)
152            .map_err(|e| anyhow!("Failed to squeeze batch dimension: {e}"))
153    }
154
155    /// Run ProteinMPNN and return per-residue pseudo-probabilities for the 21 amino acids.
156    pub fn get_pseudo_probabilities(
157        &self,
158        features: &ProteinFeatures,
159    ) -> Result<Vec<PseudoProbability>> {
160        let output = self
161            .model
162            .simple_decode(features)
163            .map_err(|e| anyhow!("ProteinMPNN forward pass failed: {e}"))?;
164        output
165            .get_pseudo_probabilities()
166            .map_err(|e| anyhow!("Failed to extract pseudo-probabilities: {e}"))
167    }
168}
ferritin_plms/ligandmpnn/pmpnn_runner.rs

ferritin_plms/ligandmpnn/
pmpnn_runner.rs
