From 34645f9564cbc23971a2d89dfa9daaa02e2adc33 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Wed, 11 Dec 2024 17:40:39 -0500
Subject: [PATCH] add deepmd-kit multi-task parameters

---
 generator.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/generator.py b/generator.py
index 6cb4550..39a4ac3 100644
--- a/generator.py
+++ b/generator.py
@@ -5,7 +5,10 @@
 import dpgen
 import dpgen2
 import dpdispatcher
-from deepmd.utils.argcheck import gen_args as deepmd_arginfo
+from deepmd.utils.argcheck import (
+    gen_args as deepmd_arginfo,
+    gen_args_multi_task as deepmd_multi_task_arginfo,
+)
 from dpgen.generator.arginfo import run_jdata_arginfo, run_mdata_arginfo
 from dpgen.simplify.arginfo import simplify_jdata_arginfo, simplify_mdata_arginfo
 from dpgen.data.arginfo import (
@@ -96,6 +99,13 @@ def add_parameter(key : str, name: str, argument: Argument, version: str,
         }
     ],
 )
+add_parameter(
+    "deepmd-kit-multi-task",
+    "DeePMD-kit",
+    deepmd_multi_task_arginfo(),
+    deepmd.__version__,
+    "deepmodeling/deepmd-kit",
+)
 
 # dpgen
 add_parameter(