feat(pt): add model branch alias

[iProzd]

Introduces model branch alias and info fields to model configuration, adds utility functions for handling model branch dictionaries, and updates related modules to use alias-based lookup and provide detailed branch information. Enhances multi-task model usability and improves logging of available model branches.

example:

dp --pt show 0415_compat_new.pt model-branch

[2025-08-14 10:05:54,246] DEEPMD WARNING To get the best performance, it is recommended to adjust the number of threads by setting the environment variables OMP_NUM_THREADS, DP_INTRA_OP_PARALLELISM_THREADS, and DP_INTER_OP_PARALLELISM_THREADS. See https://deepmd.rtfd.io/parallelism/ for more information.
[2025-08-14 10:05:59,122] DEEPMD INFO    This is a multitask model
[2025-08-14 10:05:59,122] DEEPMD INFO    Available model branches are ['Dai2023Alloy', 'Zhang2023Cathode', 'Gong2023Cluster', 'Yang2023ab', 'UniPero', 'Huang2021Deep-PBE', 'Liu2024Machine', 'Zhang2021Phase', 'Jinag2021Accurate', 'Chen2023Modeling', 'Wen2021Specialising', 'Wang2022Classical', 'Wang2022Tungsten', 'Wu2021Deep', 'Huang2021Deep-PBEsol', 'Transition1x', 'Wang2021Generalizable', 'Wu2021Accurate', 'MPTraj', 'Li2025APEX', 'Shi2024SSE', 'Tuo2023Hybrid', 'Unke2019PhysNet', 'Shi2024Electrolyte', 'ODAC23', 'Alex2D', 'OMAT24', 'SPICE2', 'OC20M', 'OC22', 'Li2025General', 'RANDOM'], where 'RANDOM' means using a randomly initialized fitting net.
[2025-08-14 10:05:59,125] DEEPMD INFO    Detailed information:
+-----------------------+------------------------------+--------------------------------+--------------------------------+
| Model Branch          | Alias                        | description                    | observed_type                  |
+-----------------------+------------------------------+--------------------------------+--------------------------------+
| Dai2023Alloy          | Alloys, Domains_Alloy        | The dataset contains           | ['La', 'Fe', 'Ho', 'Cu', 'Sn', |
|                       |                              | structure-energy-force-virial  | 'Cd', 'Y', 'Be', 'V', 'Sm',    |
|                       |                              | data for 53 typical metallic   | 'In', 'Pr', 'Mo', 'Mn', 'Gd',  |
|                       |                              | elements in alloy systems,     | 'Ru', 'Nd', 'Li', 'Tm', 'K',   |
|                       |                              | including ~9000 intermetallic  | 'Pt', 'Ir', 'Na', 'Hf', 'Dy',  |
|                       |                              | compounds and FCC, BCC, HCP    | 'Ca', 'Nb', 'Au', 'Sr', 'Si',  |
|                       |                              | structures. It consists of two | 'Ge', 'Co', 'W', 'Cr', 'Zn',   |
|                       |                              | parts: DFT-generated relaxed   | 'Ag', 'Ti', 'Ni', 'Zr', 'Pd',  |
|                       |                              | and deformed structures, and   | 'Os', 'Ta', 'Rh', 'Sc', 'Tb',  |
|                       |                              | randomly distorted structures  | 'Al', 'Ga', 'Re', 'Lu', 'Er',  |
|                       |                              | produced covering pure metals, | 'Mg', 'Ce', 'Pb']              |
|                       |                              | solid solutions, and           |                                |
|                       |                              | intermetallics with vacancies. |                                |
+-----------------------+------------------------------+--------------------------------+--------------------------------+
| OMAT24                | Default, Materials, Omat24   | OMat24 is a large-scale open   | ['La', 'Fe', 'Cu', 'Cd', 'Be', |
|                       |                              | dataset containing over 110    | 'Ar', 'V', 'Sm', 'In', 'Pm',   |
|                       |                              | million DFT calculations       | 'Pr', 'Mn', 'Ru', 'He', 'Nd',  |
|                       |                              | spanning diverse structures    | 'Th', 'Pa', 'K', 'Pt', 'Yb',   |
|                       |                              | and compositions. It is        | 'Dy', 'Sr', 'Co', 'Np', 'Cr',  |
|                       |                              | designed to support AI-driven  | 'Tl', 'Br', 'Se', 'Ni', 'Zr',  |
|                       |                              | materials discovery by         | 'Pu', 'O', 'Xe', 'Tb', 'Ga',   |
|                       |                              | providing broad and deep       | 'Lu', 'H', 'Ne', 'Er', 'Ce',   |
|                       |                              | coverage of chemical space.    | 'I', 'Kr', 'Ho', 'Cs', 'Sn',   |
|                       |                              |                                | 'Rb', 'Y', 'N', 'F', 'Mo',     |
|                       |                              |                                | 'Gd', 'B', 'Li', 'Tm', 'Sb',   |
|                       |                              |                                | 'Ir', 'Hf', 'Na', 'Ca', 'Nb',  |
|                       |                              |                                | 'Au', 'As', 'Si', 'Ge', 'W',   |
|                       |                              |                                | 'Zn', 'Hg', 'Ag', 'Bi', 'Ti',  |
|                       |                              |                                | 'Os', 'Cl', 'Pd', 'P', 'U',    |
|                       |                              |                                | 'Tc', 'Ta', 'Ba', 'Rh', 'Sc',  |
|                       |                              |                                | 'C', 'S', 'Te', 'Al', 'Re',    |
|                       |                              |                                | 'Eu', 'Mg', 'Pb', 'Ac']        |
+-----------------------+------------------------------+--------------------------------+--------------------------------+

