IO

Functions to be used for handling the serialization of models

metatrain.utils.io.check_file_extension(filename: str | Path, extension: str) → str | Path

Check the file extension of a file name and adds if it is not present.

If filename does not end with extension the extension is added and a warning will be issued.

Parameters:

• filename (str | Path) – Name of the file to be checked.

• extension (str) – Expected file extension i.e. .txt.

Returns:

Checked and probably extended file name.

Return type:

str | Path

metatrain.utils.io.is_exported_file(path: str) → bool

Check if a saved model file has been exported to a MetatensorAtomisticModel.

The functions uses metatensor.torch.atomistic.check_atomistic_model() to verify.

Parameters:

path (str) – model path

Returns:

True if the model has been exported, False otherwise.

Return type:

bool

See also

metatensor.torch.atomistic.is_atomistic_model() to verify if an already loaded model is exported.

metatrain.utils.io.load_model(path: str | Path, extensions_directory: str | Path | None = None, token: str | None = None) → Any

Load checkpoints and exported models from an URL or a local file.

If an exported model should be loaded and requires compiled extensions, their location should be passed using the extensions_directory parameter.

After reading a checkpoint, the returned model can be exported with the model’s own export() method.

Parameters:

• path (str | Path) – local or remote path to a model. For supported URL schemes see urllib.request

• extensions_directory (str | Path | None) – path to a directory containing all extensions required by an exported model

• token (str | None) – HuggingFace API token to download (private) models from HuggingFace

Raises:

• ValueError – if path is a YAML option file and no model

• ValueError – if no archietcture_name is found in the checkpoint

• ValueError – if the architecture_name is not found in the available architectures

Return type:

Any