Add model branch alias and info support

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.

Author: iProzd

deepmd/entrypoints/show.py

@@ -7,6 +7,10 @@
 from deepmd.utils.econf_embd import (
     sort_element_type,
 )
+from deepmd.utils.model_branch_dict import (
+    OrderedDictTableWrapper,
+    get_model_dict,
+)
 
 log = logging.getLogger(__name__)
 
@@ -33,10 +37,15 @@ def show(
             )
         model_branches = list(model_params["model_dict"].keys())
         model_branches += ["RANDOM"]
+        model_alias_dict, 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(
+            "Detailed information: \n"
+            + OrderedDictTableWrapper(model_branch_dict).as_table()
+        )
     if "type-map" in ATTRIBUTES:
         if model_is_multi_task:
             model_branches = list(model_params["model_dict"].keys())

deepmd/pt/infer/deep_eval.py

@@ -67,6 +67,9 @@
 from deepmd.utils.econf_embd import (
     sort_element_type,
 )
+from deepmd.utils.model_branch_dict import (
+    get_model_dict,
+)
 
 if TYPE_CHECKING:
     import ase.neighborlist
@@ -116,15 +119,25 @@ def __init__(
             self.model_def_script = self.input_param
             self.multi_task = "model_dict" in self.input_param
             if self.multi_task:
+                model_alias_dict, model_branch_dict = get_model_dict(
+                    self.input_param["model_dict"]
+                )
                 model_keys = list(self.input_param["model_dict"].keys())
+                if head is None and "Default" in model_alias_dict:
+                    head = "Default"
                 if isinstance(head, int):
                     head = model_keys[0]
                 assert head is not None, (
-                    f"Head must be set for multitask model! Available heads are: {model_keys}"
+                    f"Head must be set for multitask model! Available heads are: {model_keys}, "
+                    f"use `dp --pt show your_model.pt model-branch` to show detail information."
                 )
-                assert head in model_keys, (
-                    f"No head named {head} in model! Available heads are: {model_keys}"
+                # replace with alias
+                assert head in model_alias_dict, (
+                    f"No head or alias named {head} in model! Available heads are: {model_keys},"
+                    f"use `dp --pt show your_model.pt model-branch` to show detail information."
                 )
+                head = model_alias_dict[head]
+
                 self.input_param = self.input_param["model_dict"][head]
                 state_dict_head = {"_extra_state": state_dict["_extra_state"]}
                 for item in state_dict:
@@ -253,6 +266,17 @@ def get_has_hessian(self):
         """Check if the model has hessian."""
         return self._has_hessian
 
+    def get_model_branch(self):
+        """Get the model branch information."""
+        if "model_dict" in self.model_def_script:
+            model_alias_dict, model_branch_dict = get_model_dict(
+                self.model_def_script["model_dict"]
+            )
+            return model_alias_dict, model_branch_dict
+        else:
+            # single-task model
+            return {"Default": "Default"}, {"Default": {"alias": [], "info": {}}}
+
     def eval(
         self,
         coords: np.ndarray,

deepmd/pt/utils/finetune.py

@@ -12,6 +12,9 @@
 from deepmd.utils.finetune import (
     FinetuneRuleItem,
 )
+from deepmd.utils.model_branch_dict import (
+    get_model_dict,
+)
 
 log = logging.getLogger(__name__)
 
@@ -44,10 +47,13 @@ def get_finetune_rule_single(
             )
         else:
             model_branch_chosen = model_branch_from
-        assert model_branch_chosen in model_dict_params, (
-            f"No model branch named '{model_branch_chosen}'! "
+        model_alias_dict, model_branch_dict = get_model_dict(model_dict_params)
+        assert model_branch_chosen in model_alias_dict, (
+            f"No model branch or alias named '{model_branch_chosen}'! "
             f"Available ones are {list(model_dict_params.keys())}."
+            f"Use `dp --pt show your_model.pt model-branch` to show detail information."
         )
+        model_branch_chosen = model_alias_dict[model_branch_chosen]
         single_config_chosen = deepcopy(model_dict_params[model_branch_chosen])
     old_type_map, new_type_map = (
         single_config_chosen["type_map"],

deepmd/utils/argcheck.py

@@ -2305,6 +2305,10 @@ def model_args(exclude_hybrid=False):
 def standard_model_args() -> Argument:
     doc_descrpt = "The descriptor of atomic environment."
     doc_fitting = "The fitting of physical properties."
+    doc_model_branch_alias = (
+        "The alias of this model branch. This is only used in the multi-task model."
+    )
+    doc_info = "The information of this model branch. This is only used in the multi-task model."
 
     ca = Argument(
         "standard",
@@ -2320,6 +2324,20 @@ def standard_model_args() -> Argument:
                 [fitting_variant_type_args()],
                 doc=doc_fitting,
             ),
+            Argument(
+                "model_branch_alias",
+                list,
+                optional=True,
+                default=[],
+                doc=doc_only_pt_supported + doc_model_branch_alias,
+            ),
+            Argument(
+                "info",
+                dict,
+                optional=True,
+                default={},
+                doc=doc_only_pt_supported + doc_info,
+            ),
         ],
         doc="Standard model, which contains a descriptor and a fitting.",
     )