File size: 8,899 Bytes
e11bc48 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | """
Protein motif tokenizer: greedy max-match trie over amino-acid n-gram clusters.
No external dependencies beyond the standard library.
Usage
-----
>>> from tokenization_protein_encoder import ZESTTokenizer
>>> tok = ZESTTokenizer.from_pretrained(".")
>>> ids = tok.encode("MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQTLGQHDFSAGEGLYTHMKALRPDEDRLSPLHSVYVDQWDWERVMGDGERQFSTLKSTVEAIWAGIKATEAAVSEEFGLAPFLPDQIHFVHSQELLSRYPDLDAKGRERAIAKDLGAVFLVGIGGKLSDGHRHDVRAPDYDDWSTPSELGHAGLNGDILVWNPSVSMEFQKIPIHRLATLKKMRHSSMCGQDKTAFGKELQDLQTELESMSGQGRFFLASTPYLRPQLNQLPGLKVNLNVIEQYVQKQNQWSTILTVYRQKGKLSAEPFQPTSHQLSAEKLNEGNDNLSLAAFVQLLNTSPTLAQATAVQVQNPIDKLPNLNQDSIQALQPEDLHQVLNLPKR")
>>> print(ids[:10])
"""
import json
import os
import re
import random
from typing import Callable, Dict, List, Optional, Tuple
class _TrieNode:
__slots__ = ["children", "token_id"]
def __init__(self):
self.children: Dict[str, "_TrieNode"] = {}
self.token_id: int = -1
class ZESTTokenizer:
"""
Greedy max-match tokenizer backed by a symbol-level trie.
Clusters biochemically substitutable amino-acid n-grams into shared tokens,
analogous to BPE but guided by substitutability rather than frequency.
"""
SPECIAL_TOKENS = ["<PAD>", "<UNK>", "<CLS>", "<EOS>", "<MASK>"]
DEFAULT_ALPHABET = list("ACDEFGHIKLMNPQRSTVWY")
def __init__(self, vocab_path: str, alphabet=None, alphabet_path=None,
verbose=False):
self.path = vocab_path
self.vocab: Dict[str, int] = {}
self.id_to_token: Dict[int, str] = {}
self._root = _TrieNode()
self._symbol_to_id: Dict[str, int] = {}
self._ws = re.compile(r"\s+")
for i, tok in enumerate(self.SPECIAL_TOKENS):
self.vocab[tok] = i
self.id_to_token[i] = tok
if alphabet is not None:
self._alphabet = list(alphabet)
elif alphabet_path is not None and os.path.exists(alphabet_path):
with open(alphabet_path) as f:
self._alphabet = json.load(f)
else:
auto = vocab_path.replace(".json", "_alphabet.json")
if os.path.exists(auto):
with open(auto) as f:
self._alphabet = json.load(f)
else:
self._alphabet = list(self.DEFAULT_ALPHABET)
offset = len(self.SPECIAL_TOKENS)
for i, sym in enumerate(self._alphabet):
tid = offset + i
self.vocab[sym] = tid
self.id_to_token[tid] = sym
self._symbol_to_id[sym] = tid
self._trie_insert([sym], tid)
with open(vocab_path) as f:
clusters: Dict[str, List[str]] = json.load(f)
offset = len(self.vocab)
for i, (centroid, members) in enumerate(clusters.items()):
tid = offset + i
self.id_to_token[tid] = centroid
self.vocab[centroid] = tid
self._trie_insert(self._pattern_to_symbols(centroid), tid)
for member in (members or []):
if member == centroid:
continue
self.vocab[member] = tid
self._trie_insert(self._pattern_to_symbols(member), tid)
self.pad_id = self.vocab["<PAD>"]
self.unk_id = self.vocab["<UNK>"]
self.cls_id = self.vocab["<CLS>"]
self.eos_id = self.vocab["<EOS>"]
self.mask_id = self.vocab["<MASK>"]
self.vocab_size = len(self.id_to_token)
if verbose:
n_many = sum(1 for m in clusters.values() if m and len(m) > 1)
print(f"ZESTTokenizer — vocab size: {self.vocab_size}")
print(f" Alphabet: {len(self._alphabet)} symbols")
print(f" Patterns: {len(clusters)} ({n_many} many-to-one)")
def _pattern_to_symbols(self, pattern):
if self._alphabet and all(len(s) == 1 for s in self._alphabet):
return list(pattern)
return pattern.split()
def _trie_insert(self, symbols, token_id):
node = self._root
for sym in symbols:
if sym not in node.children:
node.children[sym] = _TrieNode()
node = node.children[sym]
node.token_id = token_id
def _trie_match(self, symbols, start):
matches = []
node = self._root
for i in range(start, len(symbols)):
sym = symbols[i]
if sym not in node.children:
break
node = node.children[sym]
if node.token_id != -1:
matches.append((i - start + 1, node.token_id))
return matches[::-1]
def _segment(self, symbols, dropout=0.0):
segments = []
i, n = 0, len(symbols)
while i < n:
matches = self._trie_match(symbols, i)
if not matches:
segments.append((1, self.mask_id))
i += 1
continue
length, tid = matches[0]
if dropout > 0 and random.random() < dropout and len(matches) > 1:
idx = random.randint(1, len(matches) - 1)
length, tid = matches[idx]
segments.append((length, tid))
i += length
return segments
def encode(self, text: str, dropout: float = 0.0,
add_special_tokens: bool = False) -> List[int]:
"""Encode a raw amino-acid sequence string to token IDs."""
symbols = list(re.sub(r"\s+", "", text).upper())
ids = [tid for _, tid in self._segment(symbols, dropout)]
if add_special_tokens:
ids = [self.cls_id] + ids + [self.eos_id]
return ids
def decode(self, ids: List[int], skip_special: bool = True) -> str:
skip = {self.pad_id, self.cls_id, self.eos_id, self.mask_id}
parts = []
for i in ids:
if skip_special and i in skip:
continue
parts.append(self.id_to_token.get(i, ""))
if self._alphabet and len(self._alphabet[0]) == 1:
return "".join(parts)
return " | ".join(parts)
def batch_encode_plus(self, sequences: List[str], padding: bool = True,
max_length: Optional[int] = None, truncation: bool = True,
dropout: float = 0.0, return_tensors: str = "pt",
add_special_tokens: bool = False):
"""Batch encode and optionally pad/truncate."""
try:
import torch
except ImportError:
raise ImportError("PyTorch is required for return_tensors=\'pt\'")
batch = [self.encode(s, dropout=dropout) for s in sequences]
n_special = 2 if add_special_tokens else 0
processed = []
for ids in batch:
if max_length and truncation and len(ids) + n_special > max_length:
ids = ids[: max_length - n_special]
if add_special_tokens:
ids = [self.cls_id] + ids + [self.eos_id]
processed.append(ids)
if padding:
target = max_length or max(len(ids) for ids in processed)
padded, masks = [], []
for ids in processed:
pad_n = max(0, target - len(ids))
padded.append(ids + [self.pad_id] * pad_n)
masks.append([1] * len(ids) + [0] * pad_n)
else:
padded = processed
masks = [[1] * len(ids) for ids in processed]
if return_tensors == "pt":
return {
"input_ids": torch.tensor(padded, dtype=torch.long),
"attention_mask": torch.tensor(masks, dtype=torch.long),
}
return {"input_ids": padded, "attention_mask": masks}
def save_pretrained(self, directory: str):
os.makedirs(directory, exist_ok=True)
clusters: Dict[str, List[str]] = {}
offset = len(self.SPECIAL_TOKENS) + len(self._alphabet)
centroid_ids = {tid for tid in self.id_to_token if tid >= offset}
for tok, tid in self.vocab.items():
if tid not in centroid_ids or tok in self.id_to_token.values():
continue
centroid = self.id_to_token[tid]
clusters.setdefault(centroid, []).append(tok)
with open(os.path.join(directory, "vocab_map.json"), "w") as f:
json.dump(clusters, f, indent=2)
with open(os.path.join(directory, "vocab_map_alphabet.json"), "w") as f:
json.dump(self._alphabet, f, indent=2)
@classmethod
def from_pretrained(cls, directory: str, **kwargs) -> "ZESTTokenizer":
vocab_path = os.path.join(directory, "vocab_map.json")
alpha_path = os.path.join(directory, "vocab_map_alphabet.json")
if os.path.exists(alpha_path):
return cls(vocab_path, alphabet_path=alpha_path, **kwargs)
return cls(vocab_path, **kwargs)
|