BigSMILES and BigSmirk

Open In Colab Binder

BigSmirk tokenizes the BigSMILES encoding for macromolecules all the way down to their constituent elements.

Let’s see it in action!

🐍 Installation is easy with pre-build binaries on PyPI and GitHub. Just run: pip install smirk

Installing from source? See the developer guide for instructions.

!python -m pip install smirk transformers

First steps

🤗 smirk subclasses Hugging Face’s PreTrainedTokenizerBase for seamless compatibility and leverages Tokenizers for raw rust-powered speed. No need to learn another framework; everything works out of the box 🎁

from smirk import SmirkBigSmilesFast

# Just import and tokenize!
bigsmirk = SmirkBigSmilesFast()
bigsmirk("{[][$]CC[$],[$]CC(CC)[$][]}") # ethylene butene copolymer

{'input_ids': [159, 148, 150, 148, 2, 150, 45, 45, 148, 2, 150, 161, 148, 2, 150, 45, 45, 4, 45, 45, 5, 148, 2, 150, 148, 150, 160], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}

# Batch Tokenization with Padding
batch = bigsmirk([
    "[H]O{[>][<]C(=O)CCCCC(=O)[<],[>]NCCCCCCN[>][<]}[H]", # nylon-6,6
    "{[][<]OCC[>][<]}{[>][<]OC(C)C[>][]}", # block copolymer
    "{[][<]C(=O)c1ccc(cc1)C(=O)[<],[>]OCCO[>][]}", # alternation co-polymer
], padding="longest")
batch

{'input_ids': [[148, 71, 150, 102, 159, 148, 164, 150, 148, 163, 150, 45, 4, 22, 102, 5, 45, 45, 45, 45, 45, 4, 22, 102, 5, 148, 163, 150, 161, 148, 164, 150, 93, 45, 45, 45, 45, 45, 45, 93, 148, 164, 150, 148, 163, 150, 160, 148, 71, 150], [159, 148, 150, 148, 163, 150, 102, 45, 45, 148, 164, 150, 148, 163, 150, 160, 159, 148, 164, 150, 148, 163, 150, 102, 45, 4, 45, 5, 45, 148, 164, 150, 148, 150, 160, 168, 168, 168, 168, 168, 168, 168, 168, 168, 168, 168, 168, 168, 168, 168], [159, 148, 150, 148, 163, 150, 45, 4, 22, 102, 5, 153, 12, 153, 153, 153, 4, 153, 153, 12, 5, 45, 4, 22, 102, 5, 148, 163, 150, 161, 148, 164, 150, 102, 45, 45, 102, 148, 164, 150, 148, 150, 160, 168, 168, 168, 168, 168, 168, 168]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0]]}

# Back to polymers!
bigsmirk.batch_decode(batch["input_ids"], skip_special_tokens=True)

['[H]O{[>][<]C(=O)CCCCC(=O)[<],[>]NCCCCCCN[>][<]}[H]',
 '{[][<]OCC[>][<]}{[>][<]OC(C)C[>][]}',
 '{[][<]C(=O)c1ccc(cc1)C(=O)[<],[>]OCCO[>][]}']

Let’s visualize BigSMILES token boundaries for PVC (Polyvinyl chloride ) and sPP (Syndiotactic Polypropylene) by coloring each token in sequence.

Hide code cell source

 
import hashlib
from html import escape
from IPython.display import HTML


def render_colored_tokens(text, tokenizer=bigsmirk):
    tokens = tokenizer.tokenize(text)
    palette = [
        "#f94144", "#f3722c", "#f8961e", "#f9844a", "#f9c74f",
        "#90be6d", "#43aa8b", "#4d908e", "#577590", "#277da1",
    ]
    spans = []
    for tok in tokens:
        digest = hashlib.sha1(tok.encode("utf-8")).digest()
        color = palette[int.from_bytes(digest[:2], "big") % len(palette)]
        label = escape(tok)
        spans.append(
            f"<span style='display:inline-block;margin:2px;padding:3px 6px;border-radius:6px;background:{color};color:#111;font-family:monospace'>{label}</span>"
        )
    return HTML("<div style='line-height:2.2'>" + "".join(spans) + "</div>")

# Polyvinyl chloride (PVC)
render_colored_tokens("{[][$]CC(Cl)[$][]}")
{[][$]CC(Cl)[$][]}
# Syndiotactic polypropylene (sPP)
render_colored_tokens("CC{[>][<]C[C@@H](C)C[C@H](C)[>];[<]C=CC,[<]C[C@H](C)C=CC[]}")
CC{[>][<]C[C@@H](C)C[C@H](C)[>];[<]C=CC,[<]C[C@H](C)C=CC[]}

Handling the Fragment Name Definition Notation

