DemoDiffMolecularGenerator

Graph diffusion with optional molecular demonstrations and normalized score values.

Usage

DemoDiff trains on SMILES and optionally uses other molecules as demonstrations. In fit(), contexts[i] lists indices of other training molecules for target i; in generate(), contexts contains demonstration SMILES. Supply score as normalized values in [0, 1], with higher values meaning more target-like demonstrations. The target’s score is fixed to 1 internally. Use the same score definition in training and generation. The local example uses Morgan fingerprint Tanimoto similarity: 1 means identical fingerprints and 0 means no shared fingerprint bits. During generation, scores describe the demonstrations in the desired task; they do not guarantee measured similarities of generated molecules. Omit both contexts and score for unconditional training or generation.

from torch_molecule import DemoDiffMolecularGenerator

smiles = ["CCOCCN", "c1ccccc1CCO", "CC(=O)NCCO"]
model = DemoDiffMolecularGenerator(
    num_layer=2, hidden_size=64, num_head=4,
    diffusion_steps=50, epochs=20, batch_size=3,
)
model.fit(
    smiles,
    contexts=[[1], [0], [0, 1]],
    score=[[0.8], [0.8], [0.6, 0.7]],
)
generated = model.generate(
    n_samples=10, contexts=[smiles[0]], score=[0.8],
)

The scores above are illustrative. Demonstrations guide generation; they are not guaranteed substructures. The tokenizer can lose stereochemistry and formal charges, while polymer attachment atoms are supported.

from torch_molecule import DemoDiffMolecularGenerator

Examples · Methods

Constructor

class DemoDiffMolecularGenerator(num_layer=24, hidden_size=1280, num_head=16, mlp_ratio=4.0, diffusion_steps=500, diffusion_noise_schedule='cosine', context_length=150, drop_context=0.1, drop_incontext=0.0, no_context_train=False, lambda_train=None, noise_weight=None, guidance_scale=2.0, batch_size=64, epochs=200, learning_rate=2e-05, grad_clip_value=0.1, weight_decay=1e-12, verbose='none', *, device=None, model_name='DemoDiffMolecularGenerator')

Train a fragment-graph diffusion model, optionally using demonstrations.

Parameters:

• num_layer (int, default=24) – Number of transformer layers.

• hidden_size (int, default=1280) – Transformer hidden dimension.

• num_head (int, default=16) – Attention heads; each head must have an even dimension.

• mlp_ratio (float, default=4.0) – Feed-forward expansion ratio.

• diffusion_steps (int, default=500) – Number of discrete diffusion steps.

• diffusion_noise_schedule (str, default="cosine") – Noise schedule: “cosine”, “linear”, or “custom”.

• context_length (int, default=150) – Maximum total fragment tokens in a target and its demonstrations.

• drop_context (float, default=0.1) – Probability of ignoring all demonstrations for a training batch.

• drop_incontext (float, default=0.0) – Initial drop probability for individual demonstrations, decreasing with order.

• no_context_train (bool, default=False) – Train without conditioning even when demonstrations are supplied.

• lambda_train (list of float, optional) – Token, bond, and attachment-position loss weights; defaults to [1, 1, 1].

• noise_weight (str or None, default=None) – Optional loss weighting: “noiser_lower” or “noiser_higher”.

• guidance_scale (float, default=2.0) – Demonstration guidance strength during generation.

• batch_size (int, default=64) – Number of target molecules per batch.

• epochs (int, default=200) – Training epochs.

• learning_rate (float, default=2e-5) – AdamW learning rate.

• grad_clip_value (float or None, default=0.1) – Maximum gradient norm; None disables clipping.

• weight_decay (float, default=1e-12) – AdamW weight decay.

• verbose (str, default="none") – “none”, “progress_bar”, or “print_statement”.

• device (str or torch.device, optional) – Training and generation device; selected automatically when omitted.

• model_name (str, default="DemoDiffMolecularGenerator") – Model identifier.

Notes

Uses the supplied 3,000-fragment vocabulary. The original tokenizer does not preserve all stereochemistry or formal charges. Polymer attachment atoms are supported. Demonstrations guide generation but are not fixed substructures.

Example

examples/generation/demodiff.py

"""Train a small DemoDiff model on local SMILES, with optional demonstrations.

Run from the repository root:
    python examples/generation/demodiff.py
    python examples/generation/demodiff.py --with-context
    python examples/generation/demodiff.py --data data/polymer100.csv --with-context

Demonstration score uses Morgan fingerprint Tanimoto similarity in [0, 1]:
1 for identical fingerprints, 0 for no shared fingerprint bits. This is a local
example of normalized task scores; the target's score is always 1 internally.
Supply the same score definition to fit() and generate(). Scores describe the
demonstrations, and do not guarantee measured similarities of generated molecules.
Only the smiles column is used. Generation does not enforce polymer attachment
atoms, and the tokenizer can lose stereochemistry and formal charges.
"""

import argparse
import csv
import random
from pathlib import Path

import torch
from rdkit import Chem, DataStructs
from rdkit.Chem import rdFingerprintGenerator

from torch_molecule import DemoDiffMolecularGenerator


