convertion script WIP

convert_tf_checkpoint.py

@@ -10,7 +10,7 @@ import argparse
 import tensorflow as tf
 import torch
 
-from .modeling_pytorch import BertConfig, BertModel
+from modeling_pytorch import BertConfig, BertModel
 
 parser = argparse.ArgumentParser()
 
@@ -35,6 +35,10 @@ parser.add_argument("--pytorch_dump_path",
 args = parser.parse_args()
 
 def convert():
+    # Initialise PyTorch model
+    config = BertConfig.from_json_file(args.bert_config_file)
+    model = BertModel(config)
+
     # Load weights from TF model
     path = args.tf_checkpoint_path
     print("Converting TensorFlow checkpoint from {}".format(path))
@@ -49,24 +53,26 @@ def convert():
         names.append(name)
         arrays.append(array)
 
-    # Initialise PyTorch model and fill weights-in
-    config = BertConfig.from_json_file(args.bert_config_file)
-    model = BertModel(config)
     for name, array in zip(names, arrays):
         name = name[5:]  # skip "bert/"
-        assert name[-2:] == ":0"
-        name = name[:-2]
         name = name.split('/')
         pointer = model
         for m_name in name:
-            if re.fullmatch(r'[A-Za-z]+\d+', m_name):
-                l = re.split(r'(\d+)', m_name)
+            if re.fullmatch(r'[A-Za-z]+_\d+', m_name):
+                l = re.split(r'_(\d+)', m_name)
             else:
                 l = [m_name]
-            pointer = getattr(pointer, l[0])
+            if l[0] == 'kernel':
+                pointer = getattr(pointer, 'weight')
+            else:
+                pointer = getattr(pointer, l[0])
             if len(l) >= 2:
                 num = int(l[1])
                 pointer = pointer[num]
+        if m_name[-11:] == '_embeddings':
+            pointer = getattr(pointer, 'weight')
+        # elif m_name == 'kernel':
+        #     pointer = getattr(pointer, 'weight')
         try:
             assert pointer.shape == array.shape
         except AssertionError as e:
@@ -79,4 +85,3 @@ def convert():
 
 if __name__ == "__main__":
     convert()
-    return None