ferritin_plms/
registry.rs

1//! One table describing every supported model.
2//!
3//! Before this, the same facts were spread across six enums that each encoded a
4//! different subset in a different shape: `AmplifyModels` returned a
5//! [`WeightSource`], `ESM2Models` a source plus a config, `ESMCModels` a source
6//! plus a filename plus a config wrapped in `Result`, and so on. Which
7//! tokenizer a model needs lived in its runner; how many special tokens it
8//! wraps lived in its `PlmRunner` impl; whether it had ever been checked
9//! against a Python reference lived only in the test suite.
10//!
11//! [`REGISTRY`] is where that belongs. A model is a data row.
12//!
13//! # Why a const table
14//!
15//! The variation between models is data — strings, dimensions, token counts —
16//! so it stays data. The moment it becomes `dyn`, the compiler stops checking
17//! that every model is fully specified, which is exactly the property worth
18//! having: adding a row that omits a field will not compile.
19//!
20//! [`Family`] is a closed enum for the same reason. Adding a backbone should be
21//! a deliberate, reviewed act rather than something that falls out of a string.
22//!
23//! ```
24//! # use ferritin_plms::registry::{REGISTRY, lookup};
25//! let card = lookup("esm2-t6-8m").expect("registered");
26//! assert_eq!(card.metadata.d_model, 320);
27//! assert_eq!(card.source.repo_id, "facebook/esm2_t6_8M_UR50D");
28//! ```
29
30use crate::loader::WeightSource;
31use crate::plm_runner::{ModelMetadata, SpecialTokenLayout};
32
33// ── Family ────────────────────────────────────────────────────────────────────
34
35/// Architecture family a model belongs to.
36///
37/// Closed on purpose: a new family means a new loader, which is a reviewed
38/// change rather than a new string. Families are added when a loader for them
39/// lands, so every variant here has at least one row in [`REGISTRY`].
40#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
41pub enum Family {
42    /// ESM-2 and its schema-compatible relatives (ESM-1v, SaProt).
43    Esm2,
44    /// AMPLIFY.
45    Amplify,
46    /// ESM Cambrian.
47    Esmc,
48    /// ESM3, the multi-track model.
49    Esm3,
50    /// ProteinMPNN / LigandMPNN — inverse folding, not an embedding model.
51    Mpnn,
52    /// ProtT5 and the rest of the T5 encoder family.
53    T5,
54}
55
56// ── TokenizerSpec ─────────────────────────────────────────────────────────────
57
58/// Where a model's tokenizer comes from.
59///
60/// Not incidental detail: the four ported families genuinely differ here, and
61/// that knowledge used to be buried in six separate runners.
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub enum TokenizerSpec {
64    /// A `tokenizer.json` downloaded from the model's own HF repo.
65    HfJson,
66    /// A `tokenizer.json` compiled into the binary with `include_bytes!`,
67    /// named by its path under `src/`.
68    Embedded(&'static str),
69    /// A hand-written vocabulary table in Rust, named by its module path.
70    BuiltinVocab(&'static str),
71    /// A bare `vocab.txt` from the model's repo — one token per line, the line
72    /// number being the id. SaProt ships this instead of a `tokenizer.json`,
73    /// and so do the BERT-style protein models, over a completely different
74    /// alphabet — so the file format alone does not say how to read a
75    /// sequence. The [`VocabAlphabet`] does.
76    HfVocabTxt(VocabAlphabet),
77    /// No tokenizer: the model consumes structure, not sequence.
78    None,
79}
80
81/// How to read a sequence against a bare `vocab.txt`.
82///
83/// Split out from [`TokenizerSpec::HfVocabTxt`] because that variant used to
84/// mean "SaProt" in practice: both the runner and the conformance suite
85/// branched on the bare variant to pick SaProt's two-character alphabet. A
86/// second `vocab.txt` model over a one-character alphabet — ProtBert was the
87/// concrete case — would silently have been fed `M#Q#I#…`, made every residue
88/// `<unk>`, and still passed the shape assertions (ferritin-goh.11).
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub enum VocabAlphabet {
91    /// SaProt: two characters per residue over a 20x20 (amino acid, 3Di)
92    /// product alphabet, with `<cls>`/`<pad>`/`<eos>`/`<unk>` specials.
93    SaProtPairs,
94    /// One character per residue over the plain amino-acid alphabet — the
95    /// BERT-style protein models, with `[PAD]`/`[UNK]`/`[CLS]`/`[SEP]`/
96    /// `[MASK]` specials.
97    SingleResidue,
98}
99
100// ── ParityStatus ──────────────────────────────────────────────────────────────
101
102/// Whether this model's numerics have been checked against a Python reference.
103///
104/// `Unverified` is the honest default and covers most rows. Read it as "this
105/// model's output could be anything", not as "this model is probably fine".
106///
107/// That wording used to be softer — it said `Unverified` was not a statement
108/// that a model is wrong, only that nothing proved it right. ProteinMPNN then
109/// demonstrated the difference is not academic: it sat at `Unverified` through
110/// every release up to v0.3.3 while agreeing with the reference on 2 of 93
111/// positions, where chance alone over a 21-token vocabulary is about 4. The
112/// output was not approximately right, it was unrelated to what the model
113/// computes, and the registry said only "unchecked" the whole time
114/// (ferritin-100.33).
115///
116/// The lesson is about which rows are dangerous. A row whose family already has
117/// a `Verified` sibling shares a proven code path and differs mainly in
118/// weights. A row whose family has *no* verified member has an entire
119/// architecture that nothing has ever checked — that is the position
120/// ProteinMPNN was in. Prefer closing whole-family gaps over adding a second
121/// fixture to a family that already has one.
122#[derive(Debug, Clone, Copy, PartialEq, Eq)]
123pub enum ParityStatus {
124    /// A committed fixture pins this model's output against the reference
125    /// implementation. Names the fixture stem under `tests/fixtures/`.
126    Verified {
127        /// Fixture stem, without the `.safetensors` extension.
128        fixture: &'static str,
129    },
130    /// No parity fixture exists, so nothing checks this model's numerics.
131    ///
132    /// Not a weaker form of `Verified` — an absence of evidence that has
133    /// already, once, been concealing a model that was flatly wrong.
134    Unverified,
135}
136
137// ── ModelCard ─────────────────────────────────────────────────────────────────
138
139/// Everything needed to identify, fetch, and load one model.
140#[derive(Debug, Clone, Copy)]
141pub struct ModelCard {
142    /// Stable kebab-case identifier, unique across the registry.
143    pub id: &'static str,
144    /// Architecture family.
145    pub family: Family,
146    /// Which repo and revision the weights live in, and their on-disk format.
147    pub source: WeightSource,
148    /// Path of the weight file within the repo.
149    pub file: &'static str,
150    /// Where the tokenizer comes from.
151    pub tokenizer: TokenizerSpec,
152    /// How many special tokens the tokenizer wraps a sequence in.
153    pub specials: SpecialTokenLayout,
154    /// Architecture dimensions.
155    ///
156    /// For the non-embedding families these are the nearest analogue rather
157    /// than a literal reading: `d_model` is the model's hidden width and
158    /// `vocab_size` its output alphabet.
159    pub metadata: ModelMetadata,
160    /// Rough in-memory footprint at F32, from the published parameter count.
161    ///
162    /// Approximate on purpose — it exists to tier CI and to warn before a
163    /// machine starts swapping, not to be exact. Load at F16 to roughly halve
164    /// it; see `LoadOptions::with_dtype`.
165    pub approx_bytes_f32: u64,
166    /// Whether anything checks this model's numerics.
167    pub parity: ParityStatus,
168    /// Set when the model cannot currently be loaded, saying why.
169    ///
170    /// Three ported models are in this state, each for a recorded reason. A
171    /// row that carries this is present because the model is part of the
172    /// public API, not because it works.
173    pub unsupported: Option<&'static str>,
174}
175
176impl ModelCard {
177    /// Whether this model can actually be loaded today.
178    pub const fn is_loadable(&self) -> bool {
179        self.unsupported.is_none()
180    }
181
182    /// Whether this model consumes sequence and can implement
183    /// [`PlmRunner`][crate::plm_runner::PlmRunner].
184    ///
185    /// A property of the card, not of its [`Family`]: ESM3 contains both a
186    /// sequence model and a VQ-VAE structure encoder that takes backbone
187    /// coordinates, so the family alone does not settle it.
188    pub const fn is_embedding_model(&self) -> bool {
189        !matches!(self.tokenizer, TokenizerSpec::None)
190    }
191
192    /// Lowercase family name, for matching against string-keyed tables.
193    pub const fn family_str(&self) -> &'static str {
194        match self.family {
195            Family::Esm2 => "esm2",
196            Family::Amplify => "amplify",
197            Family::Esmc => "esmc",
198            Family::Esm3 => "esm3",
199            // One generator covers the whole family — ProteinMPNN,
200            // LigandMPNN and SolubleMPNN share `model_utils.py`.
201            Family::Mpnn => "mpnn",
202            Family::T5 => "t5",
203        }
204    }
205}
206
207const GB: u64 = 1024 * 1024 * 1024;
208const MB: u64 = 1024 * 1024;
209
210/// Every model the public API exposes.
211///
212/// Rows carrying [`unsupported`][ModelCard::unsupported] are still listed: they
213/// are reachable from the public API, and a registry that hid them would
214/// misrepresent what a caller can name.
215pub const REGISTRY: &[ModelCard] = &[
216    // ── ESM-2 ────────────────────────────────────────────────────────────────
217    ModelCard {
218        id: "esm2-t6-8m",
219        family: Family::Esm2,
220        source: WeightSource::safetensors("facebook/esm2_t6_8M_UR50D").at_revision("main"),
221        file: "model.safetensors",
222        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
223        specials: SpecialTokenLayout::BOS_EOS,
224        metadata: ModelMetadata {
225            d_model: 320,
226            n_layers: 6,
227            vocab_size: 33,
228            max_positions: Some(1026),
229        },
230        approx_bytes_f32: 32 * MB,
231        parity: ParityStatus::Verified {
232            fixture: "esm2_parity",
233        },
234        unsupported: None,
235    },
236    ModelCard {
237        id: "esm2-t12-35m",
238        family: Family::Esm2,
239        source: WeightSource::safetensors("facebook/esm2_t12_35M_UR50D").at_revision("main"),
240        file: "model.safetensors",
241        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
242        specials: SpecialTokenLayout::BOS_EOS,
243        metadata: ModelMetadata {
244            d_model: 480,
245            n_layers: 12,
246            vocab_size: 33,
247            max_positions: Some(1026),
248        },
249        approx_bytes_f32: 140 * MB,
250        parity: ParityStatus::Unverified,
251        unsupported: None,
252    },
253    ModelCard {
254        id: "esm2-t30-150m",
255        family: Family::Esm2,
256        source: WeightSource::safetensors("facebook/esm2_t30_150M_UR50D").at_revision("main"),
257        file: "model.safetensors",
258        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
259        specials: SpecialTokenLayout::BOS_EOS,
260        metadata: ModelMetadata {
261            d_model: 640,
262            n_layers: 30,
263            vocab_size: 33,
264            max_positions: Some(1026),
265        },
266        approx_bytes_f32: 600 * MB,
267        parity: ParityStatus::Unverified,
268        unsupported: None,
269    },
270    ModelCard {
271        id: "esm2-t33-650m",
272        family: Family::Esm2,
273        source: WeightSource::safetensors("facebook/esm2_t33_650M_UR50D").at_revision("main"),
274        file: "model.safetensors",
275        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
276        specials: SpecialTokenLayout::BOS_EOS,
277        metadata: ModelMetadata {
278            d_model: 1280,
279            n_layers: 33,
280            vocab_size: 33,
281            max_positions: Some(1026),
282        },
283        approx_bytes_f32: 2 * GB + 600 * MB,
284        parity: ParityStatus::Unverified,
285        unsupported: None,
286    },
287    ModelCard {
288        id: "esm2-t36-3b",
289        family: Family::Esm2,
290        source: WeightSource::safetensors("facebook/esm2_t36_3B_UR50D").at_revision("main"),
291        file: "model.safetensors",
292        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
293        specials: SpecialTokenLayout::BOS_EOS,
294        metadata: ModelMetadata {
295            d_model: 2560,
296            n_layers: 36,
297            vocab_size: 33,
298            max_positions: Some(1026),
299        },
300        approx_bytes_f32: 12 * GB,
301        parity: ParityStatus::Unverified,
302        unsupported: None,
303    },
304    ModelCard {
305        id: "esm2-t48-15b",
306        family: Family::Esm2,
307        source: WeightSource::safetensors("facebook/esm2_t48_15B_UR50D").at_revision("main"),
308        file: "model.safetensors",
309        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
310        specials: SpecialTokenLayout::BOS_EOS,
311        metadata: ModelMetadata {
312            d_model: 5120,
313            n_layers: 48,
314            vocab_size: 33,
315            max_positions: Some(1026),
316        },
317        approx_bytes_f32: 60 * GB,
318        parity: ParityStatus::Unverified,
319        unsupported: None,
320    },
321    // ── ESM-1: learned absolute positions rather than rotary ─────────────────
322    //
323    // The five ESM-1v UR90S members are meant to be ensembled for zero-shot
324    // variant effect prediction, so all five are registered; they are
325    // near-identical cards distinguished only by weights (ferritin-goh.4).
326    ModelCard {
327        id: "esm1v-t33-650m-ur90s-1",
328        family: Family::Esm2,
329        source: WeightSource::pth("facebook/esm1v_t33_650M_UR90S_1", None),
330        file: "pytorch_model.bin",
331        // Vocabulary is byte-identical to ESM-2's, so the embedded tokenizer
332        // serves both (verified, ferritin-goh.4).
333        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
334        specials: SpecialTokenLayout::BOS_EOS,
335        metadata: ModelMetadata {
336            d_model: 1280,
337            n_layers: 33,
338            vocab_size: 33,
339            max_positions: Some(1026),
340        },
341        approx_bytes_f32: 2 * GB + 600 * MB,
342        parity: ParityStatus::Unverified,
343        unsupported: None,
344    },
345    ModelCard {
346        id: "esm1v-t33-650m-ur90s-2",
347        family: Family::Esm2,
348        source: WeightSource::pth("facebook/esm1v_t33_650M_UR90S_2", None),
349        file: "pytorch_model.bin",
350        // Vocabulary is byte-identical to ESM-2's, so the embedded tokenizer
351        // serves both (verified, ferritin-goh.4).
352        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
353        specials: SpecialTokenLayout::BOS_EOS,
354        metadata: ModelMetadata {
355            d_model: 1280,
356            n_layers: 33,
357            vocab_size: 33,
358            max_positions: Some(1026),
359        },
360        approx_bytes_f32: 2 * GB + 600 * MB,
361        parity: ParityStatus::Unverified,
362        unsupported: None,
363    },
364    ModelCard {
365        id: "esm1v-t33-650m-ur90s-3",
366        family: Family::Esm2,
367        source: WeightSource::pth("facebook/esm1v_t33_650M_UR90S_3", None),
368        file: "pytorch_model.bin",
369        // Vocabulary is byte-identical to ESM-2's, so the embedded tokenizer
370        // serves both (verified, ferritin-goh.4).
371        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
372        specials: SpecialTokenLayout::BOS_EOS,
373        metadata: ModelMetadata {
374            d_model: 1280,
375            n_layers: 33,
376            vocab_size: 33,
377            max_positions: Some(1026),
378        },
379        approx_bytes_f32: 2 * GB + 600 * MB,
380        parity: ParityStatus::Unverified,
381        unsupported: None,
382    },
383    ModelCard {
384        id: "esm1v-t33-650m-ur90s-4",
385        family: Family::Esm2,
386        source: WeightSource::pth("facebook/esm1v_t33_650M_UR90S_4", None),
387        file: "pytorch_model.bin",
388        // Vocabulary is byte-identical to ESM-2's, so the embedded tokenizer
389        // serves both (verified, ferritin-goh.4).
390        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
391        specials: SpecialTokenLayout::BOS_EOS,
392        metadata: ModelMetadata {
393            d_model: 1280,
394            n_layers: 33,
395            vocab_size: 33,
396            max_positions: Some(1026),
397        },
398        approx_bytes_f32: 2 * GB + 600 * MB,
399        parity: ParityStatus::Unverified,
400        unsupported: None,
401    },
402    ModelCard {
403        id: "esm1v-t33-650m-ur90s-5",
404        family: Family::Esm2,
405        source: WeightSource::pth("facebook/esm1v_t33_650M_UR90S_5", None),
406        file: "pytorch_model.bin",
407        // Vocabulary is byte-identical to ESM-2's, so the embedded tokenizer
408        // serves both (verified, ferritin-goh.4).
409        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
410        specials: SpecialTokenLayout::BOS_EOS,
411        metadata: ModelMetadata {
412            d_model: 1280,
413            n_layers: 33,
414            vocab_size: 33,
415            max_positions: Some(1026),
416        },
417        approx_bytes_f32: 2 * GB + 600 * MB,
418        parity: ParityStatus::Unverified,
419        unsupported: None,
420    },
421    ModelCard {
422        id: "esm1b-t33-650m-ur50s",
423        family: Family::Esm2,
424        source: WeightSource::pth("facebook/esm1b_t33_650M_UR50S", None),
425        file: "pytorch_model.bin",
426        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
427        specials: SpecialTokenLayout::BOS_EOS,
428        metadata: ModelMetadata {
429            d_model: 1280,
430            n_layers: 33,
431            vocab_size: 33,
432            max_positions: Some(1026),
433        },
434        approx_bytes_f32: 2 * GB + 600 * MB,
435        parity: ParityStatus::Unverified,
436        unsupported: None,
437    },
438    // ── SaProt: ESM-2 architecture over a structure-aware alphabet ───────────
439    ModelCard {
440        id: "saprot-35m-af2",
441        family: Family::Esm2,
442        source: WeightSource::pth("westlake-repl/SaProt_35M_AF2", None),
443        // No safetensors in this repo; weights are a PyTorch pickle.
444        file: "pytorch_model.bin",
445        tokenizer: TokenizerSpec::HfVocabTxt(VocabAlphabet::SaProtPairs),
446        specials: SpecialTokenLayout::BOS_EOS,
447        metadata: ModelMetadata {
448            d_model: 480,
449            n_layers: 12,
450            // 20x20 (amino acid, 3Di) pairs plus specials.
451            vocab_size: 446,
452            max_positions: Some(1026),
453        },
454        approx_bytes_f32: 140 * MB,
455        parity: ParityStatus::Verified {
456            fixture: "saprot_parity",
457        },
458        unsupported: None,
459    },
460    ModelCard {
461        id: "saprot-650m-af2",
462        family: Family::Esm2,
463        source: WeightSource::pth("westlake-repl/SaProt_650M_AF2", None),
464        file: "pytorch_model.bin",
465        tokenizer: TokenizerSpec::HfVocabTxt(VocabAlphabet::SaProtPairs),
466        specials: SpecialTokenLayout::BOS_EOS,
467        metadata: ModelMetadata {
468            d_model: 1280,
469            n_layers: 33,
470            vocab_size: 446,
471            max_positions: Some(1026),
472        },
473        approx_bytes_f32: 2 * GB + 600 * MB,
474        parity: ParityStatus::Unverified,
475        unsupported: None,
476    },
477    // ── Modern checkpoints on ESM-2's exact schema (ferritin-goh.12) ─────────
478    //
479    // These three publish a config.json byte-identical in shape to
480    // facebook/esm2_t33_650M_UR50D, and their tensor names are ESM-2's, so
481    // they need no new architecture code. Layouts were verified by reading
482    // each checkpoint's header remotely rather than by downloading 2.6 GB
483    // apiece — see ferritin-goh.12 for the method.
484    ModelCard {
485        id: "fastesm2-650",
486        family: Family::Esm2,
487        // The only safetensors ESM-2-family 650M checkpoint found, which makes
488        // it the cheapest way to exercise the safetensors path at this size.
489        source: WeightSource::safetensors("Synthyra/FastESM2_650"),
490        file: "model.safetensors",
491        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
492        specials: SpecialTokenLayout::BOS_EOS,
493        metadata: ModelMetadata {
494            d_model: 1280,
495            n_layers: 33,
496            vocab_size: 33,
497            max_positions: Some(1026),
498        },
499        approx_bytes_f32: 2 * GB + 600 * MB,
500        parity: ParityStatus::Verified {
501            fixture: "fastesm2_parity",
502        },
503        unsupported: None,
504    },
505    ModelCard {
506        id: "pepmlm-650m",
507        family: Family::Esm2,
508        // Fine-tuned from ESM-2 650M for peptide binder design; its
509        // config.json still records _name_or_path facebook/esm2_t33_650M_UR50D.
510        source: WeightSource::pth("ChatterjeeLab/PepMLM-650M", None),
511        file: "pytorch_model.bin",
512        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
513        specials: SpecialTokenLayout::BOS_EOS,
514        metadata: ModelMetadata {
515            d_model: 1280,
516            n_layers: 33,
517            vocab_size: 33,
518            max_positions: Some(1026),
519        },
520        approx_bytes_f32: 2 * GB + 600 * MB,
521        parity: ParityStatus::Verified {
522            fixture: "pepmlm_parity",
523        },
524        unsupported: None,
525    },
526    ModelCard {
527        id: "dplm-650m",
528        family: Family::Esm2,
529        // DPLM is a discrete DIFFUSION model initialised from ESM-2 650M. Its
530        // hidden states are ordinary embeddings and are what this crate
531        // exposes; its lm_head denoises rather than scoring masked tokens, so
532        // reading its logits as masked-LM scores would be wrong. See the
533        // `logits` note on ESM2Runner.
534        source: WeightSource::pth("airkingbd/dplm_650m", None),
535        file: "pytorch_model.bin",
536        tokenizer: TokenizerSpec::Embedded("esm2/tokenizer.json"),
537        specials: SpecialTokenLayout::BOS_EOS,
538        metadata: ModelMetadata {
539            d_model: 1280,
540            n_layers: 33,
541            vocab_size: 33,
542            max_positions: Some(1026),
543        },
544        approx_bytes_f32: 2 * GB + 600 * MB,
545        parity: ParityStatus::Verified {
546            fixture: "dplm_parity",
547        },
548        unsupported: None,
549    },
550    // ── AMPLIFY ──────────────────────────────────────────────────────────────
551    ModelCard {
552        id: "amplify-120m",
553        family: Family::Amplify,
554        source: WeightSource::safetensors("chandar-lab/AMPLIFY_120M").at_revision("main"),
555        file: "model.safetensors",
556        // The runner downloads the repo's tokenizer.json; an embedded copy also
557        // exists for AMPLIFY::load_tokenizer.
558        tokenizer: TokenizerSpec::HfJson,
559        specials: SpecialTokenLayout::BOS_EOS,
560        metadata: ModelMetadata {
561            d_model: 640,
562            n_layers: 24,
563            vocab_size: 27,
564            max_positions: Some(2048),
565        },
566        approx_bytes_f32: 480 * MB,
567        parity: ParityStatus::Verified {
568            fixture: "amplify_parity",
569        },
570        unsupported: None,
571    },
572    ModelCard {
573        id: "amplify-350m",
574        family: Family::Amplify,
575        source: WeightSource::safetensors("chandar-lab/AMPLIFY_350M").at_revision("main"),
576        file: "model.safetensors",
577        tokenizer: TokenizerSpec::HfJson,
578        specials: SpecialTokenLayout::BOS_EOS,
579        metadata: ModelMetadata {
580            d_model: 960,
581            n_layers: 32,
582            vocab_size: 27,
583            max_positions: Some(2048),
584        },
585        approx_bytes_f32: GB + 400 * MB,
586        parity: ParityStatus::Unverified,
587        unsupported: None,
588    },
589    // ── ESM Cambrian ─────────────────────────────────────────────────────────
590    ModelCard {
591        id: "esmc-300m",
592        family: Family::Esmc,
593        source: WeightSource::pth("EvolutionaryScale/esmc-300m-2024-12", None),
594        file: "data/weights/esmc_300m_2024_12_v0.pth",
595        tokenizer: TokenizerSpec::BuiltinVocab("esmc::tokenizer::EsmSequenceTokenizer"),
596        specials: SpecialTokenLayout::BOS_EOS,
597        metadata: ModelMetadata {
598            d_model: 960,
599            n_layers: 30,
600            vocab_size: 64,
601            // Rotary positions: no hard architectural cap.
602            max_positions: None,
603        },
604        approx_bytes_f32: GB + 200 * MB,
605        parity: ParityStatus::Verified {
606            fixture: "esmc_parity",
607        },
608        unsupported: None,
609    },
610    ModelCard {
611        id: "esmc-600m",
612        family: Family::Esmc,
613        source: WeightSource::pth("EvolutionaryScale/esmc-600m-2024-12", None),
614        file: "data/weights/esmc_600m_2024_12_v0.pth",
615        tokenizer: TokenizerSpec::BuiltinVocab("esmc::tokenizer::EsmSequenceTokenizer"),
616        specials: SpecialTokenLayout::BOS_EOS,
617        metadata: ModelMetadata {
618            d_model: 1152,
619            n_layers: 36,
620            vocab_size: 64,
621            max_positions: None,
622        },
623        approx_bytes_f32: 2 * GB + 400 * MB,
624        parity: ParityStatus::Unverified,
625        unsupported: None,
626    },
627    ModelCard {
628        id: "esmc-6b",
629        family: Family::Esmc,
630        source: WeightSource::safetensors("EvolutionaryScale/esmc-6b-2024-12"),
631        // Sharded across six files; the loader follows this index
632        // (ferritin-100.24).
633        file: "model.safetensors.index.json",
634        tokenizer: TokenizerSpec::BuiltinVocab("esmc::tokenizer::EsmSequenceTokenizer"),
635        specials: SpecialTokenLayout::BOS_EOS,
636        metadata: ModelMetadata {
637            d_model: 2560,
638            n_layers: 80,
639            vocab_size: 64,
640            max_positions: None,
641        },
642        approx_bytes_f32: 24 * GB,
643        parity: ParityStatus::Unverified,
644        unsupported: None,
645    },
646    // ── ESM3 ─────────────────────────────────────────────────────────────────
647    ModelCard {
648        id: "esm3-sm-open-v1",
649        family: Family::Esm3,
650        source: WeightSource::pth("EvolutionaryScale/esm3-sm-open-v1", None),
651        file: "data/weights/esm3_sm_open_v1.pth",
652        tokenizer: TokenizerSpec::BuiltinVocab("esm3::tokenization::sequence"),
653        specials: SpecialTokenLayout::BOS_EOS,
654        metadata: ModelMetadata {
655            d_model: 1536,
656            n_layers: 48,
657            vocab_size: 64,
658            max_positions: None,
659        },
660        approx_bytes_f32: 5 * GB + 600 * MB,
661        parity: ParityStatus::Verified {
662            fixture: "esm3_parity",
663        },
664        unsupported: None,
665    },
666    ModelCard {
667        id: "esm3-structure-encoder-v0",
668        family: Family::Esm3,
669        source: WeightSource::pth("EvolutionaryScale/esm3-sm-open-v1", None),
670        file: "data/weights/esm3_structure_encoder_v0.pth",
671        // Consumes backbone coordinates and emits structure tokens.
672        tokenizer: TokenizerSpec::None,
673        specials: SpecialTokenLayout::NONE,
674        metadata: ModelMetadata {
675            d_model: 1024,
676            n_layers: 2,
677            // Codebook size: the structure-token alphabet it emits.
678            vocab_size: 4096,
679            max_positions: None,
680        },
681        approx_bytes_f32: 30 * MB,
682        parity: ParityStatus::Verified {
683            fixture: "esm3_structure_parity",
684        },
685        unsupported: None,
686    },
687    // ── ESMFold2 ─────────────────────────────────────────────────────────────
688    // ── ProtT5: a T5 encoder, architecturally independent of ESM (goh.5) ─────
689    ModelCard {
690        id: "prott5-xl-half-uniref50-enc",
691        family: Family::T5,
692        // No safetensors in this repo — the only weight file is a zip-pickle
693        // `pytorch_model.bin`, with the tensors at the root.
694        source: WeightSource::pth("Rostlab/prot_t5_xl_half_uniref50-enc", None).at_revision("main"),
695        file: "pytorch_model.bin",
696        // No `tokenizer.json` either; the repo ships a SentencePiece
697        // `spiece.model` whose reachable vocabulary is 28 pieces, transcribed
698        // into `t5::tokenizer`.
699        tokenizer: TokenizerSpec::BuiltinVocab("t5::tokenizer"),
700        // T5 appends `</s>` and prepends nothing — the only non-BOS_EOS
701        // embedding model in the registry.
702        specials: SpecialTokenLayout::EOS_ONLY,
703        metadata: ModelMetadata {
704            d_model: 1024,
705            n_layers: 24,
706            // `shared.weight` is [128, 1024]; ids 28..=127 are T5's unused
707            // `<extra_id_*>` sentinels.
708            vocab_size: 128,
709            // Relative position buckets, so no hard cap.
710            max_positions: None,
711        },
712        // 1.2B encoder parameters. Published as float16 (2.4 GB on disk), and
713        // `T5Runner::from_pretrained` loads it at F16 rather than doubling
714        // it to reach F32.
715        approx_bytes_f32: 4800 * MB,
716        parity: ParityStatus::Verified {
717            fixture: "prott5_parity",
718        },
719        unsupported: None,
720    },
721    // ── Ankh: a T5 encoder with a gated FFN (ferritin-goh.6) ─────────────────
722    //
723    // Same residue alphabet and ids as ProtT5 (A=3 .. Z=27, frequency-ordered)
724    // and the same EOS-only layout, so `t5::tokenizer` serves both. The
725    // difference is the FFN: `feed_forward_proj: "gated-gelu"`, which candle
726    // loads as a gated `T5DenseGatedActDense` rather than ProtT5's plain ReLU
727    // dense. Both models therefore share one runner.
728    ModelCard {
729        id: "ankh-base",
730        family: Family::T5,
731        // Ships a full encoder-decoder; only the encoder is loaded.
732        source: WeightSource::pth("ElnaggarLab/ankh-base", None).at_revision("main"),
733        file: "pytorch_model.bin",
734        // The repo ships a real tokenizer.json, but its Unigram vocabulary is
735        // the same alphabet at the same ids as ProtT5's SentencePiece one, so
736        // the built-in table is used for both. The parity fixture carries the
737        // reference token ids, so a divergence fails there rather than being
738        // assumed away.
739        tokenizer: TokenizerSpec::BuiltinVocab("t5::tokenizer"),
740        specials: SpecialTokenLayout::EOS_ONLY,
741        metadata: ModelMetadata {
742            d_model: 768,
743            // Encoder layers. The checkpoint also carries 24 decoder layers,
744            // which the embedding path never loads.
745            n_layers: 48,
746            vocab_size: 144,
747            // Relative position buckets, so no hard cap.
748            max_positions: None,
749        },
750        approx_bytes_f32: 2950 * MB,
751        parity: ParityStatus::Verified {
752            fixture: "ankh_parity",
753        },
754        unsupported: None,
755    },
756    ModelCard {
757        id: "ankh-large",
758        family: Family::T5,
759        source: WeightSource::pth("ElnaggarLab/ankh-large", None).at_revision("main"),
760        file: "pytorch_model.bin",
761        tokenizer: TokenizerSpec::BuiltinVocab("t5::tokenizer"),
762        specials: SpecialTokenLayout::EOS_ONLY,
763        metadata: ModelMetadata {
764            d_model: 1536,
765            n_layers: 48,
766            vocab_size: 144,
767            max_positions: None,
768        },
769        approx_bytes_f32: 7520 * MB,
770        parity: ParityStatus::Unverified,
771        unsupported: None,
772    },
773    // ── ProstT5: the only encoder-decoder here (ferritin-goh.6) ──────────────
774    //
775    // Translates between amino acids and Foldseek's 3Di structural alphabet in
776    // both directions. Its useful output is generated tokens, not an
777    // embedding, so it is driven by `ProstT5Translator` rather than PlmRunner —
778    // `is_embedding_model()` is false for it, like the structure models.
779    ModelCard {
780        id: "prostt5-fp16",
781        family: Family::T5,
782        // The F32 `Rostlab/ProstT5` is the same weights at 11.3 GB.
783        source: WeightSource::pth("Rostlab/ProstT5_fp16", None).at_revision("main"),
784        file: "pytorch_model.bin",
785        // No tokenizer here: ProstT5 consumes and produces two alphabets at
786        // once and is not a sequence-embedding model, so the special-token
787        // contract PlmRunner enforces does not apply. Its vocabulary lives in
788        // `t5::tokenizer` alongside the rest of the family.
789        tokenizer: TokenizerSpec::None,
790        specials: SpecialTokenLayout::NONE,
791        metadata: ModelMetadata {
792            d_model: 1024,
793            n_layers: 24,
794            // 128 ProtT5 ids + 20 lowercase 3Di states + the two direction
795            // tokens <fold2AA> (148) and <AA2fold> (149).
796            vocab_size: 150,
797            max_positions: None,
798        },
799        approx_bytes_f32: 11280 * MB,
800        parity: ParityStatus::Verified {
801            fixture: "prostt5_parity",
802        },
803        unsupported: None,
804    },
805    // ── ProteinMPNN ──────────────────────────────────────────────────────────
806    ModelCard {
807        id: "proteinmpnn-v48-020",
808        family: Family::Mpnn,
809        source: WeightSource::pth("zcpbx/ligandmpnn-weights", Some("model_state_dict"))
810            .at_revision("main"),
811        file: "model_params/proteinmpnn_v_48_020.pt",
812        // Consumes structure; residues are encoded by the featurizer.
813        tokenizer: TokenizerSpec::None,
814        specials: SpecialTokenLayout::NONE,
815        metadata: ModelMetadata {
816            d_model: 128,
817            // 3 encoder + 3 decoder.
818            n_layers: 6,
819            // 20 amino acids plus X.
820            vocab_size: 21,
821            max_positions: None,
822        },
823        approx_bytes_f32: 7 * MB,
824        parity: ParityStatus::Verified {
825            fixture: "proteinmpnn_parity",
826        },
827        unsupported: None,
828    },
829    ModelCard {
830        id: "ligandmpnn-v32-020-25",
831        family: Family::Mpnn,
832        source: WeightSource::pth("zcpbx/ligandmpnn-weights", Some("model_state_dict"))
833            .at_revision("main"),
834        file: "model_params/ligandmpnn_v_32_020_25.pt",
835        // Consumes structure and ligand atoms; no sequence tokenizer.
836        tokenizer: TokenizerSpec::None,
837        specials: SpecialTokenLayout::NONE,
838        metadata: ModelMetadata {
839            d_model: 128,
840            // 3 encoder + 3 decoder, plus 2 context and 2 ligand-graph
841            // rounds that ProteinMPNN does not have.
842            n_layers: 10,
843            vocab_size: 21,
844            max_positions: None,
845        },
846        approx_bytes_f32: 11 * MB,
847        parity: ParityStatus::Verified {
848            fixture: "ligandmpnn_parity",
849        },
850        unsupported: None,
851    },
852];
853
854// ── Support matrix ────────────────────────────────────────────────────────────
855
856/// Render [`REGISTRY`] as a markdown support matrix.
857///
858/// The column that matters is **Parity**. "Does it compile" and even "does it
859/// load" are not what a user needs to know before trusting a number — what
860/// they need is whether anyone has ever compared this port's output against
861/// the reference implementation. Today two models have, and the table says so
862/// rather than leaving it to be inferred from which fixtures happen to exist
863/// (ferritin-100.13).
864///
865/// The copy embedded in the crate docs is checked against this function by
866/// `test_lib_rs_support_matrix_is_current`, so the two cannot drift.
867pub fn support_matrix_markdown() -> String {
868    let mut out = String::new();
869    out.push_str("| Model | Family | Weights | Parity | Status |\n");
870    out.push_str("|---|---|---|---|---|\n");
871
872    for card in REGISTRY {
873        let format = match card.source.format {
874            crate::loader::Format::Safetensors => "safetensors",
875            crate::loader::Format::Pth { .. } => "pth",
876        };
877        let parity = match card.parity {
878            ParityStatus::Verified { fixture } => {
879                format!("verified (`{fixture}`)")
880            }
881            ParityStatus::Unverified => "**not checked**".to_string(),
882        };
883        let status = match card.unsupported {
884            None => "supported".to_string(),
885            Some(reason) => {
886                // Keep the table readable; the full reason lives on the card.
887                let short = reason.split(" (ferritin-").next().unwrap_or(reason);
888                let issue = reason
889                    .rsplit_once("(ferritin-")
890                    .map(|(_, tail)| tail.trim_end_matches(')'))
891                    .unwrap_or("");
892                let first = short.split(&[',', ':'][..]).next().unwrap_or(short);
893                if issue.is_empty() {
894                    format!("**unsupported** — {first}")
895                } else {
896                    format!("**unsupported** — {first} (ferritin-{issue})")
897                }
898            }
899        };
900        out.push_str(&format!(
901            "| `{}` | {:?} | `{}` ({format}) | {parity} | {status} |\n",
902            card.id, card.family, card.source.repo_id,
903        ));
904    }
905    out
906}
907
908// ── Lookups ───────────────────────────────────────────────────────────────────
909
910/// Find a model by its registry id.
911pub fn lookup(id: &str) -> Option<&'static ModelCard> {
912    REGISTRY.iter().find(|c| c.id == id)
913}
914
915/// Every card in a family.
916pub fn by_family(family: Family) -> impl Iterator<Item = &'static ModelCard> {
917    REGISTRY.iter().filter(move |c| c.family == family)
918}
919
920/// Every card that can actually be loaded today.
921pub fn loadable() -> impl Iterator<Item = &'static ModelCard> {
922    REGISTRY.iter().filter(|c| c.is_loadable())
923}
924
925// ── Tests ─────────────────────────────────────────────────────────────────────
926
927#[cfg(test)]
928mod tests {
929    use super::*;
930    use std::collections::HashSet;
931
932    #[test]
933    fn test_ids_are_unique() {
934        let mut seen = HashSet::new();
935        for card in REGISTRY {
936            assert!(seen.insert(card.id), "duplicate registry id: {}", card.id);
937        }
938    }
939
940    /// Ids are the registry's public handle, so they stay kebab-case and free
941    /// of the underscores and capitals the upstream repo names use.
942    #[test]
943    fn test_ids_are_kebab_case() {
944        for card in REGISTRY {
945            assert!(
946                card.id
947                    .chars()
948                    .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-'),
949                "id {:?} should be lowercase kebab-case",
950                card.id
951            );
952        }
953    }
954
955    /// Every repo id must be well formed, or the failure only shows up as a
956    /// confusing 404 after a network round-trip.
957    #[test]
958    fn test_every_source_repo_is_well_formed() {
959        for card in REGISTRY {
960            assert!(
961                card.source.repo_id.split('/').count() == 2
962                    && !card.source.repo_id.starts_with('/')
963                    && !card.source.repo_id.ends_with('/'),
964                "{}: malformed repo id {:?}",
965                card.id,
966                card.source.repo_id
967            );
968            assert!(!card.file.is_empty(), "{}: empty weight filename", card.id);
969        }
970    }
971
972    #[test]
973    fn test_metadata_dimensions_are_plausible() {
974        for card in REGISTRY {
975            assert!(card.metadata.d_model > 0, "{}: zero d_model", card.id);
976            assert!(card.metadata.n_layers > 0, "{}: zero n_layers", card.id);
977            assert!(
978                card.approx_bytes_f32 > 0,
979                "{}: zero approx_bytes_f32",
980                card.id
981            );
982        }
983    }
984
985    /// Embedding models tokenize and therefore have a vocabulary; the
986    /// structure models deliberately have neither.
987    #[test]
988    fn test_embedding_models_have_a_tokenizer_and_vocab() {
989        for card in REGISTRY {
990            if card.is_embedding_model() {
991                assert_ne!(
992                    card.tokenizer,
993                    TokenizerSpec::None,
994                    "{}: an embedding model needs a tokenizer",
995                    card.id
996                );
997                assert!(
998                    card.metadata.vocab_size > 0,
999                    "{}: an embedding model needs a vocabulary",
1000                    card.id
1001                );
1002            } else {
1003                assert_eq!(
1004                    card.metadata.max_positions, None,
1005                    "{}: a structure model has no token positions to cap",
1006                    card.id
1007                );
1008            }
1009        }
1010    }
1011
1012    /// Every family variant must have at least one row. Family is closed so
1013    /// that adding a backbone is deliberate; a variant with no models means
1014    /// the enum has drifted ahead of the loaders.
1015    #[test]
1016    fn test_every_family_has_at_least_one_model() {
1017        for family in [
1018            Family::Esm2,
1019            Family::Amplify,
1020            Family::Esmc,
1021            Family::Esm3,
1022            Family::Mpnn,
1023        ] {
1024            assert!(
1025                by_family(family).next().is_some(),
1026                "{family:?} has no models; drop the variant or add its loader"
1027            );
1028        }
1029    }
1030
1031    /// The three unsupported models are the ones with recorded reasons. This
1032    /// pins the set so that a model silently becoming unloadable — or quietly
1033    /// staying that way after a fix — shows up here.
1034    #[test]
1035    fn test_unsupported_models_are_the_known_set() {
1036        let mut unsupported: Vec<&str> = REGISTRY
1037            .iter()
1038            .filter(|c| !c.is_loadable())
1039            .map(|c| c.id)
1040            .collect();
1041        unsupported.sort_unstable();
1042        // Empty since ferritin-100.17 deleted the ESMFold2 port: a row that
1043        // can never load is a promise the crate cannot keep, so the model is
1044        // absent from the registry rather than listed as broken.
1045        assert!(
1046            unsupported.is_empty(),
1047            "every registered model should load; got {unsupported:?}"
1048        );
1049
1050        for card in REGISTRY.iter().filter(|c| !c.is_loadable()) {
1051            let reason = card.unsupported.unwrap();
1052            assert!(
1053                reason.contains("ferritin-"),
1054                "{}: an unsupported reason should cite its tracking issue; got: {reason}",
1055                card.id
1056            );
1057        }
1058    }
1059
1060    /// Parity claims must name a fixture that the test suite actually has.
1061    /// Fourteen models are verified today.
1062    ///
1063    /// ESMC-300M closed the last whole-family gap: before it, `Family::Esmc`
1064    /// held three rows and no fixture at all (ferritin-100.33).
1065    ///
1066    /// The four `Family::Esm2` rows that are not stock ESM-2 — SaProt-35M,
1067    /// PepMLM-650M, DPLM-650M and FastESM2-650 — were then verified for a
1068    /// different reason: not because their family was uncovered, but because
1069    /// the family tag is a claim about the BACKBONE and each of them diverges
1070    /// somewhere else (ferritin-100.34). SaProt reads two characters per
1071    /// residue over 446 tokens from a bare `vocab.txt`; the other three are all
1072    /// run through `ESM2Config::t33_650m()` despite DPLM being a diffusion
1073    /// model and FastESM2 declaring `model_type: fast_esm`. All four agree with
1074    /// the reference, so those assumptions were sound — but they were
1075    /// assumptions until checked.
1076    #[test]
1077    fn test_verified_models_name_a_real_fixture() {
1078        let mut verified: Vec<(&str, &str)> = REGISTRY
1079            .iter()
1080            .filter_map(|c| match c.parity {
1081                ParityStatus::Verified { fixture } => Some((c.id, fixture)),
1082                ParityStatus::Unverified => None,
1083            })
1084            .collect();
1085        verified.sort_unstable();
1086        assert_eq!(
1087            verified,
1088            [
1089                ("amplify-120m", "amplify_parity"),
1090                ("ankh-base", "ankh_parity"),
1091                ("dplm-650m", "dplm_parity"),
1092                ("esm2-t6-8m", "esm2_parity"),
1093                ("esm3-sm-open-v1", "esm3_parity"),
1094                ("esm3-structure-encoder-v0", "esm3_structure_parity"),
1095                ("esmc-300m", "esmc_parity"),
1096                ("fastesm2-650", "fastesm2_parity"),
1097                ("ligandmpnn-v32-020-25", "ligandmpnn_parity"),
1098                ("pepmlm-650m", "pepmlm_parity"),
1099                ("prostt5-fp16", "prostt5_parity"),
1100                ("proteinmpnn-v48-020", "proteinmpnn_parity"),
1101                ("prott5-xl-half-uniref50-enc", "prott5_parity"),
1102                ("saprot-35m-af2", "saprot_parity"),
1103            ],
1104            "the set of parity-verified models changed; that is a deliberate act"
1105        );
1106    }
1107
1108    #[test]
1109    fn test_lookup_finds_and_misses() {
1110        assert_eq!(lookup("esm2-t6-8m").map(|c| c.id), Some("esm2-t6-8m"));
1111        assert!(lookup("no-such-model").is_none());
1112    }
1113
1114    #[test]
1115    fn test_loadable_excludes_unsupported() {
1116        assert!(loadable().all(|c| c.unsupported.is_none()));
1117        assert_eq!(
1118            loadable().count(),
1119            REGISTRY.len(),
1120            "no model is currently unsupported, so loadable() should be every row"
1121        );
1122    }
1123}
1124
1125#[cfg(test)]
1126mod matrix {
1127    use super::*;
1128
1129    /// Regeneration helper: prints the matrix for pasting into `lib.rs`.
1130    ///
1131    /// ```shell
1132    /// cargo test -p ferritin-plms --lib print_support_matrix -- --ignored --nocapture
1133    /// ```
1134    #[test]
1135    #[ignore = "prints the matrix for copying into lib.rs"]
1136    fn print_support_matrix() {
1137        println!("{}", support_matrix_markdown());
1138    }
1139
1140    /// The copy in the crate docs must match what the registry renders.
1141    ///
1142    /// A stale matrix is worse than none: it would tell a user a model is
1143    /// parity-verified, or supported, after that stopped being true.
1144    #[test]
1145    fn test_lib_rs_support_matrix_is_current() {
1146        const LIB_RS: &str = include_str!("lib.rs");
1147        const BEGIN: &str = "//! <!-- BEGIN SUPPORT MATRIX -->";
1148        const END: &str = "//! <!-- END SUPPORT MATRIX -->";
1149
1150        let start = LIB_RS
1151            .find(BEGIN)
1152            .expect("lib.rs should carry a BEGIN SUPPORT MATRIX marker")
1153            + BEGIN.len();
1154        let end = LIB_RS
1155            .find(END)
1156            .expect("lib.rs should carry an END SUPPORT MATRIX marker");
1157
1158        let embedded: String = LIB_RS[start..end]
1159            .lines()
1160            .filter(|l| !l.trim().is_empty())
1161            .map(|l| {
1162                format!(
1163                    "{}\n",
1164                    l.trim_start().trim_start_matches("//!").trim_start()
1165                )
1166            })
1167            .collect();
1168
1169        let rendered: String = support_matrix_markdown()
1170            .lines()
1171            .map(|l| format!("{l}\n"))
1172            .collect();
1173
1174        assert_eq!(
1175            embedded, rendered,
1176            "the support matrix in lib.rs is stale. Regenerate it with:\n  \
1177             cargo test -p ferritin-plms --lib print_support_matrix -- --ignored --nocapture\n\
1178             then replace the block between the SUPPORT MATRIX markers."
1179        );
1180    }
1181
1182    /// Every model in the matrix carries an explicit parity verdict.
1183    #[test]
1184    fn test_matrix_states_parity_for_every_model() {
1185        let matrix = support_matrix_markdown();
1186        for card in REGISTRY {
1187            let row = matrix
1188                .lines()
1189                .find(|l| l.contains(&format!("`{}`", card.id)))
1190                .unwrap_or_else(|| panic!("{} missing from the matrix", card.id));
1191            assert!(
1192                row.contains("verified") || row.contains("not checked"),
1193                "{}: the matrix must state a parity verdict; got: {row}",
1194                card.id
1195            );
1196        }
1197    }
1198}
ferritin_plms/registry.rs

ferritin_plms/
registry.rs
