""" 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 = ["", "", "", "", ""] 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[""] self.unk_id = self.vocab[""] self.cls_id = self.vocab[""] self.eos_id = self.vocab[""] self.mask_id = self.vocab[""] 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)