| """ |
| 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) |
|
|
