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::types::PseudoProbability;
8use anyhow::{Result, anyhow};
9use candle_core::pickle::PthTensors;
10use candle_core::{DType, Device, Tensor};
11use candle_nn::VarBuilder;
12use hf_hub::HFClientSync;
13use std::path::Path;
14
15const PMPNN_DTYPE: DType = DType::F32;
16
17pub enum ProteinMPNNModels {
18    /// proteinmpnn_v_48_020: k_neighbors=48, dropout=0.2
19    V48_020,
20}
21
22impl ProteinMPNNModels {
23    fn hf_info(&self) -> (&'static str, &'static str, &'static str, &'static str) {
24        // (owner, repo, revision, filename)
25        match self {
26            Self::V48_020 => (
27                "zcpbx",
28                "ligandmpnn-weights",
29                "main",
30                "model_params/proteinmpnn_v_48_020.pt",
31            ),
32        }
33    }
34}
35
36pub struct ProteinMPNNRunner {
37    model: ProteinMPNN,
38}
39
40impl ProteinMPNNRunner {
41    /// Load a ProteinMPNN model from HuggingFace hub.
42    pub fn load_model(modeltype: ProteinMPNNModels, device: Device) -> Result<Self> {
43        let (owner, repo, revision, filename) = modeltype.hf_info();
44        let client = HFClientSync::new()?;
45        let hf_repo = client.model(owner, repo);
46        let weights_path = hf_repo
47            .download_file()
48            .filename(filename)
49            .revision(revision)
50            .send()
51            .map_err(|e| anyhow!("Failed to download ProteinMPNN weights from HF hub: {e}"))?;
52        Self::from_path(&weights_path, device)
53    }
54
55    /// Load from a local .pt file (e.g. from ferritin-test-data or a cached download).
56    pub fn from_path(path: impl AsRef<Path>, device: Device) -> Result<Self> {
57        let path = path.as_ref();
58        let pth = PthTensors::new(path, Some("model_state_dict"))
59            .map_err(|e| anyhow!("Failed to open {}: {e}", path.display()))?;
60        let vb = VarBuilder::from_backend(Box::new(pth), PMPNN_DTYPE, device);
61        let config = ProteinMPNNConfig::proteinmpnn();
62        let model = ProteinMPNN::load(vb, &config)
63            .map_err(|e| anyhow!("Failed to load ProteinMPNN weights: {e}"))?;
64        Ok(Self { model })
65    }
66
67    /// Run ProteinMPNN and return a (L, 21) log-probability tensor for all positions.
68    ///
69    /// Useful for numerical parity tests against a Python reference.  Values are
70    /// log-softmax of the raw logits from a single structure-conditioned forward pass
71    /// (the same computation as `simple_decode`).
72    pub fn get_log_probs(&self, features: &ProteinFeatures) -> Result<Tensor> {
73        let output = self
74            .model
75            .simple_decode(features)
76            .map_err(|e| anyhow!("ProteinMPNN forward pass failed: {e}"))?;
77        // log_probs shape: (1, L, 21) — squeeze the batch dimension
78        output
79            .get_log_probs()
80            .squeeze(0)
81            .map_err(|e| anyhow!("Failed to squeeze batch dimension: {e}"))
82    }
83
84    /// Run ProteinMPNN and return per-residue pseudo-probabilities for the 21 amino acids.
85    pub fn get_pseudo_probabilities(
86        &self,
87        features: &ProteinFeatures,
88    ) -> Result<Vec<PseudoProbability>> {
89        let output = self
90            .model
91            .simple_decode(features)
92            .map_err(|e| anyhow!("ProteinMPNN forward pass failed: {e}"))?;
93        output
94            .get_pseudo_probabilities()
95            .map_err(|e| anyhow!("Failed to extract pseudo-probabilities: {e}"))
96    }
97}
ferritin_plms/ligandmpnn/pmpnn_runner.rs

ferritin_plms/ligandmpnn/
pmpnn_runner.rs
