Train an MPNN model and predict with it

An MPNN (message-passing neural network, Chemprop-style) learns to predict a property from molecular structure: a number such as solubility or potency, a yes/no outcome such as active or inactive, or one of a few classes. Train one on your data with the train-mpnn job, then use it for inference like any public model.

1. Train

The training set is a list of SMILES (input_data) and one numeric value per molecule (targets), in the same order. Give the model a model_name you'll use to call it later.

from cognichem_client import CogniChem
 
client = CogniChem.from_env()
 
payload = {
    "model_name": "aq-solubility-v1",
    "input_data": ["CCO", "c1ccccc1O", "CC(=O)Oc1ccccc1C(=O)O"],   # your molecules
    "input_format": "smiles",
    "targets": [1.10, -0.04, -1.72],                               # one value per molecule
    "unit": "logS",                                                # optional, returned with predictions
    "model_description": "Aqueous solubility, in-house shake-flask data",  # optional
    "model_params": {"n_epochs": 50, "split_method": "scaffold"},
}
 
job = client.jobs.submit(job_name="train-aq-solubility-v1", job_type="train-mpnn", payload=payload)
status = client.jobs.wait(job.process_id)
print(status.status)

model_params is optional. It sets the architecture (mp_hidden_dim, mp_depth, n_ffn_layers, ffn_hidden_dim), the training (n_epochs, learning_rate, batch_size), how the data are split (split_method, validation_size, test_size), and random_seed. Defaults suit a first try; the MPNN Models tool page lists every field and its limits.

learning_rate (default 0.001) is the peak of Chemprop's learning-rate schedule: the rate rises linearly from a tenth of it to the full value over the first 5 epochs, then decays exponentially back to a tenth by the last epoch.

The saved model is the epoch with the lowest validation loss, not the last epoch, and the reported metrics come from that epoch (best_epoch).

Classification and several targets

Set task to choose what the model predicts:

payload = {
    "model_name": "herg-blocker-v1",
    "task": "binary",
    "binary_threshold": 6,          # pIC50 >= 6 counts as a blocker
    "input_data": smiles,
    "input_format": "smiles",
    "targets": pic50_values,
}

A model can also learn several properties at once. Send one row of values per molecule, with None where you have no measurement, and name the columns:

payload = {
    "model_name": "adme-panel-v1",
    "input_data": ["CCO", "c1ccccc1O", "CC(=O)Oc1ccccc1C(=O)O"],
    "input_format": "smiles",
    "targets": [[1.10, -0.31], [-0.04, None], [None, 1.19]],
    "target_names": ["logS", "logP"],
    "target_units": ["log(mol/L)", ""],
}

