Addressing PR comments

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

deepmd/dpmodel/infer/deep_eval.py

@@ -393,7 +393,7 @@ def get_model_def_script(self) -> dict:
         """Get model definition script."""
         return json.loads(self.dp.get_model_def_script())
 
-    def get_model(self):
+    def get_model(self) -> "BaseModel":
         """Get the dpmodel BaseModel.
 
         Returns

deepmd/infer/deep_eval.py

@@ -342,7 +342,7 @@ def get_observed_types(self) -> dict:
         raise NotImplementedError("Not implemented in this backend.")
 
     @abstractmethod
-    def get_model(self):
+    def get_model(self) -> Any:
         """Get the model module implemented by the deep learning framework.
 
         For PyTorch, this returns the nn.Module. For Paddle, this returns
@@ -700,7 +700,7 @@ def get_observed_types(self) -> dict:
         """Get observed types (elements) of the model during data statistics."""
         return self.deep_eval.get_observed_types()
 
-    def get_model(self):
+    def get_model(self) -> Any:
         """Get the model module implemented by the deep learning framework.
 
         For PyTorch, this returns the nn.Module. For Paddle, this returns

deepmd/jax/infer/deep_eval.py

@@ -421,7 +421,7 @@ def get_model_def_script(self) -> dict:
         """Get model definition script."""
         return json.loads(self.dp.get_model_def_script())
 
-    def get_model(self):
+    def get_model(self) -> Any:
         """Get the JAX model as BaseModel.
 
         Returns

deepmd/pd/infer/deep_eval.py

@@ -46,6 +46,10 @@
 if TYPE_CHECKING:
     import ase.neighborlist
 
+    from deepmd.pd.model.model.model import (
+        BaseModel,
+    )
+
 
 class DeepEval(DeepEvalBackend):
     """Paddle backend implementation of DeepEval.
@@ -506,7 +510,7 @@ def get_model_size(self) -> dict:
             "total": sum_param_des + sum_param_fit,
         }
 
-    def get_model(self):
+    def get_model(self) -> "BaseModel":
         """Get the Paddle model.
 
         Returns

deepmd/pt/infer/deep_eval.py

@@ -75,6 +75,10 @@
 if TYPE_CHECKING:
     import ase.neighborlist
 
+    from deepmd.pt.model.model.model import (
+        BaseModel,
+    )
+
 log = logging.getLogger(__name__)
 
 
@@ -706,7 +710,7 @@ def get_observed_types(self) -> dict:
             "observed_type": sort_element_type(observed_type_list),
         }
 
-    def get_model(self):
+    def get_model(self) -> "BaseModel":
         """Get the PyTorch model.
 
         Returns

deepmd/tf/infer/deep_eval.py

@@ -1126,7 +1126,7 @@ def get_model_def_script(self) -> dict:
         model_def_script = script.decode("utf-8")
         return json.loads(model_def_script)["model"]
 
-    def get_model(self):
+    def get_model(self) -> "tf.Graph":
         """Get the TensorFlow graph.
 
         Returns