nvalchemi.models.demo.DemoModelWrapper

class nvalchemi.models.demo.DemoModelWrapper(num_atom_types=100, hidden_dim=64)

Wrapper for the demo model that implements the BaseModelMixin interface.

This wrapper is used to implement the BaseModelMixin interface for the demo model.

This demonstrates how BaseModelMixin adds elements to the underlying model that help standardize the model’s interface and behavior, and allow for the original model to be used without modification.

Parameters:

• num_atom_types (int)

• hidden_dim (int)

adapt_input(data, **kwargs)

Adapts the input data to the model’s expected format.

This method is used to convert the input data to the format expected by the model.

Parameters:

• data (AtomicData | Batch)

• kwargs (Any)

Return type:

dict[str, Any]

adapt_output(model_output, data)

Adapts the model output to the framework’s expected format.

The super() implementation will provide the initial OrderedDict with keys that are expected to be present in the model output. This method will then map the model outputs to this OrderedDict.

Technically, this is not necessary for the demo model, but it is included to demonstrate how to override the super() implementation.

Parameters:

• model_output (Any)

• data (AtomicData | Batch)

Return type:

OrderedDict[str, Float[Tensor, ‘B 1’] | Float[Tensor, ‘V 3’] | Float[Tensor, ‘V 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3’] | None]

compute_embeddings(data, **kwargs)

Computes the embeddings for the demo model.

This method is used to compute the embeddings for the demo model.

This implementation assumes that the underlying model does not expose embeddings, and relies on the developer to effectively duplicate the logic of the forward call and extract embeddings.

For models that do expose embeddings, the implementation can be simplified by calling forward or other methods.

Parameters:

• data (AtomicData | Batch) – Input data to compute embeddings for.

• kwargs (Any)

Returns:

Input data with computed embeddings.

Return type:

AtomicData | Batch

property dtype: dtype

Returns the projection layer’s datatype for casting.

property embedding_shapes: dict[str, tuple[int, ...]]

Returns the expected shapes of the embeddings for the demo model.

This serves as an immutable specification of the model’s expected embedding shapes.

Returns:

Expected shapes of the embeddings for the demo model.

Return type:

dict[str, tuple[int, …]]

export_model(path, as_state_dict=False)

Export the demo model without the BaseModelMixin interface.

This is not intended to be used in production, but more as a demonstration of how to export the model as you may with an actual subclass.

Parameters:

• path (Path) – Path to export the model to.

• as_state_dict (bool, optional) – Whether to export the model as a state dictionary. Defaults to False, which pickles the model entirely.

Return type:

None

forward(data, **kwargs)

Forward pass of the demo model wrapper.

This wraps the underlying model’s forward pass, adapting both the input and outputs to the framework’s expected format.

Parameters:

• data (AtomicData | Batch) – Input data to use for the forward pass.

• kwargs (Any)

Returns:

Output data from the forward pass, formatted with the correct key naming conventions.

Return type:

ModelOutputs

property model_card: ModelCard

Returns the model card for the demo model.

This serves as an immutable specification of the model’s capabilities and requirements.

Fields will all default to False, and are expected to be set to True if the model expects an input, or directly outputs a particular property.

Returns:

Model card for the demo model.

Return type:

ModelCard