Add workaround to fix keras loading

Author: CyberZHG

keras_bert/__init__.py

@@ -5,4 +5,4 @@
 from .util import *
 from .datasets import *
 
-__version__ = '0.87.0'
+__version__ = '0.88.0'

keras_bert/datasets/pretrained.py

@@ -5,6 +5,7 @@
 import shutil
 from collections import namedtuple
 from keras_bert.backend import keras
+from tensorflow.keras.utils import get_file
 
 __all__ = ['PretrainedInfo', 'PretrainedList', 'get_pretrained']
 
@@ -35,7 +36,7 @@ def get_pretrained(info):
     path = info
     if isinstance(info, PretrainedInfo):
         path = info.url
-    path = keras.utils.get_file(fname=os.path.split(path)[-1], origin=path, extract=True)
+    path = get_file(fname=os.path.split(path)[-1], origin=path, extract=True)
     base_part, file_part = os.path.split(path)
     file_part = file_part.split('.')[0]
     if isinstance(info, PretrainedInfo):

tests/optimizers/test_warmup.py

@@ -1,8 +1,10 @@
 import os
 import tempfile
 from unittest import TestCase
+
 import numpy as np
-from keras_bert.backend import keras, TF_KERAS
+
+from keras_bert.backend import keras
 from keras_bert import AdamWarmup
 
 
@@ -31,7 +33,10 @@ def _test_fit(self, optmizer):
 
         model_path = os.path.join(tempfile.gettempdir(), 'keras_warmup_%f.h5' % np.random.random())
         model.save(model_path)
-        model = keras.models.load_model(model_path, custom_objects={'AdamWarmup': AdamWarmup})
+
+        from tensorflow.python.keras.utils.generic_utils import CustomObjectScope
+        with CustomObjectScope({'AdamWarmup': AdamWarmup}):  # Workaround for incorrect global variable used in keras
+            model = keras.models.load_model(model_path, custom_objects={'AdamWarmup': AdamWarmup})
 
         results = model.predict(x).argmax(axis=-1)
         diff = np.sum(np.abs(y - results))
@@ -78,17 +83,10 @@ def test_fit_embed(self):
 
         x = np.random.randint(0, 5, (1024, 15))
         y = (x[:, 1] > 2).astype('int32')
-        model.fit(x, y, epochs=10)
+        model.fit(x, y, epochs=10, verbose=1)
 
         model_path = os.path.join(tempfile.gettempdir(), 'test_warmup_%f.h5' % np.random.random())
         model.save(model_path)
-        keras.models.load_model(model_path, custom_objects={'AdamWarmup': AdamWarmup})
-
-    def test_legacy(self):
-        opt = AdamWarmup(
-            decay_steps=10000,
-            warmup_steps=5000,
-            learning_rate=1e-3,
-        )
-        if not TF_KERAS:
-            opt.lr = opt.lr
+        from tensorflow.python.keras.utils.generic_utils import CustomObjectScope
+        with CustomObjectScope({'AdamWarmup': AdamWarmup}):  # Workaround for incorrect global variable used in keras
+            keras.models.load_model(model_path, custom_objects={'AdamWarmup': AdamWarmup})

tests/test_bert.py

@@ -18,10 +18,12 @@ def test_sample(self):
         )
         model_path = os.path.join(tempfile.gettempdir(), 'keras_bert_%f.h5' % np.random.random())
         model.save(model_path)
-        model = keras.models.load_model(
-            model_path,
-            custom_objects=get_custom_objects(),
-        )
+        from tensorflow.python.keras.utils.generic_utils import CustomObjectScope
+        with CustomObjectScope(get_custom_objects()):  # Workaround for incorrect global variable used in keras
+            model = keras.models.load_model(
+                model_path,
+                custom_objects=get_custom_objects(),
+            )
         model.summary(line_length=200)
 
     def test_task_embed(self):
@@ -38,10 +40,12 @@ def test_task_embed(self):
         model = keras.models.Model(inputs, outputs)
         model_path = os.path.join(tempfile.gettempdir(), 'keras_bert_%f.h5' % np.random.random())
         model.save(model_path)
-        model = keras.models.load_model(
-            model_path,
-            custom_objects=get_custom_objects(),
-        )
+        from tensorflow.python.keras.utils.generic_utils import CustomObjectScope
+        with CustomObjectScope(get_custom_objects()):  # Workaround for incorrect global variable used in keras
+            model = keras.models.load_model(
+                model_path,
+                custom_objects=get_custom_objects(),
+            )
         model.summary(line_length=200)
 
     def test_save_load_json(self):
@@ -82,10 +86,12 @@ def test_fit(self):
         token_list = list(token_dict.keys())
         if os.path.exists(model_path):
             steps_per_epoch = 10
-            model = keras.models.load_model(
-                model_path,
-                custom_objects=get_custom_objects(),
-            )
+            from tensorflow.python.keras.utils.generic_utils import CustomObjectScope
+            with CustomObjectScope(get_custom_objects()):  # Workaround for incorrect global variable used in keras
+                model = keras.models.load_model(
+                    model_path,
+                    custom_objects=get_custom_objects(),
+                )
         else:
             steps_per_epoch = 1000
             model = get_model(