IO

Functions to be used for handling the serialization of models

metatensor.models.utils.io.check_suffix(filename: str | Path, suffix: str) → str | Path

Check the suffix of a file name and adds if it not existing.

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

Parameters:

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

• suffix (str) – Expected filesuffix i.e. .txt.

Returns:

Checked and probably extended file name.

Return type:

str | Path

metatensor.models.utils.io.save(model: Module, path: str | Path) → None

Saves a model to a checkpoint file.

Along with the model all the metadata needed to load it is saved as well. Checkpointed models will be saved with a .ckpt file ending. If path does not end with this file extensions .ckpt will be added and a warning emitted.

Parameters:

• model (Module) – The model to save.

• path (str | Path) – The path to the file where the model should be saved.

Return type:

None

metatensor.models.utils.io.load(path: str | Path) → Module | RecursiveScriptModule

Loads a checkpoint or an exported model from a file.

Parameters:

path (str | Path) – path to load the model

Returns:

the loaded model

Return type:

Module | RecursiveScriptModule

metatensor.models.utils.io.export(model: Module, path: str | Path) → None

Export a trained model to allow it to make predictions.

This includes predictions within molecular simulation engines. Exported models will be saved with a .pt file ending. If path does not end with this file extensions .pt will be added and a warning emitted.

Parameters:

• model (Module) – model to be exported

• path (str | Path) – path to save the exported model

Return type:

None

metatensor.models.utils.io.is_exported(model: Any)

Check if a model has been exported.

Parameters:

model (Any) – The model to check

Returns:

True if the model has been exported, False otherwise