Summary by CodeRabbit

• New Features

  • Alias-based multi-task branch selection for evaluation and fine-tuning; new API to query model alias/branch info; show now prints a detailed model-branch table.

• Documentation

  • Model config gains optional fields to declare branch aliases and per-branch info (PyTorch-only).

• Examples

  • Added a two-task PyTorch example demonstrating aliases, shared components, and per-branch info.

• Tests

  • Tests include the new example and now filter out table-like show output.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds alias-aware model-branch utilities and integrates them into show, DeepEval, and finetune flows; extends model argument schema with optional alias/info fields; adds an ASCII table renderer for branch details; updates examples and tests to cover alias-based multi-task setups.

Changes

Cohort / File(s)	Summary
Model-branch utilities deepmd/utils/model_branch_dict.py	New get_model_dict(model_dict) returning (model_alias_dict, model_branch_dict) and OrderedDictTableWrapper for ASCII table rendering of branch info.
Entrypoint: show enhancements deepmd/entrypoints/show.py	Imports get_model_dict and OrderedDictTableWrapper; when showing "model-branch" uses get_model_dict(...), logs a "Detailed information" ASCII table, and prefers branch info.observed_type over always calling DeepEval.
PT inference alias resolution deepmd/pt/infer/deep_eval.py	Adds alias-aware multitask head resolution using get_model_dict; supports Default alias, integer head fallback, case-insensitive alias matching, remaps state_dict keys for branch-prefixed params, updates error messaging, and adds public DeepEval.get_model_branch().
PT fine-tune alias resolution deepmd/pt/utils/finetune.py	Uses get_model_dict to validate/map chosen alias to canonical branch for multi-task finetuning; updates assertions and error guidance to reference aliases/model-branch.
Argument schema updates deepmd/utils/argcheck.py	Adds optional model_branch_alias (list) and info (dict) to standard_model_args with local docs (doc_model_branch_alias, doc_info), marked PyTorch-only.
Examples / Tests examples/water_multi_task/.../input_torch_with_alias.json, source/tests/common/test_examples.py, source/tests/pt/test_dp_show.py	Adds a two-task PyTorch example using aliases/info and includes it in multi-task test inputs; tightens dp show test filtering to ignore table-like output and "Detailed information" lines.

Sequence Diagram(s)

sequenceDiagram
  participant User
  participant ShowCmd as "dp show"
  participant DeepEval
  participant Utils as get_model_dict
  Note over ShowCmd,Utils: Show "model-branch" path
  ShowCmd->>DeepEval: load model for inspection
  DeepEval->>Utils: get_model_dict(model_dict)
  Utils-->>DeepEval: (model_alias_dict, model_branch_dict)
  DeepEval-->>ShowCmd: model_branch_dict
  ShowCmd->>Utils: OrderedDictTableWrapper(model_branch_dict).as_table()
  Utils-->>ShowCmd: ASCII table ("Detailed information")
  ShowCmd-->>User: prints branches + table

