Skip to content

README.md

Latest commit

 

History

History
241 lines (198 loc) · 9.73 KB

README.md

File metadata and controls

241 lines (198 loc) · 9.73 KB
Raw

ChemLM: De Novo Drug Design with Chemical Language Models

Python PyTorch HuggingFace RDKit License: MIT Tests Code style: black

ChemLM is a PyTorch + 🤗 Transformers framework for training and sampling SMILES-based chemical language models (CLMs) for de novo molecular generation, with support for LSTM and Transformer architectures, beam search decoding, perplexity-based molecule ranking, and bioactivity- conditioned fine-tuning.

Motivation

Generative chemical language models treat molecules as strings — most commonly SMILES — and learn to generate novel, synthesizable, drug-like structures by modeling the distribution of valid molecular strings token-by-token. This repurposes decades of sequence-modeling progress (RNNs, and now Transformers) for molecular design.

ChemLM is built around the two-stage transfer-learning strategy popularized by the ETH MODLAB group: (1) pretrain a chemical language model broadly on a large, unlabeled corpus of drug-like molecules to learn the "grammar" of chemistry, then (2) fine-tune on a smaller set of molecules labeled with a bioactivity signal (active/inactive, pIC50, etc.) so the model's generative distribution shifts toward molecules with a desired biological profile. Perplexity — the model's own confidence in a generated sequence — is then used as a lightweight, label-free proxy for ranking and filtering candidates before any downstream scoring or synthesis.

This project is a from-scratch, educational/research reference implementation of that pipeline — not an official release from ETH MODLAB — intended as a clear, hackable starting point for de novo design experiments.

Features

  • 🧬 SMILES tokenizer — regex-based, chemistry-aware tokenization (multi-character atoms, ring-closure digits, stereochemistry, charges) with a vocabulary built directly from your training corpus.
  • 🔁 Two model architectures, trained and sampled through a common interface:
    • SMILESLSTM — a stacked LSTM sequence model (REINVENT-style).
    • SMILESTransformer — a causal Transformer decoder built on HuggingFace's GPT2LMHeadModel.
  • 🎯 Beam search sampling for both architectures, plus temperature/ top-k stochastic sampling for diverse candidate generation.
  • 📊 Perplexity-based molecule ranking — rank and filter generated candidates by the model's own per-token log-likelihood.
  • 💊 Bioactivity-conditioned fine-tuning — condition generation on an <active> / <inactive> prefix token learned during fine-tuning on a labeled bioactivity dataset, following the hybrid CLM strategy of Moret et al. (2023).
  • 📈 Standard generation benchmarks — validity, uniqueness, novelty, and internal diversity (RDKit-based, with graceful fallback if RDKit is unavailable).
  • ✅ Unit-tested core components (tokenizer, models, generation) via pytest.
  • 📓 Notebooks for data exploration and an end-to-end training walkthrough.

Installation

git clone https://github.com/Islamomar-1/ChemLM.git
cd ChemLM

python -m venv venv
source venv/bin/activate  # Windows: venv\Scripts\activate

pip install -r requirements.txt
# or, for an editable/development install:
pip install -e .

Note: RDKit is used for molecule validity checks, canonicalization, and diversity metrics. If it isn't available on your platform, ChemLM degrades gracefully (validity is assumed rather than verified) so the rest of the pipeline still runs.

Quick Start

# 1. Pretrain an LSTM chemical language model on the bundled sample corpus
python train.py --data data/sample_smiles.txt --arch lstm --epochs 20

# 2. Generate molecules with beam search, ranked by perplexity
python generate.py \
    --checkpoint checkpoints/lstm_pretrained.pt \
    --vocab checkpoints/vocab.json \
    --num_samples 50 --method beam --beam_size 10 \
    --filter_valid --output generated_molecules.smi

# 3. Evaluate validity / uniqueness / novelty / diversity
python evaluate.py --generated generated_molecules.smi --train_data data/sample_smiles.txt

# 4. Fine-tune on a bioactivity-labeled dataset (SMILES,label CSV)
python train.py \
    --bioactivity data/bioactivity.csv --arch lstm --finetune \
    --init_checkpoint checkpoints/lstm_pretrained.pt

# 5. Sample molecules conditioned on predicted activity
python generate.py \
    --checkpoint checkpoints/lstm_finetuned.pt \
    --vocab checkpoints/vocab.json \
    --condition active --num_samples 20

