diff --git a/docs/src/getting-started/usage.rst b/docs/src/getting-started/usage.rst index 06242589f..dca27c522 100644 --- a/docs/src/getting-started/usage.rst +++ b/docs/src/getting-started/usage.rst @@ -58,7 +58,7 @@ the current directory and type Evaluation ########## -The sub-command to evaluate a already trained model is +The sub-command to evaluate an already trained model is .. code-block:: bash @@ -66,7 +66,23 @@ The sub-command to evaluate a already trained model is .. literalinclude:: ../../../examples/usage.sh :language: bash - :lines: 9- + :lines: 9-25 + + +Exporting +######### + +Exporting a model is very useful if you want to use it in other frameworks, +especially in molecular dynamics simulations. +The sub-command to export an already trained model is + +.. code-block:: bash + + metatensor-models export + +.. literalinclude:: ../../../examples/usage.sh + :language: bash + :lines: 25- In the next tutorials we show how adjust the dataset section of ``options.yaml`` file to use it for your own datasets. diff --git a/examples/usage.sh b/examples/usage.sh index 279d1bd23..84eac8d61 100644 --- a/examples/usage.sh +++ b/examples/usage.sh @@ -22,3 +22,9 @@ head -n 20 output.xyz # All command line flags of the eval sub-command can be listed via metatensor-models eval --help + +# For example, the following command + +metatensor-models export model.pt + +# creates an `exported-model.pt` file that contains the exported model. diff --git a/src/metatensor/models/cli/export_model.py b/src/metatensor/models/cli/export_model.py index ce4415a09..48d08e961 100644 --- a/src/metatensor/models/cli/export_model.py +++ b/src/metatensor/models/cli/export_model.py @@ -47,4 +47,7 @@ def export_model(model: str, output: Optional[str]) -> None: # Export the model wrapper = MetatensorAtomisticModel(loaded_model.eval(), loaded_model.capabilities) - wrapper.export("exported-model.pt") + if output is None: + wrapper.export("exported-model.pt") + else: + wrapper.export(output) diff --git a/tests/cli/test_export_model.py b/tests/cli/test_export_model.py new file mode 100644 index 000000000..ba13c4e0e --- /dev/null +++ b/tests/cli/test_export_model.py @@ -0,0 +1,32 @@ +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + + +# Execute the setup script which will make sum_over_samples saveable. +current_dir = os.path.dirname(__file__) +setup_path = os.path.join(current_dir, "..", "..", "scripts", "setup.py") +exec(open(setup_path).read()) + + +RESOURCES_PATH = Path(__file__).parent.resolve() / ".." / "resources" + + +@pytest.mark.parametrize("output", [None, "exported.pt"]) +def test_export(monkeypatch, tmp_path, output): + """Test that the export cli runs without an error raise.""" + monkeypatch.chdir(tmp_path) + shutil.copy(RESOURCES_PATH / "bpnn-model.pt", "bpnn-model.pt") + + command = ["metatensor-models", "export", "bpnn-model.pt"] + + if output is not None: + command += ["-o", output] + else: + output = "exported-model.pt" + + subprocess.check_call(command) + assert Path(output).is_file() diff --git a/tests/resources/bpnn-model.pt b/tests/resources/bpnn-model.pt index 0fed0ed14..2a0158978 100644 Binary files a/tests/resources/bpnn-model.pt and b/tests/resources/bpnn-model.pt differ diff --git a/tests/resources/model.pt b/tests/resources/model.pt deleted file mode 100644 index a8aa5a3d3..000000000 Binary files a/tests/resources/model.pt and /dev/null differ