#!/usr/bin/env python3 """DeepChem 2.4.0 regex SMILES tokenizer for MolGemma / dpp4-molgemma-finetuned. This is the exact tokenizer the MolGemma model was pretrained and fine-tuned with (recovered from the reference notebook ``MolMistral-Moses-KerasNLP.ipynb``). It must be used for any inference with ``Delower/MolGemma`` or ``Delower/dpp4-molgemma-finetuned`` — no other tokenizer can address the model's 591-token embedding. Vocabulary ---------- ``vocab_deepchem_2.4.0.txt`` holds one token per line; a token's id is its line index (0-based). 591 tokens total. Special tokens (fixed by the vocabulary file): =========== ===== token id =========== ===== ``[PAD]`` 0 ``[UNK]`` 11 ``[CLS]`` 12 ``[SEP]`` 13 ``[MASK]`` 14 =========== ===== Conventions (identical to training and generation) * Training sequences are the RAW token ids of the SMILES string: no BOS/EOS/[CLS]/[SEP] tokens are added. The target is the input shifted by one position (standard causal LM). * The model inputs are ``token_ids`` and ``padding_mask`` (``padding_mask = (token_ids != PAD_ID)``). * Generation is seeded with the single token ``'C'`` and sampled autoregressively; the sampled ids are detokenized and the special tokens ``[PAD]``/``[CLS]``/``[SEP]`` are stripped from the string. * A SMILES containing a character the regex splits into a token that is not in the vocabulary is out-of-vocabulary (OOV). During fine-tune data preparation such rows were dropped (0 of 9,338 rows were OOV). :func:`encode` raises on OOV by default. Usage ----- >>> tok = DeepChemTokenizer("vocab_deepchem_2.4.0.txt") >>> ids = tok.encode("CC#N") >>> tok.decode(ids) 'CC#N' Run as a script (``python tokenizer.py``) to execute the self-test: vocab size, special ids, and an alogliptin round-trip. """ import os import re __all__ = ["DeepChemTokenizer", "PATTERN", "PAD_ID", "UNK_ID", "CLS_ID", "SEP_ID"] # The DeepChem 2.4.0 SMILES regex — do not modify: the model's # embedding was trained against exactly this segmentation. PATTERN = (r"(\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#" r"|-|\+|\\|\/|:|~|@|\?|>>?|\*|\$|\%[0-9]{2}|[0-9])") _REGEX = re.compile(PATTERN) PAD_ID = 0 UNK_ID = 11 CLS_ID = 12 SEP_ID = 13 MASK_ID = 14 # Special tokens removed from detokenized generation output. _STRIP_TOKENS = ("[PAD]", "[CLS]", "[SEP]") def default_vocab_path(): """Locate vocab_deepchem_2.4.0.txt next to this file or in the cwd.""" here = os.path.dirname(os.path.abspath(__file__)) for cand in (os.path.join(here, "vocab_deepchem_2.4.0.txt"), os.path.join(os.getcwd(), "vocab_deepchem_2.4.0.txt")): if os.path.exists(cand): return cand return os.path.join(here, "vocab_deepchem_2.4.0.txt") class DeepChemTokenizer: """Regex SMILES tokenizer over the DeepChem 2.4.0 vocabulary.""" def __init__(self, vocab_path=None): self.vocab_path = vocab_path or default_vocab_path() with open(self.vocab_path) as f: self.tokens = f.read().splitlines() if len(self.tokens) != 591: raise ValueError( f"expected 591 vocab tokens, got {len(self.tokens)} " f"from {self.vocab_path}") self.stoi = {tok: i for i, tok in enumerate(self.tokens)} self.itos = {i: tok for tok, i in self.stoi.items()} # Sanity: the special ids are defined by the vocab file itself. for name, idx in (("[PAD]", PAD_ID), ("[UNK]", UNK_ID), ("[CLS]", CLS_ID), ("[SEP]", SEP_ID), ("[MASK]", MASK_ID)): if self.tokens[idx] != name: raise ValueError( f"vocab mismatch: line {idx} is {self.tokens[idx]!r}, " f"expected {name!r}") def __len__(self): return len(self.tokens) def tokenize(self, smiles): """Split a SMILES string into DeepChem regex tokens.""" return _REGEX.findall(str(smiles)) def detokenize(self, tokens): """Join tokens back into a SMILES string (exact inverse of :meth:`tokenize` for in-vocabulary token sequences).""" return "".join(tokens) def encode(self, smiles, strict=True): """Encode a SMILES string to raw token ids (no BOS/EOS added — the training convention). With ``strict=True`` (default) an out-of-vocabulary token raises ``ValueError``; with ``strict=False`` OOV tokens are skipped.""" ids = [] for tok in self.tokenize(smiles): idx = self.stoi.get(tok) if idx is None: if strict: raise ValueError( f"out-of-vocabulary token {tok!r} in {smiles!r}") continue ids.append(idx) return ids def decode(self, ids): """Decode token ids to a SMILES string. ``[PAD]``/``[CLS]``/ ``[SEP]`` are stripped, mirroring the generation pipeline.""" s = "".join(self.itos[int(i)] for i in ids) for special in _STRIP_TOKENS: s = s.replace(special, "") return s def in_vocab(self, smiles): """True if every regex token of the SMILES is in the vocab.""" return all(t in self.stoi for t in self.tokenize(smiles)) def _self_test(): tok = DeepChemTokenizer() assert len(tok) == 591, len(tok) assert tok.stoi["[PAD]"] == PAD_ID == 0 assert tok.stoi["[CLS]"] == CLS_ID == 12 assert tok.stoi["[SEP]"] == SEP_ID == 13 print(f"vocab OK: 591 tokens | [PAD]={PAD_ID} [CLS]={CLS_ID} " f"[SEP]={SEP_ID} | vocab file: {tok.vocab_path}") # Alogliptin (FDA-12 reference SMILES, as used in the project). alogliptin = ("Cn1c(=O)cc(N2CCC[C@@H](N)C2)n" "(Cc2ccccc2C#N)c1=O") tokens = tok.tokenize(alogliptin) assert tok.detokenize(tokens) == alogliptin, "token round-trip failed" ids = tok.encode(alogliptin) assert tok.decode(ids) == alogliptin, "id round-trip failed" assert tok.in_vocab(alogliptin) print(f"alogliptin round-trip OK: {len(tokens)} tokens, " f"ids {ids[:8]}... -> {tok.decode(ids)!r}") # Generation seed convention: sampling starts from token 'C'. assert tok.encode("C") == [tok.stoi["C"]] print(f"generation seed OK: 'C' -> id {tok.stoi['C']}") print("SELF-TEST PASSED") if __name__ == "__main__": _self_test()