Allow exporting models from remote locations

docs/src/dev-docs/utils/index.rst

@@ -5,24 +5,8 @@ This is the API for the ``utils`` module of ``metatrain``.
 
 .. toctree::
    :maxdepth: 1
+   :glob:
 
    additive/index
    data/index
-   architectures
-   devices
-   dtype
-   errors
-   evaluate_model
-   external_naming
-   export
-   io
-   jsonschema
-   logging
-   loss
-   metrics
-   neighbor_lists
-   omegaconf
-   output_gradient
-   per_atom
-   transfer
-   units
+   ./*

docs/src/getting-started/checkpoints.rst

@@ -25,20 +25,30 @@ or
     mtt train options.yaml -c model.ckpt
 
 Checkpoints can also be turned into exported models using the ``export`` sub-command.
+The command requires the *architecture name* and the saved checkpoint *path* as
+positional arguments
 
 .. code-block:: bash
 
-    mtt export model.ckpt -o model.pt
+    mtt export experimental.soap_bpnn model.ckpt -o model.pt
 
 or
 
 .. code-block:: bash
 
-    mtt export model.ckpt --output model.pt
+    mtt export experimental.soap_bpnn model.ckpt --output model.pt
+
+For a export of distribution of models the ``export`` command also supports parsing
+models from remote locations. To export a remote model you can provide a URL instead of
+a file path.
+
+.. code-block:: bash
+
+    mtt export experimental.soap_bpnn https://my.url.com/model.ckpt --output model.pt
 
 Keep in mind that a checkpoint (``.ckpt``) is only a temporary file, which can have
 several dependencies and may become unusable if the corresponding architecture is
-updated. In constrast, exported models (``.pt``) act as standalone files.
-For long-term usage, you should export your model! Exporting a model is also necessary
-if you want to use it in other frameworks, especially in molecular simulations
-(see the :ref:`tutorials`).
+updated. In constrast, exported models (``.pt``) act as standalone files. For long-term
+usage, you should export your model! Exporting a model is also necessary if you want to
+use it in other frameworks, especially in molecular simulations (see the
+:ref:`tutorials`).

examples/programmatic/llpr/llpr.py

@@ -28,18 +28,19 @@
 #
 
 import torch
-from metatensor.torch.atomistic import load_atomistic_model
+
+from metatrain.utils.io import load_model
 
 
 # %%
 #
-# Exported models can be loaded using the `load_atomistic_model` function from the
-# metatensor.torch.atomistic` module. The function requires the path to the exported
-# model and, for many models, also the path to the respective extensions directory.
-# Both are produced during the training process.
+# Models can be loaded using the :func:`metatrain.utils.io.load_model` function from
+# the. For already exported models The function requires the path to the exported model
+# and, for many models, also the path to the respective extensions directory. Both are
+# produced during the training process.
 
 
-model = load_atomistic_model("model.pt", extensions_directory="extensions/")
+model = load_model("model.pt", extensions_directory="extensions/")
 
 # %%
 #

pyproject.toml

@@ -12,7 +12,7 @@ dependencies = [
     "ase < 3.23.0",
     "metatensor-learn==0.3.0",
     "metatensor-operations==0.3.0",
-    "metatensor-torch==0.6.0",
+    "metatensor-torch==0.6.1",
     "jsonschema",
     "omegaconf",
     "python-hostlist",

src/metatrain/cli/eval.py

@@ -22,6 +22,7 @@
 )
 from ..utils.errors import ArchitectureError
 from ..utils.evaluate_model import evaluate_model
+from ..utils.io import load_model
 from ..utils.logging import MetricLogger
 from ..utils.metrics import MAEAccumulator, RMSEAccumulator
 from ..utils.neighbor_lists import (
@@ -94,7 +95,8 @@ def _add_eval_model_parser(subparser: argparse._SubParsersAction) -> None:
 def _prepare_eval_model_args(args: argparse.Namespace) -> None:
     """Prepare arguments for eval_model."""
     args.options = OmegaConf.load(args.options)
-    args.model = metatensor.torch.atomistic.load_atomistic_model(
+    # models for evaluation are already exported. Don't have to pass the `name` argument
+    args.model = load_model(
         path=args.__dict__.pop("path"),
         extensions_directory=args.__dict__.pop("extensions_directory"),
     )

src/metatrain/cli/export.py

@@ -3,11 +3,10 @@
 from pathlib import Path
 from typing import Any, Union
 
-import torch
+from metatensor.torch.atomistic import is_atomistic_model
 
-from ..utils.architectures import find_all_architectures, import_architecture
-from ..utils.export import is_exported
-from ..utils.io import check_file_extension
+from ..utils.architectures import find_all_architectures
+from ..utils.io import check_file_extension, load_model
 from .formatter import CustomHelpFormatter
 
 
@@ -40,7 +39,10 @@ def _add_export_model_parser(subparser: argparse._SubParsersAction) -> None:
     parser.add_argument(
         "path",
         type=str,
-        help="Saved model which should be exported",
+        help=(
+            "Saved model which should be exported. Path can be either a URL or a "
+            "local file."
+        ),
     )
     parser.add_argument(
         "-o",
@@ -55,32 +57,31 @@ def _add_export_model_parser(subparser: argparse._SubParsersAction) -> None:
 
 def _prepare_export_model_args(args: argparse.Namespace) -> None:
     """Prepare arguments for export_model."""
-    architecture_name = args.__dict__.pop("architecture_name")
-    architecture = import_architecture(architecture_name)
-
-    args.model = architecture.__model__.load_checkpoint(args.__dict__.pop("path"))
+    args.model = load_model(
+        path=args.__dict__.pop("path"),
+        architecture_name=args.__dict__.pop("architecture_name"),
+    )
 
 
 def export_model(model: Any, output: Union[Path, str] = "exported-model.pt") -> None:
-    """Export a trained model to allow it to make predictions.
+    """Export a trained model allowing 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.
 
     :param model: model to be exported
-    :param output: path to save the exported model
+    :param output: path to save the model
     """
-    path = str(check_file_extension(filename=output, extension=".pt"))
+    path = str(
+        Path(check_file_extension(filename=output, extension=".pt"))
+        .absolute()
+        .resolve()
+    )
+    extensions_path = str(Path("extensions/").absolute().resolve())
 
-    if is_exported(model):
-        logger.info(f"The model is already exported. Saving it to `{path}`.")
-        torch.jit.save(model, path)
-    else:
-        extensions_path = "extensions/"
-        logger.info(
-            f"Exporting model to '{path}' and extensions to '{extensions_path}'"
-        )
-        mts_atomistic_model = model.export()
-        mts_atomistic_model.save(path, collect_extensions=extensions_path)
-        logger.info("Model exported successfully")
+    if not is_atomistic_model(model):
+        model = model.export()
+
+    model.save(path, collect_extensions=extensions_path)
+    logger.info(f"Model exported to '{path}' and extensions to '{extensions_path}'")

src/metatrain/cli/train.py

@@ -11,7 +11,6 @@
 
 import numpy as np
 import torch
-from metatensor.torch.atomistic import load_atomistic_model
 from omegaconf import DictConfig, OmegaConf
 
 from .. import PACKAGE_ROOT
@@ -31,7 +30,7 @@
 from ..utils.devices import pick_devices
 from ..utils.distributed.logging import is_main_process
 from ..utils.errors import ArchitectureError
-from ..utils.io import check_file_extension
+from ..utils.io import check_file_extension, load_model
 from ..utils.jsonschema import validate
 from ..utils.omegaconf import BASE_OPTIONS, check_units, expand_dataset_config
 from .eval import _eval_targets
@@ -454,8 +453,9 @@ def train_model(
     # EVALUATE FINAL MODEL ####
     ###########################
 
-    mts_atomistic_model = load_atomistic_model(
-        str(output_checked), extensions_directory=extensions_path
+    mts_atomistic_model = load_model(
+        path=output_checked,
+        extensions_directory=extensions_path,
     )
     mts_atomistic_model = mts_atomistic_model.to(final_device)
 

src/metatrain/experimental/alchemical_model/model.py

@@ -7,6 +7,7 @@
 from metatensor.torch.atomistic import (
     MetatensorAtomisticModel,
     ModelCapabilities,
+    ModelMetadata,
     ModelOutput,
     NeighborListOptions,
     System,
@@ -16,7 +17,6 @@
 from ...utils.additive import ZBL
 from ...utils.data.dataset import DatasetInfo
 from ...utils.dtype import dtype_to_str
-from ...utils.export import export
 from .utils import systems_to_torch_alchemical_batch
 
 
@@ -185,7 +185,7 @@ def export(self) -> MetatensorAtomisticModel:
             dtype=dtype_to_str(dtype),
         )
 
-        return export(model=self, model_capabilities=capabilities)
+        return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities)
 
     def set_composition_weights(
         self,

src/metatrain/experimental/gap/model.py

@@ -14,6 +14,7 @@
 from metatensor.torch.atomistic import (
     MetatensorAtomisticModel,
     ModelCapabilities,
+    ModelMetadata,
     ModelOutput,
     System,
 )
@@ -22,7 +23,6 @@
 from metatrain.utils.data.dataset import DatasetInfo
 
 from ...utils.additive import ZBL, CompositionModel
-from ...utils.export import export
 
 
 class GAP(torch.nn.Module):
@@ -253,7 +253,7 @@ def export(self) -> MetatensorAtomisticModel:
             self._subset_of_regressors.export_torch_script_model()
         )
 
-        return export(model=self, model_capabilities=capabilities)
+        return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities)
 
     def set_composition_weights(
         self,

src/metatrain/experimental/pet/model.py

@@ -8,6 +8,7 @@
 from metatensor.torch.atomistic import (
     MetatensorAtomisticModel,
     ModelCapabilities,
+    ModelMetadata,
     ModelOutput,
     NeighborListOptions,
     System,
@@ -20,7 +21,6 @@
 
 from ...utils.additive import ZBL
 from ...utils.dtype import dtype_to_str
-from ...utils.export import export
 from .utils import systems_to_batch_dict
 
 
@@ -195,4 +195,4 @@ def export(self) -> MetatensorAtomisticModel:
             supported_devices=["cpu", "cuda"],  # and not __supported_devices__
             dtype=dtype_to_str(dtype),
         )
-        return export(model=self, model_capabilities=capabilities)
+        return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities)

src/metatrain/experimental/pet/tests/test_exported.py

@@ -12,7 +12,6 @@
 from metatrain.experimental.pet import PET as WrappedPET
 from metatrain.utils.architectures import get_default_hypers
 from metatrain.utils.data import DatasetInfo, TargetInfo
-from metatrain.utils.export import export
 from metatrain.utils.neighbor_lists import (
     get_requested_neighbor_lists,
     get_system_with_neighbor_lists,
@@ -57,7 +56,7 @@ def test_to(device):
         supported_devices=["cpu", "cuda"],
     )
 
-    exported = export(model, capabilities)
+    exported = model.export()
     exported.to(device=device, dtype=dtype)
 
     system = System(

src/metatrain/experimental/soap_bpnn/model.py

@@ -8,6 +8,7 @@
 from metatensor.torch.atomistic import (
     MetatensorAtomisticModel,
     ModelCapabilities,
+    ModelMetadata,
     ModelOutput,
     System,
 )
@@ -18,7 +19,6 @@
 
 from ...utils.additive import ZBL, CompositionModel
 from ...utils.dtype import dtype_to_str
-from ...utils.export import export
 
 
 class Identity(torch.nn.Module):
@@ -339,7 +339,7 @@ def export(self) -> MetatensorAtomisticModel:
             dtype=dtype_to_str(dtype),
         )
 
-        return export(model=self, model_capabilities=capabilities)
+        return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities)
 
     def add_output(self, output_name: str) -> None:
         """Add a new output to the self."""

src/metatrain/utils/evaluate_model.py

@@ -9,11 +9,11 @@
     ModelEvaluationOptions,
     ModelOutput,
     System,
+    is_atomistic_model,
     register_autograd_neighbors,
 )
 
 from .data import TargetInfo
-from .export import is_exported
 from .output_gradient import compute_gradient
 
 
@@ -221,7 +221,7 @@ def _strain_gradients_to_block(gradients_list):
 def _get_outputs(
     model: Union[torch.nn.Module, torch.jit._script.RecursiveScriptModule]
 ):
-    if is_exported(model):
+    if is_atomistic_model(model):
         return model.capabilities().outputs
     else:
         return model.outputs
@@ -237,7 +237,7 @@ def _get_model_outputs(
     targets: Dict[str, TargetInfo],
     check_consistency: bool,
 ) -> Dict[str, TensorMap]:
-    if is_exported(model):
+    if is_atomistic_model(model):
         # put together an EvaluationOptions object
         options = ModelEvaluationOptions(
             length_unit="",  # this is only needed for unit conversions in MD engines