feat: Add delete overlap option to pth2json and improve model saving logic (#268)

Author: QG-phy

dptb/entrypoints/main.py

@@ -202,6 +202,12 @@ def main_parser() -> argparse.ArgumentParser:
         default=None,
         help="The pth ckpt to be transfered to json.",
     )
+    parser_pth2json.add_argument(
+        "-dels", 
+        "--deleteoverlap",
+        help="Transfer to no overlap version.",
+        action="store_true"
+    )
     parser_pth2json.add_argument(
         "-o",
         "--outdir",

dptb/entrypoints/pth2json.py

@@ -38,6 +38,24 @@ def pth2json(
 
     json_dict = nnsk.to_json()
 
+    if kwargs.get("deleteoverlap", False):
+        log.info("The deleteoverlap option is set. The model will be converted to orthogonal version.")
+        # convert the model to orthogonal version
+        log.info("Converting the model to orthogonal version.")
+        # remove the overlap key in the json_dict
+        if "overlap" in json_dict['model_params']:
+            del json_dict['model_params']["overlap"]
+        json_dict['common_options']['overlap'] = False
+
+        if json_dict['model_options']['nnsk'].get('freeze',False) == ['overlap']:
+            json_dict['model_options']['nnsk']['freeze'] = False
+        elif isinstance(json_dict['model_options']['nnsk'].get('freeze',False), list):
+            json_dict['model_options']['nnsk']['freeze'] = [x for x in json_dict['model_options']['nnsk']['freeze'] if x != 'overlap']
+        
+        # turn off the push option in the ckpt json
+        if isinstance(json_dict['model_options']['nnsk'].get('push',False), dict):
+            json_dict['model_options']['nnsk']['push'] = False
+
     # dump the json file
     json_file = Path(outdir) / "ckpt.json"
     with open(json_file, "w") as f:

dptb/plugins/saver.py

@@ -63,11 +63,6 @@ def iteration(self, **kwargs):
         name = self.trainer.model.name+suffix
         self.latest_quene.append(name)
 
-        if len(self.latest_quene) > max_ckpt:
-            delete_name = self.latest_quene.pop(0)
-            delete_path = os.path.join(self.checkpoint_path, delete_name+".pth")
-            os.remove(delete_path)
-
         if len(self.latest_quene) > max_ckpt:
             delete_name = self.latest_quene.pop(0)
             delete_path = os.path.join(self.checkpoint_path, delete_name+".pth")

dptb/utils/argcheck.py

@@ -1679,7 +1679,10 @@ def collect_cutoffs(jdata):
     r_max, er_max, oer_max = get_cutoffs_from_model_options(model_options)
 
     if model_options.get("nnsk", None) is not None:
-        if model_options["nnsk"]["push"]:
+        if model_options["nnsk"]["push"] and \
+            abs(model_options["nnsk"]["push"]['rs_thr']) + \
+            abs(model_options["nnsk"]["push"]['rc_thr']) + \
+            abs(model_options["nnsk"]["push"]['w_thr']) > 1e-8:
             assert jdata.get("data_options",None) is not None, "data_options should be provided in jdata for nnsk push"
             assert jdata['data_options'].get("r_max") is not None, "r_max should be provided in data_options for nnsk push"
             log.info('YOU ARE USING NNSK PUSH MODEL, r_max will be used from data_options. Be careful! check the value in data options and model options. r_max or rs/rc !')