chore(cc): merge get backend codes (#4355)

Fix #4308.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Introduced a new function to dynamically determine the backend
framework based on the model file type.
  
- **Improvements**
- Enhanced backend detection logic in multiple classes, allowing for
more flexible model initialization.
- Simplified control flow in the initialization methods of various
components.

- **Bug Fixes**
- Improved error handling for unsupported backends and model formats
during initialization processes.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jinzhe Zeng <[email protected]>

source/api_cc/include/common.h

@@ -15,6 +15,13 @@ namespace deepmd {
 typedef double ENERGYTYPE;
 enum DPBackend { TensorFlow, PyTorch, Paddle, JAX, Unknown };
 
+/**
+ * @brief Get the backend of the model.
+ * @param[in] model The model name.
+ * @return The backend of the model.
+ **/
+DPBackend get_backend(const std::string& model);
+
 struct NeighborListData {
   /// Array stores the core region atom's index
   std::vector<int> ilist;

source/api_cc/src/DataModifier.cc

@@ -28,8 +28,7 @@ void DipoleChargeModifier::init(const std::string& model,
               << std::endl;
     return;
   }
-  // TODO: To implement detect_backend
-  DPBackend backend = deepmd::DPBackend::TensorFlow;
+  const DPBackend backend = get_backend(model);
   if (deepmd::DPBackend::TensorFlow == backend) {
 #ifdef BUILD_TENSORFLOW
     dcm = std::make_shared<deepmd::DipoleChargeModifierTF>(model, gpu_rank,

source/api_cc/src/DeepPot.cc

@@ -39,17 +39,7 @@ void DeepPot::init(const std::string& model,
               << std::endl;
     return;
   }
-  DPBackend backend;
-  if (model.length() >= 4 && model.substr(model.length() - 4) == ".pth") {
-    backend = deepmd::DPBackend::PyTorch;
-  } else if (model.length() >= 3 && model.substr(model.length() - 3) == ".pb") {
-    backend = deepmd::DPBackend::TensorFlow;
-  } else if (model.length() >= 11 &&
-             model.substr(model.length() - 11) == ".savedmodel") {
-    backend = deepmd::DPBackend::JAX;
-  } else {
-    throw deepmd::deepmd_exception("Unsupported model file format");
-  }
+  const DPBackend backend = get_backend(model);
   if (deepmd::DPBackend::TensorFlow == backend) {
 #ifdef BUILD_TENSORFLOW
     dp = std::make_shared<deepmd::DeepPotTF>(model, gpu_rank, file_content);

source/api_cc/src/DeepSpin.cc

@@ -36,14 +36,7 @@ void DeepSpin::init(const std::string& model,
               << std::endl;
     return;
   }
-  DPBackend backend;
-  if (model.length() >= 4 && model.substr(model.length() - 4) == ".pth") {
-    backend = deepmd::DPBackend::PyTorch;
-  } else if (model.length() >= 3 && model.substr(model.length() - 3) == ".pb") {
-    backend = deepmd::DPBackend::TensorFlow;
-  } else {
-    throw deepmd::deepmd_exception("Unsupported model file format");
-  }
+  const DPBackend backend = get_backend(model);
   if (deepmd::DPBackend::TensorFlow == backend) {
 #ifdef BUILD_TENSORFLOW
     dp = std::make_shared<deepmd::DeepSpinTF>(model, gpu_rank, file_content);

source/api_cc/src/DeepTensor.cc

@@ -30,8 +30,7 @@ void DeepTensor::init(const std::string &model,
               << std::endl;
     return;
   }
-  // TODO: To implement detect_backend
-  DPBackend backend = deepmd::DPBackend::TensorFlow;
+  const DPBackend backend = get_backend(model);
   if (deepmd::DPBackend::TensorFlow == backend) {
 #ifdef BUILD_TENSORFLOW
     dt = std::make_shared<deepmd::DeepTensorTF>(model, gpu_rank, name_scope_);

source/api_cc/src/common.cc

@@ -1399,3 +1399,15 @@ void deepmd::print_summary(const std::string& pre) {
             << "set tf inter_op_parallelism_threads: " << num_inter_nthreads
             << std::endl;
 }
+
+deepmd::DPBackend deepmd::get_backend(const std::string& model) {
+  if (model.length() >= 4 && model.substr(model.length() - 4) == ".pth") {
+    return deepmd::DPBackend::PyTorch;
+  } else if (model.length() >= 3 && model.substr(model.length() - 3) == ".pb") {
+    return deepmd::DPBackend::TensorFlow;
+  } else if (model.length() >= 11 &&
+             model.substr(model.length() - 11) == ".savedmodel") {
+    return deepmd::DPBackend::JAX;
+  }
+  throw deepmd::deepmd_exception("Unsupported model file format");
+}