| |
| """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"] |
|
|
| |
| |
| 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 |
|
|
| |
| _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()} |
| |
| 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 = ("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}") |
|
|
| |
| 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() |
|
|