Loading

sequenceDiagram
  participant Caller
  participant DeepEval
  participant Utils as get_model_dict
  Note over Caller,DeepEval: Head resolution for evaluation
  Caller->>DeepEval: init(head?)
  DeepEval->>Utils: get_model_dict(model_dict) [if multitask]
  Utils-->>DeepEval: (model_alias_dict, model_branch_dict)
  DeepEval->>DeepEval: map/validate head (alias / int / Default / case-insensitive)
  DeepEval->>DeepEval: remap state_dict keys for chosen branch
  DeepEval-->>Caller: proceed with evaluation

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~35 minutes

Possibly related PRs

• feat(pt): Add command to check the available model branches in multi-task pre-trained model(Issue #3742) #3796 — Overlaps in adding model-branch listing/show functionality and tests.
• feat(pt): add observed-type option for dp show #4820 — Related changes to observed-type handling and dp show / DeepEval interactions.
• feat: DeepEval.get_model_def_script and common dp show #4131 — Related modifications to DeepEval/show pathways for exposing model-definition/branch information.

Suggested reviewers

• njzjz
• wanghan-iapcm
• Chengqian-Zhang

Tip

🔌 Remote MCP (Model Context Protocol) integration is now available!

Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats.

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

📥 Commits

Reviewing files that changed from the base of the PR and between 3421fa5 and 3780f90.

📒 Files selected for processing (1)

• deepmd/utils/argcheck.py (2 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• deepmd/utils/argcheck.py

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (15)

• GitHub Check: Build wheels for cp311-win_amd64
• GitHub Check: Build wheels for cp310-manylinux_aarch64
• GitHub Check: Build wheels for cp311-manylinux_x86_64
• GitHub Check: Build wheels for cp311-macosx_arm64
• GitHub Check: Test C++ (false)
• GitHub Check: Build C++ (cuda120, cuda)
• GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
• GitHub Check: Test Python (2, 3.12)
• GitHub Check: Test Python (3, 3.12)
• GitHub Check: Build C++ (cpu, cpu)
• GitHub Check: Test Python (4, 3.9)
• GitHub Check: Test Python (3, 3.9)
• GitHub Check: Build C++ (rocm, rocm)
• GitHub Check: Build C++ (cuda, cuda)
• GitHub Check: Build C++ (clang, clang)

✨ Finishing Touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR/Issue comments)

Type @coderabbitai help to get the list of available commands.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Status, Documentation and Community

• Visit our Status Page to check the current availability of CodeRabbit.
• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (5)

deepmd/utils/argcheck.py (1)

2308-2340: Constrain types for model_branch_alias and clarify schema

model_branch_alias is declared as a bare list, which weakens validation. Since aliases are strings, prefer list[str] for stricter schema checking. info can remain as dict (free-form metadata), which aligns with the intended flexibility.

Apply this diff:

-            Argument(
-                "model_branch_alias",
-                list,
-                optional=True,
-                default=[],
-                doc=doc_only_pt_supported + doc_model_branch_alias,
-            ),
+            Argument(
+                "model_branch_alias",
+                list[str],
+                optional=True,
+                default=[],
+                doc=doc_only_pt_supported + doc_model_branch_alias,
+            ),

deepmd/utils/model_branch_dict.py (3)

88-96: Avoid dict.keys() in iteration for simplicity

Use direct iteration over the dict to satisfy linters and improve clarity.

Apply this diff:

-        for _, payload in self.data.items():
-            info = payload.get("info") or {}
-            for k in info.keys():
+        for _, payload in self.data.items():
+            info = payload.get("info") or {}
+            for k in info:
                 if k not in seen:
                     seen.add(k)
                     self.info_keys.append(k)

---

166-176: Remove unused loop index in enumerate

The loop index i is unused. Drop enumerate for readability.

Apply this diff:

-        for i, row_cells in enumerate(wrapped_rows):
+        for row_cells in wrapped_rows:
             # Determine the maximum number of wrapped lines in this row
             max_lines = max(len(cell) for cell in row_cells)

---

44-46: Remove meta comment

The comment “generated with GPT for formatted print” is not actionable and can be confusing in source control history.

Apply this diff:

-# generated with GPT for formatted print

deepmd/entrypoints/show.py (1)

40-48: Surface aliases alongside branches in the log

Since selection accepts aliases, include them explicitly in the “Available” message to reduce confusion.

Apply this diff:

-        log.info(
-            f"Available model branches are {model_branches}, "
-            f"where 'RANDOM' means using a randomly initialized fitting net."
-        )
+        aliases = [k for k in model_alias_dict.keys() if k not in model_params["model_dict"]]
+        log.info(
+            f"Available model branches: {model_branches}; aliases: {aliases}. "
+            f"'RANDOM' means using a randomly initialized fitting net."
+        )

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 88b71e8 and d1ba115.

📒 Files selected for processing (5)

• deepmd/entrypoints/show.py (2 hunks)
• deepmd/pt/infer/deep_eval.py (3 hunks)
• deepmd/pt/utils/finetune.py (2 hunks)
• deepmd/utils/argcheck.py (2 hunks)
• deepmd/utils/model_branch_dict.py (1 hunks)

🧰 Additional context used

🪛 Ruff (0.12.2)

deepmd/utils/model_branch_dict.py

92-92: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

---

166-166: Loop control variable i not used within loop body

Rename unused i to _i

(B007)

🔇 Additional comments (1)

deepmd/pt/infer/deep_eval.py (1)

269-279: New public API get_model_branch is a useful addition

The method cleanly exposes alias and branch info and falls back correctly for single-task models. Good addition.

[codecov]

Codecov Report

❌ Patch coverage is 80.00000% with 24 lines in your changes missing coverage. Please review.
✅ Project coverage is 84.29%. Comparing base (3f0087b) to head (3780f90).
⚠️ Report is 74 commits behind head on devel.

Files with missing lines	Patch %	Lines
deepmd/pt/infer/deep_eval.py	45.00%	11 Missing ⚠️
deepmd/utils/model_branch_dict.py	87.20%	11 Missing ⚠️
deepmd/entrypoints/show.py	75.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##            devel    #4883      +/-   ##
==========================================
- Coverage   84.29%   84.29%   -0.01%     
==========================================
  Files         703      704       +1     
  Lines       68728    68842     +114     
  Branches     3572     3573       +1     
==========================================
+ Hits        57934    58027      +93     
- Misses       9653     9675      +22     
+ Partials     1141     1140       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[coderabbitai]

Actionable comments posted: 0

🔭 Outside diff range comments (3)

source/tests/pt/test_dp_show.py (3)

59-61: Fix assertions that accidentally test only the last substring.

Using chained string "and" returns the last string, so the test only checks one token, weakening coverage.

Apply this diff:

-        assert (
-            "{'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut': 4.0"
-        ) in results[2]
+        assert all(
+            s in results[2]
+            for s in ("{'type': 'se_e2_a'", "'sel': [46, 92, 4]", "'rcut': 4.0")
+        )

---

185-195: Fix multitask descriptor assertions to check all expected substrings.

Same issue: chained constants via "and" reduce to the last string.

Apply this diff:

-        assert (
-            "model_1"
-            and "'type': 'se_e2_a'"
-            and "'sel': [46, 92, 4]"
-            and "'rcut_smth': 0.5"
-        ) in results[4]
+        assert all(
+            s in results[4]
+            for s in ("model_1", "'type': 'se_e2_a'", "'sel': [46, 92, 4]", "'rcut_smth': 0.5")
+        )
-        assert (
-            "model_2"
-            and "'type': 'se_e2_a'"
-            and "'sel': [46, 92, 4]"
-            and "'rcut_smth': 0.5"
-        ) in results[5]
+        assert all(
+            s in results[5]
+            for s in ("model_2", "'type': 'se_e2_a'", "'sel': [46, 92, 4]", "'rcut_smth': 0.5")
+        )

---

227-228: Fix singletask (frozen) descriptor assertion to check all substrings.

Same chained "and" issue here.

Apply this diff:

-        assert (
-            "'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut_smth': 0.5"
-        ) in results[2]
+        assert all(
+            s in results[2]
+            for s in ("'type': 'se_e2_a'", "'sel': [46, 92, 4]", "'rcut_smth': 0.5")
+        )

🧹 Nitpick comments (2)

source/tests/pt/test_dp_show.py (2)

169-175: Robustly filter table output in multitask show tests.

Good call filtering ASCII tables and headers to stabilize assertions across richer show outputs.

For reuse and clarity, consider extracting the filter into a small helper to avoid repeating magic substrings:

def _filter_show_output(lines: list[str]) -> list[str]:
    banned = ("DEEPMD WARNING", "|", "+-", "Detailed information")
    return [ln for ln in lines if all(b not in ln for b in banned)]

Then use:

results = _filter_show_output(f.getvalue().split("\n")[:-1])

---

163-175: Optional: add a targeted assertion that the alias table is emitted.

Given the new alias-aware table in show, consider a small focused test that asserts the presence of "Detailed information" and expected alias headers when model_dict contains alias/info, to guard against regressions in the table path. Keep the existing filtered test for the non-table content.

I can draft a self-contained test method that runs dp --pt show with model-branch and asserts the table header and a couple of alias values without depending on exact column widths. Want me to propose it?

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these settings in your CodeRabbit configuration.

📥 Commits

Reviewing files that changed from the base of the PR and between d1ba115 and b5740e5.

📒 Files selected for processing (4)

• deepmd/entrypoints/show.py (3 hunks)
• examples/water_multi_task/pytorch_example/input_torch_with_alias.json (1 hunks)
• source/tests/common/test_examples.py (1 hunks)
• source/tests/pt/test_dp_show.py (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• deepmd/entrypoints/show.py

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (9)

• GitHub Check: Build wheels for cp311-macosx_x86_64
• GitHub Check: Build wheels for cp310-manylinux_aarch64
• GitHub Check: Build wheels for cp311-win_amd64
• GitHub Check: Build wheels for cp311-manylinux_x86_64
• GitHub Check: Build wheels for cp311-manylinux_x86_64
• GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
• GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
• GitHub Check: Analyze (c-cpp)
• GitHub Check: Analyze (python)

🔇 Additional comments (4)

examples/water_multi_task/pytorch_example/input_torch_with_alias.json (3)

1-148: Well-structured multi-task PT example with aliases. LGTM.

The example cleanly demonstrates model_branch_alias and info usage with shared components. The configuration is consistent and should pass arg checks and normalizations.

---

61-64: No action needed — observed_type values are subsets of shared type_map_all

Checked with the provided script; output: "OK: observed_type values are subsets of the shared type map."

• Verified locations:
  • examples/water_multi_task/pytorch_example/input_torch_with_alias.json — lines 61–64
  • examples/water_multi_task/pytorch_example/input_torch_with_alias.json — lines 71–74

---

60-65: Aliases verified — unique (case-insensitive) and no reserved name used

Ran the provided validation script against examples/water_multi_task/pytorch_example/input_torch_with_alias.json. Output: "OK: aliases are unique and do not use reserved names."

• Files/locations checked: examples/water_multi_task/pytorch_example/input_torch_with_alias.json (lines 60–65, 70–75)
• Action: no changes required.

source/tests/common/test_examples.py (1)

71-72: Add alias-based multi-task input to examples test. LGTM.

Including input_torch_with_alias.json in input_files_multi ensures arg checking covers the new alias/info fields.

[wanghan-iapcm]

1. should the branch names case sensitive or not ?
2. we may adjust first and 2nd col width of the printed table to the maximum size of the branch name and alias name.

[coderabbitai]

Actionable comments posted: 0

♻️ Duplicate comments (1)

deepmd/entrypoints/show.py (1)

40-40: Past suggestion addressed: unpack from get_model_dict correctly

Using "_, model_branch_dict = get_model_dict(...)" matches the earlier review suggestion. LGTM.

🧹 Nitpick comments (4)

deepmd/entrypoints/show.py (4)

31-49: Nice integration of alias-aware branch listing with a detailed table

• Properly guards the "model-branch" path to multitask models.
• The human-readable list plus the ASCII table is a good UX improvement.

One nit: you list "RANDOM" in the short list but it’s absent from the table. Consider adding a row for clarity (see next comment with a concrete diff).

---

45-48: Optionally include a 'RANDOM' row in the table for consistency

Right now "RANDOM" appears in the plain list but not in the table, which may confuse users. You can inject a minimal placeholder row before rendering.

Apply this diff:

         _, model_branch_dict = get_model_dict(model_params["model_dict"])
         log.info(
             f"Available model branches are {model_branches}, "
             f"where 'RANDOM' means using a randomly initialized fitting net."
         )
-        log.info(
+        # Optionally reflect RANDOM in the table as a placeholder
+        model_branch_dict = dict(model_branch_dict)
+        model_branch_dict.setdefault(
+            "RANDOM",
+            {"alias": [], "info": {"description": "Randomly initialized fitting net"}},
+        )
+        log.info(
             "Detailed information: \n"
             + OrderedDictTableWrapper(model_branch_dict).as_table()
         )

---

91-106: Validate and normalize info.observed_type; dedupe to avoid double-counting

If info.observed_type is mis-typed (e.g., a string) or contains duplicates, logs and totals become misleading. Light validation and normalization would make this robust while preserving order.

Apply this diff:

-                if (
-                    model_params["model_dict"][branch]
-                    .get("info", {})
-                    .get("observed_type", None)
-                    is not None
-                ):
-                    observed_type_list = model_params["model_dict"][branch]["info"][
-                        "observed_type"
-                    ]
-                    observed_types = {
-                        "type_num": len(observed_type_list),
-                        "observed_type": observed_type_list,
-                    }
+                if (
+                    model_params["model_dict"][branch]
+                    .get("info", {})
+                    .get("observed_type", None)
+                    is not None
+                ):
+                    observed_type_list = model_params["model_dict"][branch]["info"]["observed_type"]
+                    # Normalize: accept a single string, preserve order, remove dups
+                    if isinstance(observed_type_list, (str, bytes)):
+                        observed_type_list = [observed_type_list]
+                    if not isinstance(observed_type_list, (list, tuple)):
+                        log.warning(
+                            "Branch %s: info.observed_type must be a list of strings; falling back to model introspection.",
+                            branch,
+                        )
+                        tmp_model = DeepEval(INPUT, head=branch, no_jit=True)
+                        observed_types = tmp_model.get_observed_types()
+                    else:
+                        unique_observed = list(dict.fromkeys(observed_type_list))
+                        observed_types = {
+                            "type_num": len(unique_observed),
+                            "observed_type": unique_observed,
+                        }
                 else:
                     tmp_model = DeepEval(INPUT, head=branch, no_jit=True)
                     observed_types = tmp_model.get_observed_types()

---

105-106: Avoid repeated model instantiation when falling back to introspection

Creating a new DeepEval per branch can be expensive. If the API allows, consider reusing a single instance and switching heads, or caching results per branch to cut load time.

If you want, I can sketch a small cache layer keyed by branch to avoid redundant DeepEval(...) calls when many branches miss observed_type in info.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

📥 Commits

Reviewing files that changed from the base of the PR and between bf31822 and fba956a.

📒 Files selected for processing (1)

• deepmd/entrypoints/show.py (3 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

deepmd/entrypoints/show.py (1)

deepmd/utils/model_branch_dict.py (3)

• OrderedDictTableWrapper (45-177)
• get_model_dict (11-41)
• as_table (115-177)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)

• GitHub Check: Test Python (5, 3.9)
• GitHub Check: Test Python (6, 3.12)
• GitHub Check: Test Python (3, 3.9)
• GitHub Check: Test Python (4, 3.9)
• GitHub Check: Test Python (2, 3.12)
• GitHub Check: Test Python (1, 3.12)
• GitHub Check: Test Python (6, 3.9)
• GitHub Check: Test Python (2, 3.9)
• GitHub Check: Test Python (4, 3.12)
• GitHub Check: Test Python (3, 3.12)
• GitHub Check: Test Python (5, 3.12)
• GitHub Check: Test Python (1, 3.9)
• GitHub Check: Test C++ (true)
• GitHub Check: Test C++ (false)
• GitHub Check: Build wheels for cp311-manylinux_x86_64
• GitHub Check: Build wheels for cp311-macosx_x86_64
• GitHub Check: Build C++ (cuda120, cuda)
• GitHub Check: Build wheels for cp310-manylinux_aarch64
• GitHub Check: Build wheels for cp311-macosx_arm64
• GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
• GitHub Check: Build wheels for cp311-win_amd64
• GitHub Check: Build wheels for cp311-manylinux_x86_64
• GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
• GitHub Check: Build C++ (clang, clang)
• GitHub Check: Build C++ (rocm, rocm)
• GitHub Check: Analyze (python)
• GitHub Check: Build C++ (cuda, cuda)
• GitHub Check: Build C++ (cpu, cpu)
• GitHub Check: Analyze (c-cpp)

🔇 Additional comments (1)

deepmd/entrypoints/show.py (1)

10-13: Imports look correct and align with new utilities

Brings in OrderedDictTableWrapper and get_model_dict from the new helper module; usage below is consistent with their APIs. No concerns about ordering or circularity from this file’s perspective.

[iProzd]

1. should the branch names case sensitive or not ?
2. we may adjust first and 2nd col width of the printed table to the maximum size of the branch name and alias name.

1. For inner design, case sensitive is better, which is assumed in both state keys in checkpoint or model keys in multi-task. While for user interface, maybe it's better we first accept case sensitive names, then case-insensitive names by normalizing to lowercase and mapping to the first matching key.
2. Fixed.

[coderabbitai]

Actionable comments posted: 3

♻️ Duplicate comments (1)

deepmd/utils/model_branch_dict.py (1)

32-41: Replace assert with explicit exception and harden None handling for aliases/info

Asserts can be stripped with Python -O, turning user input errors into silent misbehavior. Also, model_branch_alias or info may be None, which would raise at iteration time.

Apply this diff:

-        alias_list = model_dict[key].get("model_branch_alias", [])
+        alias_list = model_dict[key].get("model_branch_alias") or []
         model_branch_dict[key]["alias"] = alias_list
-        branch_info = model_dict[key].get("info", {})
+        branch_info = model_dict[key].get("info") or {}
         model_branch_dict[key]["info"] = branch_info
         for alias in alias_list:
-            assert alias not in model_alias_dict, (
-                f"Alias {alias} for model_branch {key} already exists in model_branch {model_alias_dict[alias]}!"
-            )
+            if alias in model_alias_dict:
+                raise ValueError(
+                    f"Alias '{alias}' for model_branch '{key}' already exists in model_branch "
+                    f"'{model_alias_dict[alias]}'!"
+                )
             model_alias_dict[alias] = key

🧹 Nitpick comments (9)

deepmd/utils/model_branch_dict.py (6)

91-97: Drop redundant .keys() usage when iterating and checking membership

Minor cleanup to satisfy linters (SIM118) and improve readability.

Apply this diff:

-        for _, payload in self.data.items():
-            info = payload.get("info") or {}
-            for k in info.keys():
-                if k not in seen:
-                    seen.add(k)
-                    self.info_keys.append(k)
+        for _, payload in self.data.items():
+            info = payload.get("info") or {}
+            for k in info:
+                if k not in seen:
+                    seen.add(k)
+                    self.info_keys.append(k)

---

127-131: Use direct dict iteration to compute branch width

Matches Ruff’s SIM118 and is a minor readability win.

Apply this diff:

-        for branch in self.data.keys():
+        for branch in self.data:
             branch_col_width = max(branch_col_width, len(str(branch)))

---

133-137: Alias column may wrap unexpectedly; compute width from the rendered alias string

You calculate width from the maximum single alias, but display a comma-joined string. For multiple aliases, the actual cell content can exceed the computed width, causing unintended wrapping.

Apply this diff:

-        alias_col_width = len(self.headers[1])  # "Alias"
-        for payload in self.data.values():
-            alias_list = payload.get("alias", [])
-            for alias in alias_list:
-                alias_col_width = max(alias_col_width, len(str(alias)))
+        alias_col_width = len(self.headers[1])  # "Alias"
+        for payload in self.data.values():
+            alias_str = ", ".join(map(str, payload.get("alias", [])))
+            alias_col_width = max(alias_col_width, len(alias_str))

Also applies to: 145-147

---

198-198: Rename unused loop index

The loop variable i is unused (B007). Rename to _ to signal intentional discard.

Apply this diff:

-        for i, row_cells in enumerate(wrapped_rows):
+        for _, row_cells in enumerate(wrapped_rows):

---

98-101: Nit: Header text mismatch with docstring

Docstring says “Model Branch Name” while header uses “Model Branch”. Pick one for consistency.

Example:

-        self.headers: list[str] = ["Model Branch", "Alias", *self.info_keys]
+        self.headers: list[str] = ["Model Branch Name", "Alias", *self.info_keys]

---

12-26: Add type hints for clarity and downstream tooling

get_model_dict lacks type hints; adding them improves IDE help and self-documentation.

Apply this diff:

-from typing import (
-    Any,
-    Optional,
-)
+from typing import Any, Optional, Dict, Tuple, OrderedDict as TOrderedDict
@@
-def get_model_dict(model_dict):
+def get_model_dict(
+    model_dict: Dict[str, Dict[str, Any]]
+) -> Tuple[Dict[str, str], Dict[str, Dict[str, Any]]]:

deepmd/pt/infer/deep_eval.py (3)

125-127: Drop unused variable to avoid linter noise

model_branch_dict is not used in this method.

Apply this diff:

-                model_alias_dict, model_branch_dict = get_model_dict(
+                model_alias_dict, _ = get_model_dict(
                     self.input_param["model_dict"]
                 )

---

140-148: Case-insensitive fallback should only run for string inputs

If head is not a string, calling .lower() will raise. Also, keep behavior unchanged for non-strings.

Apply this diff:

-                if head not in model_alias_dict:
-                    # preprocess with potentially case-insensitive input
-                    head_lower = head.lower()
-                    for mk in model_alias_dict:
-                        if mk.lower() == head_lower:
-                            # mapped the first matched head
-                            head = mk
-                            break
+                if isinstance(head, str) and head not in model_alias_dict:
+                    # case-insensitive fallback: map to the first matching key
+                    head_lower = head.lower()
+                    for mk in model_alias_dict:
+                        if mk.lower() == head_lower:
+                            head = mk
+                            break

---

153-163: Optional: record the resolved branch and log selection

Expose the resolved branch for downstream consumers (debugging, telemetry, or UIs) and log it once.

Apply this diff:

-                head = model_alias_dict[head]
+                head = model_alias_dict[head]
+                self.selected_model_branch = head  # for introspection
+                log.info(f"Selected multitask head: {head}")

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

📥 Commits

Reviewing files that changed from the base of the PR and between fba956a and 3421fa5.

📒 Files selected for processing (2)

• deepmd/pt/infer/deep_eval.py (4 hunks)
• deepmd/utils/model_branch_dict.py (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

deepmd/pt/infer/deep_eval.py (2)

deepmd/utils/model_branch_dict.py (1)

• get_model_dict (12-42)

deepmd/utils/finetune.py (1)

• get_model_branch (47-49)

🪛 Ruff (0.12.2)

deepmd/utils/model_branch_dict.py

93-93: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

---

129-129: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

---

198-198: Loop control variable i not used within loop body

Rename unused i to _i

(B007)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (19)

• GitHub Check: Test C++ (false)
• GitHub Check: Test C++ (true)
• GitHub Check: Test Python (6, 3.9)
• GitHub Check: Test Python (6, 3.12)
• GitHub Check: Test Python (5, 3.9)
• GitHub Check: Test Python (1, 3.9)
• GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
• GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
• GitHub Check: Build wheels for cp311-macosx_x86_64
• GitHub Check: Build wheels for cp311-win_amd64
• GitHub Check: Build wheels for cp310-manylinux_aarch64
• GitHub Check: Build wheels for cp311-manylinux_x86_64
• GitHub Check: Build wheels for cp311-manylinux_x86_64
• GitHub Check: Build wheels for cp311-macosx_arm64
• GitHub Check: Build C++ (cuda120, cuda)
• GitHub Check: Build C++ (clang, clang)
• GitHub Check: Build C++ (rocm, rocm)
• GitHub Check: Analyze (python)
• GitHub Check: Analyze (c-cpp)

🔇 Additional comments (2)

deepmd/pt/infer/deep_eval.py (2)

129-133: Good defaulting behavior with clear logging

Defaulting to the “Default” alias when present and logging the decision improves UX for multitask models.

---

283-293: LGTM: clear and useful branch introspection for callers

The new get_model_branch() method returns both alias mapping and branch metadata, and provides a sensible single-task fallback.