ferritin_plms/
lib.rs

1//! ferritin-plms
2//!
3//!
4//! ```shell
5//! cargo run --example amplify
6//! cargo run --example amplify --features metal
7//! ```
8//!
9//! # Model support matrix
10//!
11//! Generated from [`registry::REGISTRY`]; regenerate with
12//! `cargo test -p ferritin-plms --lib print_support_matrix -- --ignored --nocapture`.
13//! `test_lib_rs_support_matrix_is_current` fails if this copy drifts.
14//!
15//! The column to read first is **Parity**. "It compiles" and even "it loads"
16//! are not what you need before trusting a number — you need to know whether
17//! anyone has compared this port's output against the reference
18//! implementation.
19//!
20//! Read **not checked** as "this output could be anything". It is not a
21//! milder form of verified. ProteinMPNN sat in that column through every
22//! release up to v0.3.3 while agreeing with the reference on 2 of 93
23//! positions — worse than the ~4 you would expect from chance over a 21-token
24//! vocabulary. Its output was not approximately right, it was unrelated to
25//! what the model computes, and nothing in this table said so
26//! (ferritin-100.33).
27//!
28//! Weigh an unchecked row by how much it really shares with a verified one.
29//! A row whose family has no verified member anywhere is an entire
30//! architecture nobody has ever compared — the position ProteinMPNN was in.
31//! Every family in this table now has at least one verified member.
32//!
33//! But a shared **Family** tag is a claim about the backbone, not about the
34//! whole path, so four `Esm2` rows are verified in their own right rather than
35//! by inheritance (ferritin-100.34). `saprot-35m-af2` runs the same
36//! `EsmForMaskedLM` backbone as `esm2-t6-8m` yet reads *two* characters per
37//! residue over a 446-token (amino acid, 3Di) vocabulary from a bare
38//! `vocab.txt` — a different tokenizer path, and one that yields a perfectly
39//! well-formed tensor when read wrongly. `pepmlm-650m`, `dplm-650m` and
40//! `fastesm2-650` share the stock alphabet but are all pushed through
41//! `ESM2Config::t33_650m()`, which assumes each is plain ESM-2 650M with
42//! different weights — a real assumption, given `dplm-650m` is a diffusion
43//! model and `fastesm2-650` declares `model_type: fast_esm`. All four agree
44//! with the reference; the point is that they now do so demonstrably.
45//!
46//! The rows still unchecked are the ones genuinely shaped "same path,
47//! different weights" — the `esm2-t*` ladder, the five `esm1v` members,
48//! `esm1b`, `saprot-650m-af2`, `amplify-350m`, `ankh-large`, `esmc-600m` and
49//! `esmc-6b` — where inherited trust is reasonable.
50//!
51//! <!-- BEGIN SUPPORT MATRIX -->
52//! | Model | Family | Weights | Parity | Status |
53//! |---|---|---|---|---|
54//! | `esm2-t6-8m` | Esm2 | `facebook/esm2_t6_8M_UR50D` (safetensors) | verified (`esm2_parity`) | supported |
55//! | `esm2-t12-35m` | Esm2 | `facebook/esm2_t12_35M_UR50D` (safetensors) | **not checked** | supported |
56//! | `esm2-t30-150m` | Esm2 | `facebook/esm2_t30_150M_UR50D` (safetensors) | **not checked** | supported |
57//! | `esm2-t33-650m` | Esm2 | `facebook/esm2_t33_650M_UR50D` (safetensors) | **not checked** | supported |
58//! | `esm2-t36-3b` | Esm2 | `facebook/esm2_t36_3B_UR50D` (safetensors) | **not checked** | supported |
59//! | `esm2-t48-15b` | Esm2 | `facebook/esm2_t48_15B_UR50D` (safetensors) | **not checked** | supported |
60//! | `esm1v-t33-650m-ur90s-1` | Esm2 | `facebook/esm1v_t33_650M_UR90S_1` (pth) | **not checked** | supported |
61//! | `esm1v-t33-650m-ur90s-2` | Esm2 | `facebook/esm1v_t33_650M_UR90S_2` (pth) | **not checked** | supported |
62//! | `esm1v-t33-650m-ur90s-3` | Esm2 | `facebook/esm1v_t33_650M_UR90S_3` (pth) | **not checked** | supported |
63//! | `esm1v-t33-650m-ur90s-4` | Esm2 | `facebook/esm1v_t33_650M_UR90S_4` (pth) | **not checked** | supported |
64//! | `esm1v-t33-650m-ur90s-5` | Esm2 | `facebook/esm1v_t33_650M_UR90S_5` (pth) | **not checked** | supported |
65//! | `esm1b-t33-650m-ur50s` | Esm2 | `facebook/esm1b_t33_650M_UR50S` (pth) | **not checked** | supported |
66//! | `saprot-35m-af2` | Esm2 | `westlake-repl/SaProt_35M_AF2` (pth) | verified (`saprot_parity`) | supported |
67//! | `saprot-650m-af2` | Esm2 | `westlake-repl/SaProt_650M_AF2` (pth) | **not checked** | supported |
68//! | `fastesm2-650` | Esm2 | `Synthyra/FastESM2_650` (safetensors) | verified (`fastesm2_parity`) | supported |
69//! | `pepmlm-650m` | Esm2 | `ChatterjeeLab/PepMLM-650M` (pth) | verified (`pepmlm_parity`) | supported |
70//! | `dplm-650m` | Esm2 | `airkingbd/dplm_650m` (pth) | verified (`dplm_parity`) | supported |
71//! | `amplify-120m` | Amplify | `chandar-lab/AMPLIFY_120M` (safetensors) | verified (`amplify_parity`) | supported |
72//! | `amplify-350m` | Amplify | `chandar-lab/AMPLIFY_350M` (safetensors) | **not checked** | supported |
73//! | `esmc-300m` | Esmc | `EvolutionaryScale/esmc-300m-2024-12` (pth) | verified (`esmc_parity`) | supported |
74//! | `esmc-600m` | Esmc | `EvolutionaryScale/esmc-600m-2024-12` (pth) | **not checked** | supported |
75//! | `esmc-6b` | Esmc | `EvolutionaryScale/esmc-6b-2024-12` (safetensors) | **not checked** | supported |
76//! | `esm3-sm-open-v1` | Esm3 | `EvolutionaryScale/esm3-sm-open-v1` (pth) | verified (`esm3_parity`) | supported |
77//! | `esm3-structure-encoder-v0` | Esm3 | `EvolutionaryScale/esm3-sm-open-v1` (pth) | verified (`esm3_structure_parity`) | supported |
78//! | `prott5-xl-half-uniref50-enc` | T5 | `Rostlab/prot_t5_xl_half_uniref50-enc` (pth) | verified (`prott5_parity`) | supported |
79//! | `ankh-base` | T5 | `ElnaggarLab/ankh-base` (pth) | verified (`ankh_parity`) | supported |
80//! | `ankh-large` | T5 | `ElnaggarLab/ankh-large` (pth) | **not checked** | supported |
81//! | `prostt5-fp16` | T5 | `Rostlab/ProstT5_fp16` (pth) | verified (`prostt5_parity`) | supported |
82//! | `proteinmpnn-v48-020` | Mpnn | `zcpbx/ligandmpnn-weights` (pth) | verified (`proteinmpnn_parity`) | supported |
83//! | `ligandmpnn-v32-020-25` | Mpnn | `zcpbx/ligandmpnn-weights` (pth) | verified (`ligandmpnn_parity`) | supported |
84//! <!-- END SUPPORT MATRIX -->
85
86// The crate deliberately uses the `foo/mod.rs` + inner `mod foo` layout for
87// each model family, so module_inception is expected throughout.
88#![allow(clippy::module_inception)]
89
90pub use amplify::amplify::{AMPLIFY, AmplifyOutput};
91pub use amplify::amplify_runner::{AmplifyModels, AmplifyRunner};
92pub use amplify::config::AMPLIFYConfig;
93use candle_core::utils::{cuda_is_available, metal_is_available};
94use candle_core::{Device, Result};
95pub use esm2::esm2::{ESM2, ESM2Config};
96pub use esm2::esm2_runner::{ESM2Models, ESM2Runner};
97pub use esm3::models::esm3::ESM3Config;
98pub use esm3::pretrained::{ESM3Models, ESM3Runner};
99pub use esmc::models::esmc::{ESMC, ESMCConfig, ESMCOutput, LogitsConfig, LogitsOutput};
100pub use esmc::pretrained::{ESMCModels, ESMCRunner};
101pub use featurize::StructureFeatures;
102pub use ligandmpnn::configs::ProteinMPNNConfig;
103pub use ligandmpnn::model::ProteinMPNN;
104pub use ligandmpnn::pmpnn_runner::{ProteinMPNNModels, ProteinMPNNRunner};
105pub use t5::runner::{T5Models, T5Runner};
106pub use t5::translator::{ProstT5Models, ProstT5Translator};
107
108pub mod amplify;
109pub mod esm2;
110pub mod esm3;
111pub mod esmc;
112pub mod featurize;
113pub mod ligandmpnn;
114pub mod loader;
115pub mod plm_runner;
116pub mod registry;
117pub mod t5;
118pub mod types;
119pub mod utils;
120pub use plm_runner::{ModelMetadata, PlmRunner, SpecialTokenLayout};
121pub use registry::{Family, ModelCard, ParityStatus, REGISTRY, TokenizerSpec, VocabAlphabet};
122
123/// Returns the best available device for computation.
124///
125/// If `cpu` is true, always returns `Device::Cpu` regardless of available hardware.
126/// Otherwise prioritizes CUDA GPU if available, then Metal GPU on supported platforms,
127/// and falls back to CPU if no GPU acceleration is available.
128pub fn device(cpu: bool) -> Result<Device> {
129    if cpu {
130        return Ok(Device::Cpu);
131    }
132    if cuda_is_available() {
133        Ok(Device::new_cuda(0)?)
134    } else if metal_is_available() {
135        Ok(Device::new_metal(0)?)
136    } else {
137        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
138        {
139            println!(
140                "Running on CPU, to run on GPU(metal), build this example with `--features metal`"
141            );
142        }
143        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
144        {
145            println!("Running on CPU, to run on GPU, build this example with `--features cuda`");
146        }
147        Ok(Device::Cpu)
148    }
149}
ferritin_plms/lib.rs

ferritin_plms/
lib.rs
