Train an MPNN model and predict with it

An MPNN (message-passing neural network, Chemprop-style) learns to predict a numeric property, such as solubility or potency, from molecular structure. 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, "validation_size": 0.1},
}
 
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), the share held out for validation (validation_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).

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, plus the unit you trained with

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.

Tips

• More data beats more epochs. A few hundred measured molecules is a sensible minimum for a useful model.
• Keep a test set out of training and compare its predictions with your measurements before trusting the model.
• Delete models you no longer need with client.inference.models.mpnn.delete(name) to stay under your plan's limit.