Chemical Language Models for Drug Discovery

The transformer revolution in NLP — starting with BERT and GPT — showed that models pre-trained on vast amounts of text with a simple self-supervised objective learn representations that transfer powerfully to downstream tasks with very little labelled data. The same idea applies to molecules. SMILES strings are a linear encoding of molecular structure, and a transformer pre-trained to reconstruct masked SMILES tokens learns to capture chemical semantics: atoms, bonds, ring systems, functional groups, and even physicochemical trends.

Chemical language models (ChemLMs) are particularly valuable in drug discovery because labelled data is scarce and expensive. A model pre-trained on 77 million ChEMBL or ZINC SMILES can be fine-tuned on 500 labelled activity measurements and outperform Morgan fingerprint models trained on the same 500 points.

Why SMILES as a Language?

SMILES (Simplified Molecular-Input Line-Entry System) encodes a molecular graph as a string of characters. From a language modelling perspective:

• Each character or sub-token (atom symbols, ring closures, branch brackets) is a "word"
• The SMILES grammar defines what sequences are chemically valid (syntax)
• Chemical properties impose meaning on valid sequences (semantics)

A transformer can learn both: it learns which token sequences are grammatically valid (SMILES syntax) through masked token prediction, and it implicitly captures chemical semantics through the co-occurrence patterns of functional groups in the training corpus.

Key Pre-trained Models

Tokenization Strategies

SMILES tokenization is non-trivial. Naive character-level tokenization treats "Br" as two tokens "B"+"r", which is chemically wrong. The main strategies are:

• Atom-level tokenization: regex splits SMILES at atom boundaries. "Br", "Cl", "Si" are preserved as single tokens. Standard for most chemical language models.
• BPE (Byte Pair Encoding): subword tokenization learned from the training corpus frequency. Groups common recurring sub-structures into single tokens. Used by ChemBERTa.
• Scaffold tokenization: tokens the scaffold core first, then substituents separately. Biases the model toward scaffold-aware representations.

import re

def atom_wise_tokenize(smiles: str) -> list[str]:
    """
    Tokenize a SMILES string at atom boundaries.
    Handles multi-character atoms (Br, Cl, Si, etc.) and
    bracket atoms ([NH], [C@@H], etc.) as single tokens.
    """
    PATTERN = r"(\[[^\]]+\]|Br|Cl|Si|Se|Na|Ca|Mg|[A-Z][a-z]?|%\d{2}|[bcnops]|[=#\-\+\\\/\.\(\)@%\d])"
    tokens = re.findall(PATTERN, smiles)
    return tokens


# Example
smi = "CC(=O)Oc1ccccc1C(=O)O"
tokens = atom_wise_tokenize(smi)
print("SMILES:", smi)
print("Tokens:", tokens)
# Output: ['C', 'C', '(', '=', 'O', ')', 'O', 'c', '1', 'c', 'c', ...]

Loading ChemBERTa with HuggingFace

from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np

# Load pre-trained ChemBERTa
MODEL_NAME = "seyonec/ChemBERTa-zinc-base-v1"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)
model.eval()

def smiles_to_embedding(smiles: str, pooling: str = "cls") -> np.ndarray:
    """
    Encode a SMILES string into a fixed-size embedding vector using ChemBERTa.

    Args:
        smiles:  Input SMILES string
        pooling: 'cls' uses the [CLS] token embedding;
                 'mean' uses mean over all token embeddings

    Returns:
        numpy array of shape (hidden_size,)  — 384 for ChemBERTa-zinc-base-v1
    """
    inputs = tokenizer(smiles, return_tensors="pt", truncation=True, max_length=128)

    with torch.no_grad():
        outputs = model(**inputs)

    hidden_states = outputs.last_hidden_state  # shape: (1, seq_len, hidden_size)

    if pooling == "cls":
        embedding = hidden_states[:, 0, :]   # [CLS] token
    else:
        # Mean pool over non-padding tokens
        attention_mask = inputs["attention_mask"]  # shape: (1, seq_len)
        mask_expanded = attention_mask.unsqueeze(-1).float()
        embedding = (hidden_states * mask_expanded).sum(dim=1) / mask_expanded.sum(dim=1)

    return embedding.squeeze().numpy()


