-
Notifications
You must be signed in to change notification settings - Fork 5
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Dump and load extensions #222
Conversation
src/metatensor/models/cli/export.py
Outdated
@@ -63,4 +63,4 @@ def export_model(model: Any, output: Union[Path, str] = "exported-model.pt") -> | |||
torch.jit.save(model, path) | |||
else: | |||
mts_atomistic_model = model.export() | |||
mts_atomistic_model.export(path) | |||
mts_atomistic_model.export(path, collect_extensions="extensions") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should we inform the user in the end where the extensions are saved?
Related to #208
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think this should work now! Very nice @frostedoyster
Maybe add some tests that extensions are saved and loaded correctly, Via CLI and Python would be great.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can you also add the logger.info
referring to the location as in train.
@PicoCentauri Now there are tests that both |
@@ -60,7 +64,10 @@ def export_model(model: Any, output: Union[Path, str] = "exported-model.pt") -> | |||
path = str(check_suffix(filename=output, suffix=".pt")) | |||
|
|||
if is_exported(model): | |||
logger.info(f"The model is already exported. Saving it to `{path}`.") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Cool. I am always struggeling if we should also add a test if these infos are displayed... But I think it is fine for now.
Reopening #201
📚 Documentation preview 📚: https://metatensor-models--222.org.readthedocs.build/en/222/