The BigSMILES line notation allows some portions of the BigSMILES representations be replaced by more abstract but compact proxy fragment names, for example, the names of repeating units. The smirk BigSMILES tokenizers handles this '[' + '#' + fragment_name + ']' syntax by replacing the fragment name with its definition and then tokenizing the expanded BigSMILES. The definition should be provided within the BigSMILES as specified by the line notation i.e.: BigSMILES_string + '[' + '#' + fragment_name + ']' + BigSMILES_string + '.' + '{' + '#' + fragment_name + '=' + BigSMILES_string + '}'

For example:

C([#Arm1])([#Arm2]).{#Arm1=CO}.{#Arm2=N} --> C(CO)(N) ---> 'C', '(', 'C', 'O', ')', '(', 'N', ')'

More examples of valid and invalid use of the fragment name definition notation are provided below.

abstract_label_examples = {
    "defined [#label] placeholder": "C([#Arm])([#Arm])([#Arm])[#Arm].{#Arm=CO{[<][>]CCO[<][>]}}",
    "multiple [#label] placeholders": "C([#Arm1])([#Arm2]).{#Arm1=CO}.{#Arm2=N}",
    "[#label] in stochastic object": "{[][$]CC(C)([#Side])[$][]}.{#Side=C(=O)OCC}",
    "undefined [#label] placeholder": "C([#Arm])([#Arm])",
    "bare labels stay unknown": r"A([$1[<1]1])R(A'[$1[>1]1]).{#A=C}.{#R=C}",
}

for label, text in abstract_label_examples.items():
    encoded = bigsmirk(text, add_special_tokens=False)
    tokens = bigsmirk.convert_ids_to_tokens(encoded["input_ids"])
    decoded = bigsmirk.decode(encoded["input_ids"], skip_special_tokens=True)

    print(label)
    print("input:", text)
    print("tokens:", tokens)
    print("unknowns:", tokens.count(bigsmirk.unk_token))
    print("decoded:", decoded)
    print()

defined [#label] placeholder
input: C([#Arm])([#Arm])([#Arm])[#Arm].{#Arm=CO{[<][>]CCO[<][>]}}
tokens: ['C', '(', 'C', 'O', '{', '[', '<', ']', '[', '>', ']', 'C', 'C', 'O', '[', '<', ']', '[', '>', ']', '}', ')', '(', 'C', 'O', '{', '[', '<', ']', '[', '>', ']', 'C', 'C', 'O', '[', '<', ']', '[', '>', ']', '}', ')', '(', 'C', 'O', '{', '[', '<', ']', '[', '>', ']', 'C', 'C', 'O', '[', '<', ']', '[', '>', ']', '}', ')', 'C', 'O', '{', '[', '<', ']', '[', '>', ']', 'C', 'C', 'O', '[', '<', ']', '[', '>', ']', '}']
unknowns: 0
decoded: C(CO{[<][>]CCO[<][>]})(CO{[<][>]CCO[<][>]})(CO{[<][>]CCO[<][>]})CO{[<][>]CCO[<][>]}

multiple [#label] placeholders
input: C([#Arm1])([#Arm2]).{#Arm1=CO}.{#Arm2=N}
tokens: ['C', '(', 'C', 'O', ')', '(', 'N', ')']
unknowns: 0
decoded: C(CO)(N)

[#label] in stochastic object
input: {[][$]CC(C)([#Side])[$][]}.{#Side=C(=O)OCC}
tokens: ['{', '[', ']', '[', '$', ']', 'C', 'C', '(', 'C', ')', '(', 'C', '(', '=', 'O', ')', 'O', 'C', 'C', ')', '[', '$', ']', '[', ']', '}']
unknowns: 0
decoded: {[][$]CC(C)(C(=O)OCC)[$][]}

undefined [#label] placeholder
input: C([#Arm])([#Arm])
tokens: ['C', '(', '[', '[UNK]', ']', ')', '(', '[', '[UNK]', ']', ')']
unknowns: 2
decoded: C([])([])

bare labels stay unknown
input: A([$1[<1]1])R(A'[$1[>1]1]).{#A=C}.{#R=C}
tokens: ['[UNK]', '(', '[', '$', '1', '[', '<', '1', ']', '1', ']', ')', '[UNK]', '(', '[UNK]', '[', '$', '1', '[', '>', '1', ']', '1', ']', ')']
unknowns: 3
decoded: ([$1[<1]1])([$1[>1]1])

Zero to Polymer Foundation Model with Smirk!

Let’s train a small RoBERTa model on polymers from S. Choi et al., 2024 using Hugging Face and smirk.

!python -m pip install accelerate datasets torch

Dataset Preprocessing

This tutorial uses the BigSMILES dataset generated by Choi et al and published on Figshare. The hidden setup cell below downloads and extracts the required CSV automatically.

Hide code cell source

 
from pathlib import Path
from urllib.request import urlretrieve
from zipfile import ZipFile

from datasets import load_dataset
import truststore

truststore.inject_into_ssl()
data_dir = Path("with_Tg")
data_file = data_dir / "JCIM_sup_bigsmiles.csv"
if not data_file.exists():
    archive, _ = urlretrieve(
        "https://api.figshare.com/v2/file/download/42507037",
        "with_Tg.zip",
    )
    data_dir.mkdir(exist_ok=True)
    with ZipFile(archive) as zip_file:
        zip_file.extract("JCIM_sup_bigsmiles.csv", data_dir)

dataset = (
    load_dataset("csv", data_files=str(data_file))["train"]
    .select_columns(["BigSMILES"])
    .train_test_split(test_size=0.2, seed=42)
)
dataset = dataset.map(bigsmirk, input_columns=["BigSMILES"], desc="Tokenizing")

💡 Hugging Face/ Tokenizers may raise a warning about being forked as we’ve already used our tokenizers (this isn’t a smirk issue). It’s harmless, but when actually training it’s best to avoid tokenization until after the fork to benefit from the rust-level parallelism

🎉 That’s it! We’ve tokenized all of the BigSMILES dataset using smirk!

dataset["train"].to_pandas().head()
BigSMILES input_ids attention_mask
0 {<c1ccc(n1CC)>} [159, 163, 153, 12, 153, 153, 153, 4, 154, 12,... [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
1 {<NC(CO)C(=O)>} [159, 163, 93, 45, 4, 45, 102, 5, 45, 4, 22, 1... [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
2 {<c1cccc(c1)N2C(=O)c3ccc(cc3C2=O)Oc4ccc5c(c4)C... [159, 163, 153, 12, 153, 153, 153, 153, 4, 153... [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...
3 {<CCCCCCCCCCOC(=O)C(=O)O>} [159, 163, 45, 45, 45, 45, 45, 45, 45, 45, 45,... [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...
4 {<CCOCCOCCOC(=O)c1ccc(cc1)C(=O)O>} [159, 163, 45, 45, 102, 45, 45, 102, 45, 45, 1... [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...

Training

Once we’ve tokenized the dataset, training the model is just a matter of configuration.

from transformers import Trainer, RobertaForMaskedLM, RobertaConfig, DataCollatorForLanguageModeling

# A very small model for demonstrating training a molecular foundation model with smirk 
config = RobertaConfig(
    vocab_size=len(bigsmirk),
    hidden_size=256,
    intermediate_size=1024,
    num_hidden_layers=4,
    num_attention_heads=4,
)
model = RobertaForMaskedLM(config)

# Setup up the trainer to use our dataset
trainer = Trainer(
    model=model,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    processing_class=bigsmirk,
    data_collator=DataCollatorForLanguageModeling(bigsmirk), # The data collator needs to know about our tokenizer
)

trainer.train()

MLM Example: Predict a Masked Token

Mask one token in a BigSMILES string and ask the trained model for top predictions.

import torch

inference_model = trainer.model
inference_model.eval()
device = next(inference_model.parameters()).device

sample = dataset["test"][5]["BigSMILES"]

# Encode and choose a position to mask
encoded = bigsmirk(sample, add_special_tokens=False)
input_ids = encoded["input_ids"]
tokens = bigsmirk.convert_ids_to_tokens(input_ids)
mask_pos = len(tokens) // 2

masked_ids = input_ids.copy()
masked_ids[mask_pos] = bigsmirk.mask_token_id
masked_tokens = bigsmirk.convert_ids_to_tokens(masked_ids)

inputs = {
    "input_ids": torch.tensor([masked_ids],device=device),
    "attention_mask": torch.ones((1, len(masked_ids)),  device=device),
}

with torch.no_grad():
    logits = inference_model(**inputs).logits[0, mask_pos].detach().cpu()
    probs = torch.softmax(logits, dim=-1)

top_k = 5
top_ids = torch.topk(probs, k=top_k).indices.tolist()
top_tokens = bigsmirk.convert_ids_to_tokens(top_ids)

print("Original:", sample)
print("Masked :", "".join(masked_tokens))
print(f"Masked token index: {mask_pos} (original token: {tokens[mask_pos]})")
print("\nTop predictions:")

for rank, (tok_id, tok) in enumerate(zip(top_ids, top_tokens), start=1):
    candidate_ids = masked_ids.copy()
    candidate_ids[mask_pos] = tok_id
    candidate = bigsmirk.decode(candidate_ids, skip_special_tokens=True)
    score = probs[tok_id].item()
    print(f"{rank}. token={tok!r:>4}  p={score:.4f}  ->  {candidate}")

Original: {<CCc1ccc(c2cccnc12)>}
Masked : {<CCc1ccc(c[MASK]cccnc12)>}
Masked token index: 11 (original token: 2)

Top predictions:
1. token='Zr'  p=0.0127  ->  {<CCc1ccc(cZrcccnc12)>}
2. token='Yb'  p=0.0108  ->  {<CCc1ccc(cYbcccnc12)>}
3. token= 'B'  p=0.0103  ->  {<CCc1ccc(cBcccnc12)>}
4. token='Sm'  p=0.0103  ->  {<CCc1ccc(cSmcccnc12)>}
5. token='Es'  p=0.0100  ->  {<CCc1ccc(cEscccnc12)>}