# Batch encoding
def encode_smiles_batch(smiles_list: list[str], batch_size: int = 64) -> np.ndarray:
    """Encode a list of SMILES into a 2D embedding matrix."""
    all_embeddings = []
    for i in range(0, len(smiles_list), batch_size):
        batch = smiles_list[i:i + batch_size]
        inputs = tokenizer(
            batch,
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=128
        )
        with torch.no_grad():
            outputs = model(**inputs)
        # CLS token embeddings
        embeddings = outputs.last_hidden_state[:, 0, :].numpy()
        all_embeddings.append(embeddings)
    return np.vstack(all_embeddings)

Fine-Tuning for Property Prediction

Fine-tuning a pre-trained ChemBERTa on a labelled property dataset follows the same workflow as fine-tuning BERT on text classification, but uses SMILES as input instead of sentences.

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel, get_linear_schedule_with_warmup
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np
from sklearn.metrics import roc_auc_score

class MoleculeDataset(Dataset):
    """PyTorch dataset for SMILES + property prediction."""

    def __init__(self, smiles: list, labels: list, tokenizer, max_length: int = 128):
        self.smiles = smiles
        self.labels = torch.tensor(labels, dtype=torch.float32)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.smiles)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.smiles[idx],
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        return {
            "input_ids": encoding["input_ids"].squeeze(0),
            "attention_mask": encoding["attention_mask"].squeeze(0),
            "label": self.labels[idx],
        }


class ChemBERTaClassifier(nn.Module):
    """Binary classifier head on top of ChemBERTa."""

    def __init__(self, model_name: str, dropout: float = 0.1):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(model_name)
        hidden_size = self.encoder.config.hidden_size
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden_size, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, 1),
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls_embedding = outputs.last_hidden_state[:, 0, :]  # [CLS] token
        logit = self.classifier(cls_embedding).squeeze(-1)
        return logit


def train_one_epoch(model, loader, optimizer, scheduler, device):
    model.train()
    total_loss = 0
    criterion = nn.BCEWithLogitsLoss()
    for batch in loader:
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["label"].to(device)

        optimizer.zero_grad()
        logits = model(input_ids, attention_mask)
        loss = criterion(logits, labels)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        total_loss += loss.item()
    return total_loss / len(loader)


# --- Training Script ---
MODEL_NAME = "seyonec/ChemBERTa-zinc-base-v1"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCHS = 10
LR = 2e-5
BATCH_SIZE = 32

df = pd.read_csv("activity_dataset.csv")   # columns: smiles, label (0/1)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Scaffold split (use scaffold_split from scaffold-analysis guide)
from sklearn.model_selection import train_test_split
train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)

train_ds = MoleculeDataset(train_df["smiles"].tolist(), train_df["label"].tolist(), tokenizer)
test_ds  = MoleculeDataset(test_df["smiles"].tolist(), test_df["label"].tolist(), tokenizer)
train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
test_loader  = DataLoader(test_ds, batch_size=BATCH_SIZE)

model = ChemBERTaClassifier(MODEL_NAME).to(DEVICE)
optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01)
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    num_training_steps=len(train_loader) * EPOCHS,
)

for epoch in range(1, EPOCHS + 1):
    train_loss = train_one_epoch(model, train_loader, optimizer, scheduler, DEVICE)
    print(f"Epoch {epoch}: train loss = {train_loss:.4f}")

# Evaluate
model.eval()
all_preds, all_labels = [], []
with torch.no_grad():
    for batch in test_loader:
        logits = model(batch["input_ids"].to(DEVICE), batch["attention_mask"].to(DEVICE))
        all_preds.extend(torch.sigmoid(logits).cpu().numpy())
        all_labels.extend(batch["label"].numpy())