Swap --arch lstm for --arch transformer at any step to use the Transformer variant instead.

Stack

Component Technology
Modeling PyTorch, 🤗 Transformers (GPT2LMHeadModel backbone)
Cheminformatics RDKit (validity, canonicalization, fingerprints)
Data handling Custom torch.utils.data.Dataset classes
Testing pytest
Experiment tracking (optional) Weights & Biases
Notebooks Jupyter

Repository Structure

chemlm/
├── data/
│   ├── __init__.py
│   ├── tokenizer.py         # SMILESTokenizer: regex tokenization + vocab
│   ├── dataset.py           # SMILESDataset, BioactivityDataset
│   └── sample_smiles.txt    # small demo corpus for quick start / tests
├── models/
│   ├── __init__.py
│   ├── lstm_model.py        # SMILESLSTM (stacked LSTM LM)
│   └── transformer_model.py # SMILESTransformer (GPT2-style decoder)
├── notebooks/
│   ├── 01_data_exploration.ipynb
│   └── 02_model_training.ipynb
├── tests/
│   ├── test_tokenizer.py
│   ├── test_models.py
│   └── test_generate.py
├── train.py                 # pretraining + bioactivity-conditioned fine-tuning
├── generate.py               # beam search / sampling + perplexity ranking
├── evaluate.py                # validity / uniqueness / novelty / diversity
├── requirements.txt
├── setup.py
├── LICENSE
└── .gitignore

Benchmark

Measured results (actually run, not simulated) from training each architecture end-to-end on the bundled 20-molecule demo corpus (data/sample_smiles.txt). This corpus exists purely to validate that the pipeline runs correctly on CPU in seconds — train on a real corpus (e.g. ChEMBL, GDB-13, or a curated patent set as in Moret et al. 2023) for meaningful generative quality; a 20-molecule vocabulary will always overfit and produce a middling validity rate like the one below.

Model Params Epochs Val. Perplexity
SMILESLSTM (3×512) 5.79M 30 3.19
SMILESTransformer (6L/8H, d=256) 4.81M 10 4.15

Generation metrics (60 samples, temperature sampling, top_k=8, from the LSTM above, measured by evaluate.py with RDKit installed):

Validity Uniqueness (of valid) Novelty (of unique) Internal diversity
38.3% 100.0% 91.3% 0.889

Reproduce with:

python train.py --data data/sample_smiles.txt --arch lstm --epochs 30 --batch_size 4
python generate.py --checkpoint checkpoints/lstm_pretrained.pt --vocab checkpoints/vocab.json \
    --num_samples 60 --method sample --temperature 1.0 --top_k 8 --output generated_raw.smi
python evaluate.py --generated generated_raw.smi --train_data data/sample_smiles.txt

Contributing

Contributions are welcome! To get started:

  1. Fork the repository and create a feature branch: git checkout -b feature/my-improvement
  2. Set up a development environment: pip install -e ".[dev]"
  3. Run the test suite before opening a PR: pytest tests/ -v
  4. Follow existing code style (type hints, docstrings on public functions/classes) and keep PRs focused on a single change.
  5. Open a pull request describing the motivation and summarizing your changes; link any related issues.

Bug reports and feature requests are equally welcome via GitHub Issues — please include a minimal reproduction (SMILES input, command run, and full traceback) when reporting bugs.

Citation

If you use this repository, please cite the foundational work it builds on:

@article{moret2023leveraging,
  title   = {Leveraging molecular structure and bioactivity with chemical
             language models for de novo drug design},
  author  = {Moret, Michael and Pachon Angona, Irene and Cotos, Leandro
             and Yan, Shen and Atz, Kenneth and Brunner, Cyrill
             and Baumgartner, Martin and Grisoni, Francesca
             and Schneider, Gisbert},
  journal = {Nature Communications},
  volume  = {14},
  number  = {1},
  pages   = {114},
  year    = {2023},
  doi     = {10.1038/s41467-022-35692-6},
  url     = {https://doi.org/10.1038/s41467-022-35692-6}
}

Moret, M., Pachon Angona, I., Cotos, L., Yan, S., Atz, K., Brunner, C., Baumgartner, M., Grisoni, F. & Schneider, G. (2023). Leveraging molecular structure and bioactivity with chemical language models for de novo drug design. Nature Communications, 14, 114. https://doi.org/10.1038/s41467-022-35692-6

If you use this specific codebase, a link back to this repository is appreciated.

License

Released under the MIT License.

ChemLM