models/score_wrapper.py

from argparse import Namespace
from functools import partial
from pathlib import Path

import torch
import yaml
from torch_geometric.nn.data_parallel import DataParallel

from onescience.utils.diffdock.diffusion_utils import get_timestep_embedding, t_to_sigma as t_to_sigma_compl
from onescience.utils.diffdock.utils import ExponentialMovingAverage

from .aa_model import AAModel
from .cg_model import CGModel
from .old_aa_model import AAOldModel


_LM_EMBEDDING_KEYS = (
    "moad_esm_embeddings_path",
    "pdbbind_esm_embeddings_path",
    "pdbsidechain_esm_embeddings_path",
    "esm_embeddings_path",
    "esm_embeddings_model",
)


def load_model_args(model_dir):
    model_dir = Path(model_dir)
    config_path = model_dir / "model_parameters.yml"
    with config_path.open("r", encoding="utf-8") as handle:
        config = yaml.full_load(handle) or {}
    return Namespace(**config)


def model_uses_lm_embeddings(model_args):
    return any(getattr(model_args, key, None) is not None for key in _LM_EMBEDDING_KEYS)


def _has_arg(args, name):
    try:
        return name in args
    except TypeError:
        return hasattr(args, name)


def _get_arg(args, name, default=None):
    if isinstance(args, dict):
        return args.get(name, default)
    if _has_arg(args, name):
        return getattr(args, name)
    return default


def get_model(args, device, t_to_sigma, no_parallel=False, confidence_mode=False, old=False):
    timestep_emb_func = get_timestep_embedding(
        embedding_type=_get_arg(args, "embedding_type", "sinusoidal"),
        embedding_dim=args.sigma_embed_dim,
        embedding_scale=_get_arg(args, "embedding_scale", 10000),
    )

    all_atoms = _get_arg(args, "all_atoms", False)
    if old and not all_atoms:
        raise NotImplementedError(
            "The old coarse-grained DiffDock model path is not migrated yet. "
            "Use old=True only with all_atoms=True, or migrate old_cg_model.py first."
        )

    lm_embedding_type = None
    if (
        _get_arg(args, "moad_esm_embeddings_path") is not None
        or _get_arg(args, "pdbbind_esm_embeddings_path") is not None
        or _get_arg(args, "pdbsidechain_esm_embeddings_path") is not None
        or _get_arg(args, "esm_embeddings_path") is not None
    ):
        lm_embedding_type = "precomputed"
    if _get_arg(args, "esm_embeddings_model") is not None:
        lm_embedding_type = args.esm_embeddings_model

    if old:
        model_class = AAOldModel
    elif all_atoms:
        model_class = AAModel
    else:
        model_class = CGModel

    model_kwargs = dict(
        t_to_sigma=t_to_sigma,
        device=device,
        no_torsion=args.no_torsion,
        timestep_emb_func=timestep_emb_func,
        num_conv_layers=args.num_conv_layers,
        lig_max_radius=args.max_radius,
        scale_by_sigma=args.scale_by_sigma,
        sigma_embed_dim=args.sigma_embed_dim,
        norm_by_sigma=_get_arg(args, "norm_by_sigma", False),
        ns=args.ns,
        nv=args.nv,
        distance_embed_dim=args.distance_embed_dim,
        cross_distance_embed_dim=args.cross_distance_embed_dim,
        batch_norm=not args.no_batch_norm,
        dropout=args.dropout,
        use_second_order_repr=args.use_second_order_repr,
        cross_max_distance=args.cross_max_distance,
        dynamic_max_cross=args.dynamic_max_cross,
        smooth_edges=_get_arg(args, "smooth_edges", False),
        odd_parity=_get_arg(args, "odd_parity", False),
        lm_embedding_type=lm_embedding_type,
        confidence_mode=confidence_mode,
        confidence_dropout=_get_arg(args, "confidence_dropout", 0.0),
        confidence_no_batchnorm=_get_arg(args, "confidence_no_batchnorm", False),
        affinity_prediction=_get_arg(args, "affinity_prediction", False),
        parallel=_get_arg(args, "parallel", 1),
        num_confidence_outputs=(
            len(args.rmsd_classification_cutoff) + 1
            if isinstance(_get_arg(args, "rmsd_classification_cutoff"), list)
            else 1
        ),
        atom_num_confidence_outputs=(
            len(args.atom_rmsd_classification_cutoff) + 1
            if isinstance(_get_arg(args, "atom_rmsd_classification_cutoff"), list)
            else 1
        ),
        parallel_aggregators=_get_arg(args, "parallel_aggregators", ""),
        fixed_center_conv=not _get_arg(args, "not_fixed_center_conv", False),
        no_aminoacid_identities=_get_arg(args, "no_aminoacid_identities", False),
        include_miscellaneous_atoms=_get_arg(args, "include_miscellaneous_atoms", False),
        sh_lmax=_get_arg(args, "sh_lmax", 2),
        differentiate_convolutions=not _get_arg(args, "no_differentiate_convolutions", False),
        tp_weights_layers=_get_arg(args, "tp_weights_layers", 2),
        num_prot_emb_layers=_get_arg(args, "num_prot_emb_layers", 0),
        reduce_pseudoscalars=_get_arg(args, "reduce_pseudoscalars", False),
        embed_also_ligand=_get_arg(args, "embed_also_ligand", False),
        atom_confidence=_get_arg(args, "atom_confidence_loss_weight", 0.0) > 0.0,
        sidechain_pred=(
            (_has_arg(args, "sidechain_loss_weight") and args.sidechain_loss_weight > 0)
            or (_has_arg(args, "backbone_loss_weight") and args.backbone_loss_weight > 0)
        ),
        depthwise_convolution=_get_arg(args, "depthwise_convolution", False),
    )
    if model_class is AAModel:
        model_kwargs["crop_beyond"] = _get_arg(args, "crop_beyond", None)
    elif model_class is AAOldModel:
        for key in (
            "atom_num_confidence_outputs",
            "differentiate_convolutions",
            "tp_weights_layers",
            "num_prot_emb_layers",
            "reduce_pseudoscalars",
            "embed_also_ligand",
            "atom_confidence",
            "sidechain_pred",
            "depthwise_convolution",
        ):
            model_kwargs.pop(key, None)
        model_kwargs["lm_embedding_type"] = (
            "esm" if _get_arg(args, "esm_embeddings_path") is not None else None
        )
        model_kwargs["use_old_atom_encoder"] = _get_arg(args, "use_old_atom_encoder", True)

    model = model_class(**model_kwargs)

    if device.type == "cuda" and not no_parallel and _get_arg(args, "dataset") != "torsional":
        model = DataParallel(model)
    model.to(device)
    return model