auc = roc_auc_score(all_labels, all_preds)
print(f"Test ROC-AUC: {auc:.4f}")

Comparing ChemBERTa vs. Fingerprint Baselines

ChemLMs are most advantageous when: (1) training sets are very small (<1000 compounds), (2) scaffold diversity is high, and (3) the property is governed by global molecular features rather than local substructures. For larger datasets (>10,000 compounds), GNNs (AttentiveFP, MPNN) often match or exceed ChemLM performance.

Using ChemBERTa Embeddings as Frozen Features

Even without fine-tuning, ChemBERTa CLS embeddings are powerful molecular features. They can be used as drop-in replacements for Morgan fingerprints in any downstream model:

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import StratifiedKFold

# Load pre-computed ChemBERTa embeddings (save with encode_smiles_batch)
df = pd.read_csv("dataset.csv")   # smiles, label
X = encode_smiles_batch(df["smiles"].tolist())   # shape: (n, 384)
y = df["label"].values

# 5-fold CV with frozen ChemBERTa + RF
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
auc_scores = []
for train_idx, test_idx in cv.split(X, y):
    clf = RandomForestClassifier(n_estimators=300, n_jobs=-1, random_state=42)
    clf.fit(X[train_idx], y[train_idx])
    preds = clf.predict_proba(X[test_idx])[:, 1]
    auc_scores.append(roc_auc_score(y[test_idx], preds))

print(f"ChemBERTa + RF: AUC = {np.mean(auc_scores):.3f} ± {np.std(auc_scores):.3f}")

Multi-Task Fine-Tuning

Multi-task fine-tuning is particularly powerful for chemical language models when you have several related tasks (e.g., ADMET endpoints) with overlapping training compounds. A shared ChemBERTa encoder with task-specific heads can leverage data across all tasks simultaneously:

class MultiTaskChemBERTa(nn.Module):
    """
    Shared ChemBERTa encoder with multiple prediction heads.
    Each head can be regression or classification.
    """

    def __init__(self, model_name: str, task_names: list[str], dropout: float = 0.1):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(model_name)
        hidden_size = self.encoder.config.hidden_size
        self.dropout = nn.Dropout(dropout)
        self.heads = nn.ModuleDict({
            task: nn.Linear(hidden_size, 1)
            for task in task_names
        })

    def forward(self, input_ids, attention_mask, task: str = None):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls_emb = self.dropout(outputs.last_hidden_state[:, 0, :])
        if task is not None:
            return self.heads[task](cls_emb).squeeze(-1)
        return {task: self.heads[task](cls_emb).squeeze(-1) for task in self.heads}


# Example: ADMET multitask model
tasks = ["solubility", "permeability", "metabolic_stability", "herg_inhibition"]
model = MultiTaskChemBERTa("seyonec/ChemBERTa-zinc-base-v1", tasks)

Key Takeaways

• Chemical language models (ChemBERTa, MolBERT) pre-trained with masked SMILES objectives learn transferable molecular representations from millions of unlabelled structures.
• For low-data regimes (<500 labelled samples), fine-tuned ChemBERTa outperforms Morgan fingerprint baselines by 0.02–0.05 AUC on MoleculeNet benchmarks.
• ChemBERTa CLS embeddings used as frozen features with a Random Forest classifier are a competitive, fast baseline that requires no fine-tuning GPU compute.
• Multi-task fine-tuning on related ADMET endpoints significantly improves individual task performance through shared representation learning.
• For large datasets (>10,000 compounds), GNNs like AttentiveFP have comparable or better performance than ChemLMs, and are often faster to train and serve.

Try it live

Explore a related model hands-on: open Molecule Screening Engine — Similarity Search Live and change its parameters while it runs. Everything stays in your browser.