From a table, pass activity_table with target_columns instead (empty cells are missing values). Predictions then include predictions_by_target; predictions holds the first target, which is the one library scoring ranks by unless you choose another with score_target. A multiclass model can learn several class-label columns too, each with its own classes; it returns labels_by_target and probabilities_by_target, and library scoring ranks by the probability of a class you name with score_class. For a table of censored measurements, such as curated bioactivity with "> 10 µM" values, censored_policy drops them (drop, the default), trains on the bound as if it were exact (as_threshold), or, for a regression model with one target, keeps them with a loss that only penalises a prediction on the wrong side of the bound (bounded, read from the table's relation column).

How much to trust a prediction

Every regression model returns a 90% prediction interval (intervals) with each prediction, calibrated on held-out data. Every model also reports how close each molecule is to its training set (ad_similarity, the highest Tanimoto similarity) and whether it is inside the model's applicability domain (in_domain), using a threshold calibrated on the validation set.

For a regression model, set uncertainty_head in model_params to mve (mean-variance estimation) or evidential to have one model predict its own uncertainty for each molecule, at no extra training cost. It comes back as uncertainty, and intervals widen for molecules the model is less sure about. An ensemble (below) costs more but stays more reliable for molecules unlike the training set.

payload["model_params"] = {"n_epochs": 50, "split_method": "scaffold", "uncertainty_head": "mve"}

Set ensemble_size in model_params to train several models (up to 10) on the same split and serve their average. Their disagreement comes back as uncertainty (added to each model's own when there is an uncertainty_head), and intervals widen for molecules where the models disagree. Each model is a full training run, so training costs that many times as much; 5 is a good choice.

payload["model_params"] = {"n_epochs": 50, "split_method": "scaffold", "ensemble_size": 5}

Very large architectures are refused when you submit: a saved model, with all its ensemble members, must fit in 1 GiB.

The validation set also picks the saved epoch, so intervals calibrated on it can come out slightly narrow. Set calibration_size (for example 0.1) to set aside a separate calibration set from the training data, drawn with the same split_method; the intervals and the applicability-domain threshold are then calibrated on it. With split_method: "provided", label those rows calibration instead.

Splits and the test set

Training sets aside two shares of the data: a validation set (validation_size, default 10%) that picks the saved epoch, and a test set (test_size, default 10%; 0 skips it) that the model never sees. split_method decides who goes where:

Expect lower test metrics from scaffold and cluster than from random. They are the more honest estimate when you will predict new chemotypes.

Cross-validation and y-randomisation

Two optional checks, both off by default:

• cv_folds (3 to 10) trains one extra model per fold of the non-test molecules and reports the mean and spread of their scores. On a small dataset a single test set is a noisy estimate; the spread shows how much it depends on which molecules were held out. Scaffold and cluster groups stay in one fold when there are at least as many groups as folds; otherwise the folds are drawn at random and the model card says so. Not available with provided or time splits.
• y_randomization (1 to 10) trains that many models on shuffled training targets and reports how far a single model trained on the real targets beats them. A model that does not clearly beat them has learned chance correlations.

payload["model_params"] = {"n_epochs": 50, "split_method": "scaffold", "cv_folds": 5, "y_randomization": 3}

Neither changes the saved model. Each fold and each run is one more training run, so the cost grows accordingly.

Model card

Every model gets a model card (model_card.json). It holds:

• the test-set metrics: RMSE, MAE and R² for regression; ROC AUC, PR AUC, balanced accuracy and MCC for binary; accuracy, balanced accuracy and macro-F1 for multiclass;
• the same metrics for two baselines: the training mean (or the share of actives, or the most common class), and the average of the 5 most similar training molecules;
• for several targets, each target's metrics next to their average;
• how often the 80% and 90% intervals contained the test values (regression), or how well the predicted probabilities match observed rates (classifiers);
• the applicability-domain threshold and how the test error differs inside and outside it;
• how similar the test molecules are to the training set;
• when you asked for them, the cross-validation folds and the shuffled-target comparison;
• the split, the architecture, a data summary and library versions, each section labelled with the OECD QSAR validation principle it covers.

A model that does not beat both baselines has not learned anything useful. On the My Models page, the Evaluation panel shows these numbers and downloads the card.

Training is a job, charged to your wallet by runtime. More molecules and epochs take longer; estimate first with client.jobs.estimate. Each plan allows a set number of custom models (Plans and limits).

2. Find the model

print(client.inference.models.mpnn.user())   # your trained models

3. Predict

Use the model's name, with is_public_model false (the default):

result = client.inference.run(
    "mpnn",
    "aq-solubility-v1",
    {"input_data": ["CCN", "c1ccncc1"], "input_format": "smiles"},
)
print(result.data)   # predictions, the unit, labels for classifiers, intervals and in_domain

Preprocessing follows the model's training settings: leave add_hydrogens and remove_hydrogens out of the request (a value different from training is rejected).

Predictions are not charged to your wallet: each molecule counts toward your monthly inference allowance.

Which atoms drive a prediction

Set return_contributions to get one signed value per atom (atom_contributions, in the input's atom order). Positive atoms raise the prediction and negative atoms lower it. Each value is measured against a copy of the molecule in which every atom is the molecule's average atom, so an atom's value is what it adds by being itself. A molecule's values add up to its prediction minus attribution_base, that copy's prediction.

result = client.inference.run(
    "mpnn",
    "aq-solubility-v1",
    {"input_data": ["c1ccccc1O"], "input_format": "smiles", "return_contributions": True},
)
print(result.data["atom_contributions"][0])   # one value per atom

A classifier's values explain P(class 1), or the predicted class's probability for multiclass. A model with several targets also returns atom_contributions_by_target. Contributions take extra compute, so a request that asks for them can hold at most 100 molecules. On the MPNN Models page, Predict colours each atom by its contribution on a colour-blind-friendly Viridis scale, from the molecule's lowest value (dark purple) to its highest (yellow); a mark on the legend shows where zero falls.

Tips

• More data beats more epochs. A few hundred measured molecules is a sensible minimum for a useful model.
• Check the model card before trusting the model: the test RMSE should clearly beat both baselines. Use scaffold if you will predict new series.
• Delete models you no longer need with client.inference.models.mpnn.delete(name) to stay under your plan's limit.