Implement ModelMetadata storing in checkpoints
#592
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
There was a problem hiding this comment.
Pull Request Overview
This PR introduces functionality to attach and load ModelMetadata objects within model checkpoints and ensures that the export and load procedures for checkpoints incorporate metadata seamlessly.
- Metadata is now stored directly under the "metadata" key in checkpoints.
- The naming convention for default metadata has been updated (from "default_metadata" to "metadata") across different models and trainers.
- The command‐line export interface has been updated to support metadata attachment for checkpoint files.
Reviewed Changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/cli/test_export_model.py | Updated tests to use checkpoint-based model files instead of direct model instantiation. |
| src/metatrain/utils/metadata.py | Revised metadata merging to avoid duplicating references. |
| src/metatrain/soap_bpnn/trainer.py, pet/trainer.py, etc. | Updated checkpoint saving to include metadata from model.metadata. |
| src/metatrain/soap_bpnn/model.py, pet/model.py, gap/model.py, nanopet/model.py | Renamed default_metadata to metadata and updated metadata handling in load/export routines. |
| src/metatrain/cli/export.py | Modified argument handling and export_model logic to support metadata on both checkpoint (.ckpt) and .pt exports. |
Comments suppressed due to low confidence (3)
src/metatrain/soap_bpnn/model.py:173
- [nitpick] The renaming from default_metadata to metadata improves clarity; ensure that the change is consistently reflected in documentation and all related modules.
__metadata__ = ModelMetadata(
tests/cli/test_export_model.py:18
- The removal of the unused model import cleans up the test file; ensure all tests only use the checkpoint-based export approach as intended.
-from metatrain.soap_bpnn import __model__
src/metatrain/cli/export.py:144
- [nitpick] Verify that the merge order in the metadata update is intended; currently, the new metadata is being merged into the checkpoint metadata and then assigned back, which may be confusing. Consider adding a comment or refactoring to make the merge precedence explicit.
append_metadata_references(metadata, current_metadata)
Luthaf
reviewed
May 13, 2025
| @@ -17,6 +17,8 @@ def append_metadata_references(self: ModelMetadata, other: ModelMetadata) -> Non | |||
| if key not in self_dict["references"]: | |||
| self_dict["references"][key] = values | |||
| else: | |||
| self_dict["references"][key] += values | |||
There was a problem hiding this comment.
Should we turn this function into a more general merge_metadata(existing, new)? So more than just the references can be merged when fine-tuning
| output: Union[Path, str], | ||
| extensions: Union[Path, str] = "extensions/", | ||
| token: Optional[str] = None, |
There was a problem hiding this comment.
This should mention that it is specific to HF.
Suggested change
| token: Optional[str] = None, | |
| hf_token: Optional[str] = None, |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
As the title suggests, this PR adds the functionality to store the
ModelMetadataobjects inside the model checkpoints, and allows to runmtt export model.ckpt -o model-with-metadata.ckpt -m metadata.yamlas a valid command to attach the metadata to the existing.ckptfile.mtt exportstill works properly for exporting in.ptfileModel.load_checkpoint()properly stores the metadata that was read from fileTrainer.save_checkpoint()saves themodel.__metadata__Contributor (creator of pull-request) checklist
Reviewer checklist
📚 Documentation preview 📚: https://metatrain--592.org.readthedocs.build/en/592/