Upload folder using huggingface_hub

f15d29e verified 4 months ago

5.25 kB

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from typing import Callable

import torch

from onescience.datapipes.materials.mattergen.chemgraph import ChemGraph
from onescience.datapipes.materials.mattergen.types import PropertySourceId
from .denoiser import GemNetTDenoiser, get_chemgraph_from_denoiser_output
from onescience.modules.embedding.mattergen_property_embeddings import (
    ZerosEmbedding,
    get_property_embeddings,
    get_use_unconditional_embedding,
)

BatchTransform = Callable[[ChemGraph], ChemGraph]


class GemNetTAdapter(GemNetTDenoiser):
    """
    Denoiser layerwise adapter with GemNetT. On top of a model.denoiser.GemNetTDenoiser,
    additionally inputs <property_embeddings_adapt> that specifies extra conditions to be conditioned on.
    """

    def __init__(self, property_embeddings_adapt: torch.nn.ModuleDict, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # ModuleDict[PropertyName, PropertyEmbedding] -- conditions adding by this adapter
        self.property_embeddings_adapt = torch.nn.ModuleDict(property_embeddings_adapt)

        # sanity check keys are required by the adapter that already exist in the base model
        assert all(
            [
                k not in self.property_embeddings.keys()
                for k in self.property_embeddings_adapt.keys()
            ]
        ), f"One of adapter conditions {self.property_embeddings_adapt.keys()} already exists in base model {self.property_embeddings.keys()}, please remove."

        # we make the choice that new adapter fields do not alter the unconditional score
        # we therefore need the unconditional embedding for all properties added in the adapter
        # to return 0. We hack the unconditional embedding module here to achieve that
        for property_embedding in self.property_embeddings_adapt.values():
            property_embedding.unconditional_embedding_module = ZerosEmbedding(
                hidden_dim=property_embedding.unconditional_embedding_module.hidden_dim,
            )

    def forward(
        self,
        x: ChemGraph,
        t: torch.Tensor,
    ) -> ChemGraph:
        """
        augment <z_per_crystal> with <self.condition_embs_adapt>.
        """
        (frac_coords, lattice, atom_types, num_atoms, batch,) = (
            x["pos"],
            x["cell"],
            x["atomic_numbers"],
            x["num_atoms"],
            x.get_batch_idx("pos"),
        )
        # (num_atoms, hidden_dim) (num_crysts, 3)
        t_enc = self.noise_level_encoding(t).to(lattice.device)
        z_per_crystal = t_enc

        # shape = (Nbatch, sum(hidden_dim of all properties in condition_on_adapt))
        conditions_base_model: torch.Tensor = get_property_embeddings(
            property_embeddings=self.property_embeddings, batch=x
        )

        if len(conditions_base_model) > 0:
            z_per_crystal = torch.cat([z_per_crystal, conditions_base_model], dim=-1)

        # compose into a dict
        conditions_adapt_dict = {}
        conditions_adapt_mask_dict = {}
        for cond_field, property_embedding in self.property_embeddings_adapt.items():
            conditions_adapt_dict[cond_field] = property_embedding.forward(batch=x)
            try:
                conditions_adapt_mask_dict[cond_field] = get_use_unconditional_embedding(
                    batch=x, cond_field=cond_field
                )
            except KeyError:
                # no values have been provided for the conditional field,
                # interpret this as the user wanting an unconditional score
                conditions_adapt_mask_dict[cond_field] = torch.ones_like(
                    x["num_atoms"], dtype=torch.bool
                ).reshape(-1, 1)

        output = self.gemnet(
            z=z_per_crystal,
            frac_coords=frac_coords,
            atom_types=atom_types,
            num_atoms=num_atoms,
            batch=batch,
            lengths=None,
            angles=None,
            lattice=lattice,
            # we construct the graph on the fly, hence pass None for these:
            edge_index=None,
            to_jimages=None,
            num_bonds=None,
            cond_adapt=conditions_adapt_dict,
            cond_adapt_mask=conditions_adapt_mask_dict,  # when True use unconditional embedding
        )

        pred_atom_types = self.fc_atom(output.node_embeddings)

        return get_chemgraph_from_denoiser_output(
            pred_atom_types=pred_atom_types,
            pred_lattice_eps=output.stress,
            pred_cart_pos_eps=output.forces,
            training=self.training,
            element_mask_func=self.element_mask_func,
            x_input=x,
        )

    @property
    def cond_fields_model_was_trained_on(self) -> list[PropertySourceId]:
        """
        We adopt the convention that all property embeddings are stored in torch.nn.ModuleDicts of
        name property_embeddings or property_embeddings_adapt in the case of a fine tuned model.
        This function returns the list of all field names that a given score model was trained to
        condition on.
        """
        return list(self.property_embeddings) + list(self.property_embeddings_adapt)

	# Copyright (c) Microsoft Corporation.
	# Licensed under the MIT License.

	from typing import Callable

	import torch

	from onescience.datapipes.materials.mattergen.chemgraph import ChemGraph
	from onescience.datapipes.materials.mattergen.types import PropertySourceId
	from .denoiser import GemNetTDenoiser, get_chemgraph_from_denoiser_output
	from onescience.modules.embedding.mattergen_property_embeddings import (
	ZerosEmbedding,
	get_property_embeddings,
	get_use_unconditional_embedding,
	)

	BatchTransform = Callable[[ChemGraph], ChemGraph]


	class GemNetTAdapter(GemNetTDenoiser):
	"""
	Denoiser layerwise adapter with GemNetT. On top of a model.denoiser.GemNetTDenoiser,
	additionally inputs <property_embeddings_adapt> that specifies extra conditions to be conditioned on.
	"""

	def __init__(self, property_embeddings_adapt: torch.nn.ModuleDict, args, *kwargs):
	super().__init__(args, *kwargs)

	# ModuleDict[PropertyName, PropertyEmbedding] -- conditions adding by this adapter
	self.property_embeddings_adapt = torch.nn.ModuleDict(property_embeddings_adapt)

	# sanity check keys are required by the adapter that already exist in the base model
	assert all(
	[
	k not in self.property_embeddings.keys()
	for k in self.property_embeddings_adapt.keys()
	]
	), f"One of adapter conditions {self.property_embeddings_adapt.keys()} already exists in base model {self.property_embeddings.keys()}, please remove."

	# we make the choice that new adapter fields do not alter the unconditional score
	# we therefore need the unconditional embedding for all properties added in the adapter
	# to return 0. We hack the unconditional embedding module here to achieve that
	for property_embedding in self.property_embeddings_adapt.values():
	property_embedding.unconditional_embedding_module = ZerosEmbedding(
	hidden_dim=property_embedding.unconditional_embedding_module.hidden_dim,
	)

	def forward(
	self,
	x: ChemGraph,
	t: torch.Tensor,
	) -> ChemGraph:
	"""
	augment <z_per_crystal> with <self.condition_embs_adapt>.
	"""
	(frac_coords, lattice, atom_types, num_atoms, batch,) = (
	x["pos"],
	x["cell"],
	x["atomic_numbers"],
	x["num_atoms"],
	x.get_batch_idx("pos"),
	)
	# (num_atoms, hidden_dim) (num_crysts, 3)
	t_enc = self.noise_level_encoding(t).to(lattice.device)
	z_per_crystal = t_enc

	# shape = (Nbatch, sum(hidden_dim of all properties in condition_on_adapt))
	conditions_base_model: torch.Tensor = get_property_embeddings(
	property_embeddings=self.property_embeddings, batch=x
	)

	if len(conditions_base_model) > 0:
	z_per_crystal = torch.cat([z_per_crystal, conditions_base_model], dim=-1)

	# compose into a dict
	conditions_adapt_dict = {}
	conditions_adapt_mask_dict = {}
	for cond_field, property_embedding in self.property_embeddings_adapt.items():
	conditions_adapt_dict[cond_field] = property_embedding.forward(batch=x)
	try:
	conditions_adapt_mask_dict[cond_field] = get_use_unconditional_embedding(
	batch=x, cond_field=cond_field
	)
	except KeyError:
	# no values have been provided for the conditional field,
	# interpret this as the user wanting an unconditional score
	conditions_adapt_mask_dict[cond_field] = torch.ones_like(
	x["num_atoms"], dtype=torch.bool
	).reshape(-1, 1)

	output = self.gemnet(
	z=z_per_crystal,
	frac_coords=frac_coords,
	atom_types=atom_types,
	num_atoms=num_atoms,
	batch=batch,
	lengths=None,
	angles=None,
	lattice=lattice,
	# we construct the graph on the fly, hence pass None for these:
	edge_index=None,
	to_jimages=None,
	num_bonds=None,
	cond_adapt=conditions_adapt_dict,
	cond_adapt_mask=conditions_adapt_mask_dict, # when True use unconditional embedding
	)

	pred_atom_types = self.fc_atom(output.node_embeddings)

	return get_chemgraph_from_denoiser_output(
	pred_atom_types=pred_atom_types,
	pred_lattice_eps=output.stress,
	pred_cart_pos_eps=output.forces,
	training=self.training,
	element_mask_func=self.element_mask_func,
	x_input=x,
	)

	@property
	def cond_fields_model_was_trained_on(self) -> list[PropertySourceId]:
	"""
	We adopt the convention that all property embeddings are stored in torch.nn.ModuleDicts of
	name property_embeddings or property_embeddings_adapt in the case of a fine tuned model.

	This function returns the list of all field names that a given score model was trained to
	condition on.
	"""
	return list(self.property_embeddings) + list(self.property_embeddings_adapt)