Upload ASTRA trained model

astra_model.py

@@ -34,26 +34,26 @@ class ASTRAViSoBERTForMaskedLM(XLMRobertaPreTrainedModel):
             remaps = []
             for key in list(state_dict.keys()):
                 new_key = key
-                if key.startswith(prefix + "cls.dense."):
-                    new_key = key.replace(prefix + "cls.dense.", prefix + "cls_dense.")
-                elif key.startswith(prefix + "cls.layer_norm."):
-                    new_key = key.replace(prefix + "cls.layer_norm.", prefix + "cls_layer_norm.")
-                elif key.startswith(prefix + "cls.decoder."):
-                    new_key = key.replace(prefix + "cls.decoder.", prefix + "cls_decoder.")
-                elif key.startswith(prefix + "lm_head.dense."):
-                    new_key = key.replace(prefix + "lm_head.dense.", prefix + "cls_dense.")
-                elif key.startswith(prefix + "lm_head.layer_norm."):
-                    new_key = key.replace(prefix + "lm_head.layer_norm.", prefix + "cls_layer_norm.")
-                elif key.startswith(prefix + "lm_head.decoder."):
-                    new_key = key.replace(prefix + "lm_head.decoder.", prefix + "cls_decoder.")
-                elif key.startswith(prefix + "lm_head."):
-                    new_key = key.replace(prefix + "lm_head.", prefix + "cls_decoder.")
+                # Replace occurrences anywhere in the key path after this module's prefix
+                if key.startswith(prefix):
+                    suffix = key[len(prefix):]
+                    suffix = suffix.replace("cls.dense.", "cls_dense.")
+                    suffix = suffix.replace("cls.layer_norm.", "cls_layer_norm.")
+                    suffix = suffix.replace("cls.decoder.", "cls_decoder.")
+                    suffix = suffix.replace("lm_head.dense.", "cls_dense.")
+                    suffix = suffix.replace("lm_head.layer_norm.", "cls_layer_norm.")
+                    suffix = suffix.replace("lm_head.decoder.", "cls_decoder.")
+                    # bare lm_head.* -> cls_decoder.*
+                    if suffix.startswith("lm_head."):
+                        suffix = suffix.replace("lm_head.", "cls_decoder.", 1)
+                    new_key = prefix + suffix
 
                 if new_key != key:
                     remaps.append((key, new_key))
 
             for old_key, new_key in remaps:
-                state_dict[new_key] = state_dict.pop(old_key)
+                if new_key not in state_dict:
+                    state_dict[new_key] = state_dict.pop(old_key)
 
         self.register_load_state_dict_pre_hook(_remap_pre_hook)
 

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ea42663a4a066b017945bb706b7058f7c80b8097936295df4680aba04763cc3b
+oid sha256:ec0edd3975d6ba755d4070a9b887ecdb234ffea69406d2be147ea7dd16a764af
 size 393222899