ferritin_plms/t5/
mod.rs

1//! The T5 family: ProtT5, Ankh and ProstT5 (ferritin-goh.5, ferritin-goh.6).
2//!
3//! ProtT5 embeddings remain one of the most widely used protein
4//! representations, and they give the crate a second, architecturally
5//! independent embedding family to cross-check against ESM: a T5 encoder with
6//! relative position buckets rather than a BERT-style stack with rotary or
7//! learned absolute positions. Ankh is a third, arriving essentially for free.
8//!
9//! ProstT5 is the odd one out and lives in [`translator`]: it is an
10//! encoder-**decoder** whose useful output is generated tokens — a Foldseek 3Di
11//! structural string — rather than an embedding, so it stays outside
12//! `PlmRunner` instead of being contorted to fit it. Its vocabulary is still
13//! this family's: [`tokenizer`] holds the amino-acid table *and* the 3Di states
14//! that sit 125 ids above it.
15//!
16//! # One runner for the two encoders
17//!
18//! ProtT5 and Ankh differ in nothing [`runner`] cares about. Their residue
19//! alphabets are the *same letters at the same ids* — `A` = 3, `L` = 4, …
20//! `Z` = 27, frequency-ordered — and both wrap a sequence in a trailing `</s>`
21//! and no BOS. So [`tokenizer`] serves both, and the differences that remain
22//! come out of each checkpoint's own `config.json`:
23//!
24//! | | ProtT5-XL | Ankh-base | Ankh-large |
25//! |---|---|---|---|
26//! | `d_model` | 1024 | 768 | 1536 |
27//! | encoder layers | 24 | 48 | 48 |
28//! | `feed_forward_proj` | `relu` | `gated-gelu` | `gated-gelu` |
29//! | published dtype | F16 | F32 | F32 |
30//!
31//! The gated FFN is the one substantive difference: candle loads it as
32//! `T5DenseGatedActDense` rather than the plain dense ProtT5 uses.
33//!
34//! Ankh ships its own `tokenizer.json` (a Unigram vocabulary with no
35//! SentencePiece boundary marker, so it takes bare sequences), and the runner
36//! deliberately does not use it. That reuse is *checked*, not assumed: the
37//! parity fixture carries HuggingFace's own token ids for both models, so a
38//! divergence fails on the ids before it reaches the embeddings.
39//!
40//! # No architecture code
41//!
42//! `candle_transformers::models::t5` already provides `T5EncoderModel`, and
43//! `candle-transformers` was already a workspace dependency, so this module is
44//! a registry entry, a config mapping, and a tokenizer — nothing else. The
45//! weight layout read off the published checkpoint lines up exactly with what
46//! candle expects:
47//!
48//! | checkpoint key | candle |
49//! |---|---|
50//! | `shared.weight` | `T5EncoderModel::load` probes this first |
51//! | `encoder.block.{i}.layer.0.SelfAttention.{q,k,v,o}.weight` | `T5Attention` |
52//! | `encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight` | block 0 only |
53//! | `encoder.block.{i}.layer.1.DenseReluDense.{wi,wo}.weight` | `T5DenseActDense` |
54//! | `encoder.final_layer_norm.weight` | `T5Stack` |
55//!
56//! Note `q/k/v` are `[4096, 1024]`, not square: ProtT5-XL sets `d_kv: 128` with
57//! `num_heads: 32`, so the attention inner dimension is 4096 against a
58//! `d_model` of 1024. candle computes `inner_dim = num_heads * d_kv` and
59//! handles that correctly. Ankh-large is the same story in the other
60//! direction: `d_kv: 64` and `num_heads: 16` give an inner dimension of 1024
61//! against a `d_model` of 1536.
62
63pub mod runner;
64pub mod tokenizer;
65pub mod translator;
ferritin_plms/t5/mod.rs

ferritin_plms/t5/
mod.rs