PROJECT_ROOT = Path(__file__).resolve().parents[2]


def build_demonstrations(smiles):
    """Choose two most similar molecules; score is Tanimoto similarity in [0, 1]."""
    generator = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048)
    fingerprints = [generator.GetFingerprint(Chem.MolFromSmiles(value)) for value in smiles]
    contexts, score = [], []
    for target, fingerprint in enumerate(fingerprints):
        similarity = DataStructs.BulkTanimotoSimilarity(fingerprint, fingerprints)
        indices = sorted(
            (index for index in range(len(smiles)) if index != target),
            key=lambda index: similarity[index],
            reverse=True,
        )[:2]
        contexts.append(indices)
        score.append([similarity[index] for index in indices])
    return contexts, score


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--data", type=Path, default=PROJECT_ROOT / "data" / "molecule100.csv")
    parser.add_argument(
        "--with-context", action="store_true", help="Train with molecular demonstrations."
    )
    parser.add_argument("--epochs", type=int, default=25)
    parser.add_argument("--n-samples", type=int, default=10)
    parser.add_argument("--device", default=None, help="For example: cpu, cuda, or cuda:1.")
    args = parser.parse_args()

    random.seed(42)
    torch.manual_seed(42)
    with args.data.open(newline="") as handle:
        smiles = [row["smiles"] for row in csv.DictReader(handle)]

    mode = "conditioned" if args.with_context else "unconditional"
    output = (
        PROJECT_ROOT / "artifacts" / "demodiff" / "examples_similarity" / f"{args.data.stem}_{mode}"
    )
    output.mkdir(parents=True, exist_ok=True)
    print(f"Training on {len(smiles)} molecules from {args.data} ({mode}).")

    model = DemoDiffMolecularGenerator(
        num_layer=2,
        hidden_size=64,
        num_head=4,
        mlp_ratio=2,
        diffusion_steps=20,
        batch_size=16,
        epochs=args.epochs,
        learning_rate=1e-3,
        drop_context=0.1 if args.with_context else 0.0,
        guidance_scale=2.0 if args.with_context else 1.0,
        device=args.device,
        verbose="progress_bar",
    )
    if args.with_context:
        # fit() uses training indices, with one similarity score per demonstration.
        contexts, score = build_demonstrations(smiles)
        print(f"First target: contexts={contexts[0]}, score={score[0]}")
        model.fit(smiles, contexts=contexts, score=score)
    else:
        model.fit(smiles)

    checkpoint = str(output / "model.pt")
    model.save_to_local(checkpoint)
    loaded = DemoDiffMolecularGenerator(device=model.device)
    loaded.load_from_local(checkpoint)

    if args.with_context:
        # Reuse the first target's demonstrations and their similarity scores.
        generated = loaded.generate(
            n_samples=args.n_samples,
            contexts=[smiles[index] for index in contexts[0]],
            score=score[0],
        )
    else:
        generated = loaded.generate(n_samples=args.n_samples)

    with (output / "generated.csv").open("w", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["smiles"])
        writer.writerows([value] for value in generated)
    with (output / "loss.csv").open("w", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["epoch", "loss"])
        writer.writerows(enumerate(model.fitting_loss, start=1))

    print(f"Returned {len(generated)}/{args.n_samples} valid connected molecules.")
    print(f"Outputs: {output}")
    for value in generated[:5]:
        print(value)


if __name__ == "__main__":
    main()

Methods

fit	Train fresh weights on SMILES, optionally with related demonstrations.
generate	Generate SMILES, optionally guided by shared demonstration molecules.

For shared checkpoint and parameter operations, see Base Modules.

fit

DemoDiffMolecularGenerator.fit(X_train, y_train=None, *, contexts=None, score=None)

Train fresh weights on SMILES, optionally with related demonstrations.

contexts[i] lists indices of other molecules in X_train shown alongside target i. score[i] contains one finite similarity or normalized task score in [0, 1] per demonstration; higher means closer to the target. The target’s own score is fixed to 1 internally. For structural similarity, supply Tanimoto similarity directly. Scores are used unchanged; use the same definition and scale in generate(). With both arguments omitted, training is unconditional. y_train is unused; property labels are not converted into demonstrations automatically.

Parameters:

X_train (List[str])

generate

DemoDiffMolecularGenerator.generate(n_samples=10, *, contexts=None, score=None, num_nodes=None, guidance_scale=None, max_attempts=10)

Generate SMILES, optionally guided by shared demonstration molecules.

contexts contains demonstration SMILES. score supplies one finite similarity or normalized task score in [0, 1] per demonstration, using the same definition and scale as in fit(). Higher means closer to the target, whose query score is fixed to 1 internally. Scores describe the demonstrations and guide generation; they do not guarantee measured similarities or properties of the generated molecules.

num_nodes fixes the number of fragment tokens; otherwise their count is sampled from training data. Only valid, connected decoded molecules are returned. Exhausting max_attempts per batch warns and returns available results, which may include duplicates.

Parameters:

• n_samples (int)

• contexts (List[str] | None)

• score (List[float] | None)

• num_nodes (int | None)

• guidance_scale (float | None)

• max_attempts (int)

Return type:

List[str]