Workspace for pretraining models
This repo contains the tools for pretraining and finetuning chemeleon models from scratch. The user only needs a dataset of SMILES strings to start. Then, the SMILES are featurized. The featurized data is prescaled and then used to pretrain a chemprop model. There is also capability to finetune many different models with different pretraining parameters, using anvil from openadmet, which is contained in the finetuning directory. Additional analysis that was used to make design decisions is is the analysis directory, including several python notebooks under analysis/notebooks.
Additional experimental code implementing alternative message passing schemes
and an LLM-based autoresearch procedure are included in architecture.
Conda environment definitions live in conda_envs/:
conda_envs/preprocessing_environment.yaml— used for featurization and prescaling (steps 1-2).conda_envs/training_environment.yaml— used for pretraining (step 3).
Create the environment you need, for example:
conda env create -f conda_envs/preprocessing_environment.yaml
Then, with the appropriate environment activated, install the preprocessing,
pretraining, finetuning, prediction, and evaluation console commands
(from the repo root):
pip install -e .
This exposes a preprocessing command that drives both steps below from a
YAML config (see preprocessing/pipeline_config.yaml for a documented
template), a pretraining command that runs Chemprop pretraining, a
finetuning command that runs anvil finetuning jobs, a prediction command
that generates predictions from finetuned models, and an evaluation command
that scores and compares those predictions.
a) First, run download_dataset.py to download datasets from hugginface,
in this case, the MolPILE datasets.
b) If your input is a plain-text .smiles file with one SMILES per line, convert it
to parquet with python featurization/convert_smiles_to_parquet.py /path/to/input.smiles.
c) Next, configure the featurization section of your pipeline_config.yaml
(input parquet must include SMILES and id columns) and run:
preprocessing pipeline_config.yaml --only featurize
This computes the mordred descriptors, using the
conda_envs/preprocessing_environment.yaml environment.
a) Configure the prescaling section of your pipeline_config.yaml and run:
preprocessing pipeline_config.yaml --only prescale
(or omit --only to run featurization followed by prescaling in one
command, using a config with both sections). This normalizes and scales
the distributions to avoid gradient issues during training.
i) Remove nans and infs from the data, setting nans to 0 and infs to +/-
maxes for each descriptor. If any column is more than 20% infs or nans,
it is discarded.
ii) Winsorize outliers to within three standard deviations of the mean to reduce
outlier effects.
iii) Compute the correlation between feature columns and drop columns with a Pearson's
R of > 0.98 to reduce redundancy and data bias.
iv) Apply the Yeo-Johnson transform to each feature column to normalize
each column based on its underlying shape.
v) Drop any columns with zero variance after the transform to remove no-information
columns.
vi) Apply standard scaling to normalize each feature column.
a) Run the pretraining command to train chemeleon on the scaled mordred descriptors
for the dataset, using the conda_envs/training_environment.yaml environment.
Multiple GPUs are recommended for training:
pretraining /path/to/descriptors.zarr /path/to/output /path/to/smiles.parquet
b) To log pretraining runs to Weights & Biases, enable --use-wandb and optionally
set the project, for example:
pretraining /path/to/descriptors.zarr /path/to/output /path/to/smiles.parquet --use-wandb --wandb-project chemeleon-pretraining
Make sure you have authenticated first with wandb login. The provided
slurm_scripts/run_pretraining.sh batch wrapper will reuse credentials from
~/.netrc if WANDB_API_KEY is not already exported.
c) By default, pretraining runs locally in this process. Pass --slurm to
instead submit slurm_scripts/run_pretraining.sh as a single SLURM job
(not an array job) via sbatch and return immediately. --batch-size,
--num-workers, --log-every-n-steps, --use-wandb, and --wandb-project
are forwarded to the SLURM job; --devices and --num-nodes are ignored
in this case since the GPU/node allocation is fixed by the #SBATCH
directives in the script.
a) The output pytorch file (best.pt) is already in the converted format
(produced via save_save_weights.py) and can be passed directly through
to the anvil chemprop implementation. This will be in the
yaml file as a chemprop path parameter: from_foundation.
b) Write the yaml files needed to train your chosen models automatically with
anvil_tree_writer.py, altering individual parameters as needed for
individual experiments, placing each recipe in its own subdirectory of a
parent directory.
c) Run the finetuning command, pointing it at that parent directory, to run
openadmet anvil --recipe-path [yaml_file.yaml] --output-dir [output] for
every recipe found:
finetuning /path/to/parent_dir
By default, recipes are run sequentially in this process, one at a time.
Pass --slurm to submit a SLURM array job instead.
a) Once finetuning has produced anvil_training directories with trained
models, run the prediction command to generate predictions for a set of
SMILES, pointing it at the same parent directory used for finetuning and a
CSV containing a SMILES column:
prediction /path/to/parent_dir /path/to/input_smiles.csv
For each anvil_training directory under parent_dir, this runs
openadmet predict, writing predictions.csv next to each
anvil_training directory.
b) By default, predictions are run sequentially in this process, one at a
time. Pass --slurm to submit a SLURM array job instead.
c) Prediction is independent of evaluation, so it can be re-run later (e.g.
for new models or new input SMILES) without re-evaluating everything.
a) Once predictions.csv files exist (from the prediction command, run at
any time, including in previous sessions), run the evaluation command
to score them against known true values. Pass an output directory, the
true-values CSV, and one or more model directories to compare (each named
after its own directory basename, and not required to share a common
parent):
evaluation /path/to/output_dir /path/to/true_values.csv /path/to/model_A /path/to/model_B
For each anvil_training directory under each model dir (expected layout:
<model_dir>/<endpoint>/anvil_training) whose sibling predictions.csv
exists but has no regression_metrics.json yet, this computes regression
metrics and plots, written alongside that model's predictions.csv.
b) It then aggregates metrics across all given models for each endpoint into
comparison box plots, written to
output_dir/regression_metrics_boxplots/.
c) Because evaluation only reads existing predictions.csv files, it can be
run repeatedly to compare predictions generated across different
prediction runs without regenerating them.