Fix save and load of wrapper

keras_bert/layers/wrapper.py

@@ -3,8 +3,14 @@
 
 class Wrapper(keras.layers.Layer):
 
-    def __init__(self, **kwargs):
-        self.layers = {}
+    CONFIG_PREFIX = 'wrapper_layer_'
+
+    def __init__(self, layers=None, **kwargs):
+        if layers is None:
+            self.layers = {}
+        else:
+            self.layers = layers
+            self.built = True
         super(Wrapper, self).__init__(**kwargs)
 
     @property
@@ -20,3 +26,24 @@ def non_trainable_weights(self):
         for key in sorted(self.layers.keys()):
             weights += self.layers[key].non_trainable_weights
         return weights
+
+    def get_config(self):
+        config = {}
+        for name, layer in self.layers.items():
+            config[self.CONFIG_PREFIX + name] = {
+                'class_name': layer.__class__.__name__,
+                'config': layer.get_config(),
+            }
+        base_config = super(Wrapper, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        layers = {}
+        print(config)
+        keys = list(filter(lambda key: key.startswith(Wrapper.CONFIG_PREFIX), config.keys()))
+        for key in keys:
+            if key.startswith(Wrapper.CONFIG_PREFIX):
+                name = key[len(Wrapper.CONFIG_PREFIX):]
+                layers[name] = keras.layers.deserialize(config.pop(key), custom_objects=custom_objects)
+        return cls(layers=layers, **config)