Skip to main content

Module tokenizer

Module tokenizer 

Source
Expand description

The T5 family’s tokenizer, as a built-in table.

Shared by ProtT5, Ankh and ProstT5. The first two arrive in different containers — ProtT5 in a SentencePiece spiece.model, Ankh in a tokenizer.json Unigram model — but the alphabet underneath is the same letters at the same ids, so one table serves both. Ankh’s pieces carry no word-boundary marker, which is invisible here because this maps characters directly.

ProstT5 extends the same vocabulary rather than replacing it: Foldseek’s 3Di structural alphabet occupies ids 128..=147 as lowercase letters, each exactly THREE_DI_OFFSET above its uppercase residue. Case is therefore meaning, not formatting — M is methionine and m is a structural state — which is why vocab_id never folds it.

That equivalence is verified rather than assumed: both parity fixtures carry HuggingFace’s own token ids (see scripts/generate_t5_fixtures.py), and the Rust tests compare against them before comparing embeddings, so a divergence fails on the ids rather than showing up as mysteriously wrong numbers.

§Why this is hand-written rather than loaded

Rostlab/prot_t5_xl_half_uniref50-enc ships no tokenizer.json. It ships spiece.model, a SentencePiece protobuf — which would mean either a SentencePiece dependency or a conversion step, for a vocabulary that turns out to contain 28 pieces:

 0 <pad>   1 </s>    2 <unk>
 3 ▁A   4 ▁L   5 ▁G   6 ▁V   7 ▁S   8 ▁R   9 ▁E  10 ▁D  11 ▁T  12 ▁I
13 ▁P  14 ▁K  15 ▁F  16 ▁Q  17 ▁N  18 ▁Y  19 ▁M  20 ▁H  21 ▁W  22 ▁C
23 ▁X  24 ▁B  25 ▁O  26 ▁U  27 ▁Z

(Read directly out of the published spiece.model, not inferred.) Every amino-acid piece carries SentencePiece’s ▁ word-boundary marker, and nothing else in the file is reachable from a protein sequence, so the whole tokenizer is a 28-entry character lookup.

config.json says vocab_size: 128 and shared.weight really is [128, 1024]: ids 28..=127 are T5’s 100 <extra_id_*> sentinels, which a protein sequence never produces. The embedding table is larger than the reachable alphabet, which is normal for a T5 checkpoint and not a sign that this table is missing something.

§Whitespace

Because every piece is ▁<letter>, the reference pipeline has to space the residues out — HuggingFace’s T5Tokenizer would read "MKT" as one word and emit <unk>. Rostlab’s own model card does exactly that:

sequence = " ".join(list(re.sub(r"[UZOB]", "X", sequence)))

Mapping characters directly makes the spacing irrelevant, so this tokenizer skips whitespace and "MKT" and "M K T" produce identical ids. That removes the failure the bead was most worried about: getting the convention wrong shifts every embedding by one residue and produces plausible, meaningless numbers rather than an error.

The [UZOB] -> X substitution in that snippet is not applied here. Those four residues have their own ids (24..=27) and the model has embeddings for them; collapsing them is a convention of the model card, not of the tokenizer, and doing it silently would change a caller’s input behind their back. Callers reproducing published ProtT5 embeddings should apply it themselves — see replace_rare_residues.

Enums§

Direction
Which way a ProstT5 translation runs.

Constants§

AA_TO_FOLD_ID
<AA2fold> — prefixed to an amino-acid sequence to translate it to 3Di.
EOS_ID
</s> — id 1. Appended after the last residue; ProtT5 has no BOS.
FOLD_TO_AA_ID
<fold2AA> — prefixed to a 3Di sequence to translate it to amino acids.
PAD_ID
<pad> — id 0. Also T5’s decoder_start_token_id.
THREE_DI_OFFSET
How far a 3Di token sits above its amino-acid counterpart.
UNK_ID
<unk> — id 2. Any character outside the table below.
VOCAB_SIZE
Full vocabulary size declared by config.json and matched by shared.weight. Larger than the reachable alphabet; see the module docs.

Functions§

encode
Token ids for sequence, with the trailing </s> and no BOS.
encode_for_translation
Token ids for a ProstT5 translation input: the direction token, the sequence, then </s>.
replace_rare_residues
Rostlab’s [UZOB] -> X convention, for callers reproducing published ProtT5 embeddings.
residue_char
The amino-acid character for a token id, or None.
residue_count
How many residues sequence encodes, ignoring whitespace.
residue_id
The token id for one residue character, or UNK_ID if it has none.
three_di_char
The 3Di state character for a token id, or None if the id is not one.
three_di_id
The token id for one 3Di state character (lowercase), or UNK_ID.
vocab_id
The token id for any character in the shared vocabulary.