def build_score_model(
    model_args,
    device,
    no_parallel=False,
    confidence_mode=False,
    old=False,
):
    t_to_sigma = partial(t_to_sigma_compl, args=model_args)
    model = get_model(
        model_args,
        device,
        t_to_sigma=t_to_sigma,
        no_parallel=no_parallel,
        confidence_mode=confidence_mode,
        old=old,
    )
    return model, t_to_sigma


def load_score_model(
    model_dir,
    ckpt,
    device=None,
    no_parallel=True,
    confidence_mode=False,
    old=False,
    strict=True,
):
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model_args = load_model_args(model_dir)
    model, t_to_sigma = build_score_model(
        model_args=model_args,
        device=device,
        no_parallel=no_parallel,
        confidence_mode=confidence_mode,
        old=old,
    )

    checkpoint_path = Path(model_dir) / ckpt
    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))
    if isinstance(state_dict, dict) and "model" in state_dict and "optimizer" in state_dict:
        model.load_state_dict(state_dict["model"], strict=strict)
        if "ema_weights" in state_dict and getattr(model_args, "ema_rate", None) is not None:
            ema_weights = ExponentialMovingAverage(model.parameters(), decay=model_args.ema_rate)
            ema_weights.load_state_dict(state_dict["ema_weights"], device=device)
            ema_weights.copy_to(model.parameters())
    else:
        model.load_state_dict(state_dict, strict=strict)
    model = model.to(device)
    model.eval()
    return model, model_args, t_to_sigma