Update Utils/PLBERT/util.py

Utils/PLBERT/util.py

@@ -3,7 +3,6 @@ import yaml
 import torch
 from transformers import AlbertConfig, AlbertModel
 
-
 class CustomAlbert(AlbertModel):
     def forward(self, *args, **kwargs):
         # Call the original forward method
@@ -16,34 +15,35 @@ class CustomAlbert(AlbertModel):
 def load_plbert(log_dir):
     config_path = os.path.join(log_dir, "config.yml")
     plbert_config = yaml.safe_load(open(config_path))
-
-    albert_base_configuration = AlbertConfig(**plbert_config["model_params"])
+    
+    albert_base_configuration = AlbertConfig(**plbert_config['model_params'])
     bert = CustomAlbert(albert_base_configuration)
 
     files = os.listdir(log_dir)
     ckpts = []
     for f in os.listdir(log_dir):
-        if f.startswith("step_"):
-            ckpts.append(f)
-
-    iters = [
-        int(f.split("_")[-1].split(".")[0])
-        for f in ckpts
-        if os.path.isfile(os.path.join(log_dir, f))
-    ]
+        if f.startswith("step_"): ckpts.append(f)
+
+    iters = [int(f.split('_')[-1].split('.')[0]) for f in ckpts if os.path.isfile(os.path.join(log_dir, f))]
     iters = sorted(iters)[-1]
 
-    checkpoint = torch.load(log_dir + "/step_" + str(iters) + ".t7", map_location="cpu")
-    state_dict = checkpoint["net"]
+    checkpoint = torch.load(log_dir + "/step_" + str(iters) + ".t7", map_location='cpu')
+    state_dict = checkpoint['net']
     from collections import OrderedDict
-
     new_state_dict = OrderedDict()
     for k, v in state_dict.items():
-        name = k[7:]  # remove `module.`
-        if name.startswith("encoder."):
-            name = name[8:]  # remove `encoder.`
+        name = k[7:] # remove `module.`
+        if name.startswith('encoder.'):
+            name = name[8:] # remove `encoder.`
             new_state_dict[name] = v
-    del new_state_dict["embeddings.position_ids"]
+    
+    # Check if 'embeddings.position_ids' exists before attempting to delete it
+    if not hasattr(bert.embeddings, 'position_ids'):
+        del new_state_dict["embeddings.position_ids"]
+    
+    
     bert.load_state_dict(new_state_dict, strict=False)
-
+    
     return bert
+
+