CheMatE

CheMatE is a ModernBERT-based encoder for chemistry. It is pretrained with masked language modeling on SMILES-augmented scientific text (articles, educational content, reaction procedures), then contrastively fine-tuned for embeddings. The result is a single encoder that works on both molecular (SMILES) and natural-language chemistry tasks.

On property-prediction benchmarks (classification and regression) spanning both SMILES and NLP datasets, CheMatE ranks first overall and is the only model that stays competitive across both modalities. (On NLP alone it is tied with the strongest text encoders; it does not strictly dominate every per-modality split.)

Why

• SMILES-only encoders miss scientific context; text-only encoders miss molecular structure. CheMatE injects SMILES inline into scientific text so it learns both jointly.
• ModernBERT's long context (up to 8192 tokens) lets it ingest full articles and multi-step procedures with many inline SMILES.
• Frozen CheMatE embeddings feed simple downstream heads (linear/RF/GP) for property prediction and clustering.

Repository layout

src/
  encoder_training/ MLM pretraining of the ModernBERT encoder (configs, slurms, trainer)
  contrastive_training/ contrastive fine-tuning (pair dataset, losses, model, slurms)
  evaluation/
    eval_runner/ classification and regression runners
    post_processing/ CD ranking (Demsar) + Tukey HSD plots from the result CSVs
    configs/ evaluation configs (models + datasets to run)
    utils/ embedding extractors (HF / Morgan)
    slurm/ example SLURM jobs

Installation

git clone https://github.com/schwallergroup/CheMatE.git
cd CheMatE
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
pip install -e . # exposes the src/ packages

RDKit (used by the Morgan fingerprint baseline) is the only non-pure-python dep; install via pip as above or conda if you prefer.

Model & data

• Model. The eval configs point at SchwallerGroup/CheMatE-v0 (base: SchwallerGroup/CheMatE-v0-MLM). Baselines (ModernBERT, ChEmbed, ChemBERTa, SciBERT, GTE, Nomic, Morgan) are pulled from the Hub automatically.
• Evaluation datasets. Point paths.cleaned_data_dir in an eval config at the directory of cleaned parquet benchmarks (with classification/ and regression/ subfolders). The cleaned data is available on request; a public dataset release is planned.
• Pretraining / contrastive corpora. The CEMB corpora are large; the configs and slurms use /path/to/... placeholders. Set DATA_ROOT (slurms) or edit data_dirs (configs) to your copy.

Reproduce

The SLURM scripts under each slurms/ dir are CSCS Clariden examples. Set the account, container, venv and data/repo paths (or override REPO, DATA_ROOT, VENV, CONTAINER, SCRATCH from the environment) before submitting.

1. Pretrain the encoder (MLM)

cd src/encoder_training
sbatch slurms/train.slurm           # or slurms/train_balanced.slurm for the balanced sampler

2. Contrastive fine-tuning

cd src/contrastive_training
sbatch slurms/train_chemate_v0.slurm

3. Evaluation

cd src/evaluation
sbatch slurm/eval_test.slurm
sbatch slurm/eval_regression.slurm
sbatch slurm/eval_classification.slurm

Each runner also runs standalone, e.g.:

python eval_runner/eval_regression.py configs/config_regression_chemate_mean_cv20.yaml
python eval_runner/eval_regression.py configs/config_regression_chemate_mean_cv20.yaml --test

4. Plots (CD ranking + Tukey)

The runners write per-fold result CSVs and a per-run analysis (summary/ranking/pivot). Point the plot scripts at a combined per-fold CSV (columns task, source_file, dataset, model, fold, <metric>):

export PAPER_CSV=path/to/cv20_raw_folds.csv
python post_processing/cd_diagrams.py
python post_processing/tukey_plots.py

Citation

A citation will be added when the paper is public. For now, please reference this repository.

License

See LICENSE.