feat(pt): add embedding sub cli for pt models#5571
Open
OutisLi wants to merge 4 commits into
Open
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
This was referenced Jun 22, 2026
Contributor
📝 WalkthroughWalkthroughAdds a new ChangesEmbedding extraction feature
Sequence Diagram(s)sequenceDiagram
participant User
participant main_py as main.py / entrypoints/main.py
participant embedding_py as entrypoints/embedding.py
participant DeepEval as DeepEval (deep_eval.py)
participant PTDeepEval as PyTorch DeepEval (pt/infer/deep_eval.py)
participant Model as SeZMModel / DPModel
User->>main_py: dp embed -m model.pt -s ./data -o out.h5 --dtype fp32
main_py->>embedding_py: embedding(model, system, output, dtype="fp32")
loop per discovered system
embedding_py->>embedding_py: DeepmdData(system_path)
embedding_py->>DeepEval: eval_embedding(coords, cells, atype, dtype="fp32")
DeepEval->>DeepEval: _standard_input(coords, cells, atype, mixed_type)
DeepEval->>PTDeepEval: deep_eval.eval_embedding(coords, cells, atype, dtype="fp32")
PTDeepEval->>Model: model.forward_embedding(coord, atype, box)
Model-->>PTDeepEval: descriptor, atomic_feature, structural_feature
PTDeepEval-->>DeepEval: float32 numpy arrays
DeepEval-->>embedding_py: (descriptor, atomic_feature, structural_feature)
embedding_py->>embedding_py: write HDF5 group + gzip datasets + attributes
end
embedding_py-->>User: out.h5
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/entrypoints/embedding.py`:
- Around line 116-120: The datalist.read().splitlines() call in the datafile
parsing section keeps blank lines, which can result in empty strings being
passed to DeepmdData and causing failures. Filter out blank lines when reading
from the datafile by adding a filter condition that excludes empty or
whitespace-only entries from the all_sys list. This ensures only valid system
entries are processed while preserving the original structure of the code.
In `@deepmd/pt/infer/deep_eval.py`:
- Around line 1208-1260: The eval_embedding method directly calls
self.dp.model["Default"].forward_embedding(...) which bypasses the data modifier
that would normally be applied through ModelWrapper. Add a guard condition that
checks if self.modifier is not None before proceeding with the forward_embedding
call, and raise an appropriate error or fail fast if a modifier exists, since
modifier-aware embedding processing is not yet implemented for this code path.
In `@deepmd/pt/model/atomic_model/linear_atomic_model.py`:
- Around line 341-349: Add explicit `strict=False` parameter to both zip() calls
in the list comprehensions to satisfy Ruff B905. The first zip() call combines
self.get_model_rcuts() and self.get_model_nsels(), and the second zip() call
combines self.mixed_types_list, raw_nlists, and self.get_model_sels(). Update
both invocations to include the strict parameter to explicitly indicate that the
iterables may have different lengths.
In `@deepmd/pt/model/task/fitting.py`:
- Around line 855-868: The atomic_feature tensor allocation in the
return_atomic_feature block accesses self.neuron[-1] without checking if the
neuron list is empty. When neuron=[] is configured, this causes an IndexError.
Fix this by replacing the hardcoded self.neuron[-1] with a conditional check
that determines the correct feature dimension. If self.neuron is not empty, use
self.neuron[-1], otherwise determine the feature width from an alternative
source such as the output shape of ll.call_until_last(xx) or from the
descriptor's feature dimension to ensure the atomic_feature tensor is allocated
with the correct size regardless of neuron configuration.
In `@doc/inference/embedding.md`:
- Line 44: On line 44 in the embedding.md file, remove the shell prompt marker
`$ ` from the command to comply with markdownlint rule MD014. Change `$ dp embed
--help` to just `dp embed --help` to fix the lint violation while preserving the
command documentation.
In `@source/tests/pt/model/test_embedding.py`:
- Around line 260-271: The tolerance values (rtol=1e-10 and atol=1e-12) in both
np.testing.assert_allclose calls for eval_descriptor and eval_fitting_last_layer
are too strict for float32 comparisons across separate forward passes, causing
intermittent test failures. Increase both the relative tolerance (rtol) and
absolute tolerance (atol) parameters in both assert_allclose calls to more
reasonable values that account for floating-point precision variations while
still validating correctness.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: cf1c59a0-671a-4d28-8ec5-c902cadd310c
📒 Files selected for processing (30)
deepmd/entrypoints/embedding.pydeepmd/entrypoints/eval_desc.pydeepmd/entrypoints/main.pydeepmd/infer/deep_eval.pydeepmd/main.pydeepmd/pt/infer/deep_eval.pydeepmd/pt/model/atomic_model/base_atomic_model.pydeepmd/pt/model/atomic_model/dp_atomic_model.pydeepmd/pt/model/atomic_model/linear_atomic_model.pydeepmd/pt/model/atomic_model/sezm_atomic_model.pydeepmd/pt/model/descriptor/sezm_nn/dens.pydeepmd/pt/model/model/dp_model.pydeepmd/pt/model/model/make_model.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/model/task/dipole.pydeepmd/pt/model/task/ener.pydeepmd/pt/model/task/fitting.pydeepmd/pt/model/task/invar_fitting.pydeepmd/pt/model/task/polarizability.pydeepmd/pt/model/task/sezm_ener.pydoc/inference/embedding.mddoc/inference/index.rstdoc/inference/python.mddoc/model/dpa4.mddoc/test/test.mdsource/tests/infer/test_models.pysource/tests/pt/model/test_embedding.pysource/tests/pt/model/test_linear_atomic_model.pysource/tests/pt/test_eval_desc.pysource/tests/pt/test_oom_retry.py
💤 Files with no reviewable changes (4)
- source/tests/pt/test_oom_retry.py
- deepmd/pt/model/atomic_model/sezm_atomic_model.py
- deepmd/pt/model/model/dp_model.py
- deepmd/pt/model/descriptor/sezm_nn/dens.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5571 +/- ##
==========================================
+ Coverage 82.14% 82.17% +0.02%
==========================================
Files 900 901 +1
Lines 104139 104266 +127
Branches 4471 4474 +3
==========================================
+ Hits 85550 85676 +126
Misses 17178 17178
- Partials 1411 1412 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
njzjz-bot
suggested changes
Jun 22, 2026
njzjz-bot
left a comment
njzjz-bot
left a comment
Contributor
There was a problem hiding this comment.
Thanks for adding the embedding extraction path. I think this needs a few changes before merge, mainly around backend/API compatibility. The PR adds a PyTorch-only eval_embedding implementation, but the CLI/docs deprecate backend-generic descriptor APIs (dp eval-desc, DeepPot.eval_descriptor) in favor of dp embed / eval_embedding. That currently sends TensorFlow/Paddle users from a working descriptor path to an unsupported one. There is also no descriptor-only mode in the replacement command, so users who only need descriptors pay for fitting features and larger HDF5 output.
Reviewed by OpenClaw 2026.6.8 (844f405) (model: custom-chat-jinzhezeng-group/gpt-5.5).
njzjz-bot
reviewed
Jun 22, 2026
njzjz-bot
reviewed
Jun 22, 2026
njzjz-bot
reviewed
Jun 22, 2026
njzjz-bot
reviewed
Jun 22, 2026
njzjz-bot
reviewed
Jun 22, 2026
njzjz-bot
reviewed
Jun 22, 2026
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
deepmd/infer/deep_eval.py (1)
695-708:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate the dtype wording in
eval_embedding’s docstring.Line 707 still says all embeddings are returned as float32, but Lines 737-738 now allow
"fp64"and"native".Proposed docstring adjustment
- embeddings are returned as float32, which is ample for downstream - analysis. + embeddings are returned as float32 by default, which is ample for + downstream analysis. Pass ``dtype="fp64"`` or ``dtype="native"`` to + request a different output precision.Also applies to: 737-738
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/infer/deep_eval.py` around lines 695 - 708, The docstring for the eval_embedding method on line 712-713 states that all three embeddings are returned as float32, but the implementation now supports configurable dtypes including "fp64" and "native" (as shown in the dtype parameter). Update the docstring statement that mentions float32 to reflect that the returned dtype depends on the dtype parameter passed to the function, removing the hardcoded reference to float32 and indicating that the embeddings are returned in the specified dtype format.deepmd/pt/infer/deep_eval.py (1)
1151-1164:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate the float32-only wording now that
dtypeis configurable.Line 1163 says all embeddings are returned as float32, but Lines 1192-1193 now document
"fp64"and"native"as valid output dtypes.Proposed docstring adjustment
- total energy. All three embeddings are returned as float32, which is - ample for downstream analysis. + total energy. All three embeddings are returned as float32 by default, + which is ample for downstream analysis. Pass ``dtype="fp64"`` or + ``dtype="native"`` to request a different output precision.Also applies to: 1192-1193
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt/infer/deep_eval.py` around lines 1151 - 1164, The docstring for the evaluate_descriptor_atomic_feature_and_structural_feature method (or similar function name around line 1151-1164) contains outdated information stating that all three embeddings are always returned as float32. Since the function now includes a configurable dtype parameter that accepts "fp32", "fp64", and "native" options (as documented in lines 1192-1193), update the docstring to reflect that the output dtype depends on the dtype parameter passed to the function, rather than claiming they are exclusively float32. Remove or modify the sentence that incorrectly claims fixed float32 output and replace it with text indicating the dtype is configurable via the dtype parameter.
♻️ Duplicate comments (1)
source/tests/pt/model/test_embedding.py (1)
264-279:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winLoosen tolerances for float32 cross-call comparisons.
The tolerances
rtol=1e-10, atol=1e-12are tighter than float32 machine epsilon (~1e-7) and can cause flaky test failures across independent forward passes. Graph context shows float32 comparisons elsewhere in the codebase usertol=1e-4.Proposed fix
np.testing.assert_allclose( dp.eval_descriptor( self.coord_np, self.cell_np, self.atype_np, dtype="fp32" ), descriptor, - rtol=1e-10, - atol=1e-12, + rtol=1e-5, + atol=1e-6, ) np.testing.assert_allclose( dp.eval_fitting_last_layer( self.coord_np, self.cell_np, self.atype_np, dtype="fp32" ), atomic_feature, - rtol=1e-10, - atol=1e-12, + rtol=1e-5, + atol=1e-6, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt/model/test_embedding.py` around lines 264 - 279, The tolerances rtol=1e-10, atol=1e-12 are too tight for float32 precision and cause flaky test failures. Update both assert_allclose calls for eval_descriptor and eval_fitting_last_layer to use rtol=1e-4 which aligns with float32 comparison standards used elsewhere in the codebase. Change the rtol parameter from 1e-10 to 1e-4 and the atol parameter from 1e-12 to a value appropriate for float32 (typically around 1e-6 or 1e-5) in both assertions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@deepmd/infer/deep_eval.py`:
- Around line 695-708: The docstring for the eval_embedding method on line
712-713 states that all three embeddings are returned as float32, but the
implementation now supports configurable dtypes including "fp64" and "native"
(as shown in the dtype parameter). Update the docstring statement that mentions
float32 to reflect that the returned dtype depends on the dtype parameter passed
to the function, removing the hardcoded reference to float32 and indicating that
the embeddings are returned in the specified dtype format.
In `@deepmd/pt/infer/deep_eval.py`:
- Around line 1151-1164: The docstring for the
evaluate_descriptor_atomic_feature_and_structural_feature method (or similar
function name around line 1151-1164) contains outdated information stating that
all three embeddings are always returned as float32. Since the function now
includes a configurable dtype parameter that accepts "fp32", "fp64", and
"native" options (as documented in lines 1192-1193), update the docstring to
reflect that the output dtype depends on the dtype parameter passed to the
function, rather than claiming they are exclusively float32. Remove or modify
the sentence that incorrectly claims fixed float32 output and replace it with
text indicating the dtype is configurable via the dtype parameter.
---
Duplicate comments:
In `@source/tests/pt/model/test_embedding.py`:
- Around line 264-279: The tolerances rtol=1e-10, atol=1e-12 are too tight for
float32 precision and cause flaky test failures. Update both assert_allclose
calls for eval_descriptor and eval_fitting_last_layer to use rtol=1e-4 which
aligns with float32 comparison standards used elsewhere in the codebase. Change
the rtol parameter from 1e-10 to 1e-4 and the atol parameter from 1e-12 to a
value appropriate for float32 (typically around 1e-6 or 1e-5) in both
assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: d11ec34e-9512-4db8-8896-57b97f9f5cc0
📒 Files selected for processing (9)
deepmd/entrypoints/embedding.pydeepmd/entrypoints/eval_desc.pydeepmd/infer/deep_eval.pydeepmd/main.pydeepmd/pt/infer/deep_eval.pydoc/inference/embedding.mddoc/inference/python.mddoc/test/test.mdsource/tests/pt/model/test_embedding.py
✅ Files skipped from review due to trivial changes (2)
- doc/test/test.md
- doc/inference/embedding.md
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/entrypoints/embedding.py
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
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.
Summary by CodeRabbit
New Features
dp embedCLI command to export model embeddings (descriptor, atomic feature, structural feature) to HDF5 format.DeepPot.eval_embedding()Python API for embedding inference with configurable output dtype (fp32/fp64/native).--headflag.Documentation