Load weights instead of full model when loading fails due to bad marshal data (model saved with different Python version)

2 files changed, 31 insertions, 18 deletions

diff --git a/bicleaner_ai/decomposable_attention.py b/bicleaner_ai/decomposable_attention.py
index 51f1c84..c312bd2 100644
--- a/bicleaner_ai/decomposable_attention.py
+++ b/bicleaner_ai/decomposable_attention.py
@@ -22,7 +22,7 @@ except (SystemError, ImportError):
     from metrics import MatthewsCorrCoef
     from layers import TokenAndPositionEmbedding
 
-def build_model(vectors, settings):
+def build_model(vectors, settings, compile=True):
     max_length = settings["maxlen"]
     nr_hidden = settings["n_hidden"]
     nr_class = settings["n_classes"]
@@ -31,7 +31,10 @@ def build_model(vectors, settings):
     input2 = layers.Input(shape=(max_length,), dtype="int32", name="words2")
 
     # embeddings (projected)
-    embed = create_embedding(vectors, max_length, nr_hidden, settings["emb_trainable"])
+    embed = create_embedding(vectors, settings["emb_dim"],
+                             settings["vocab_size"],
+                             max_length, nr_hidden,
+                             settings["emb_trainable"])
 
     a = embed(input1)
     b = embed(input2)
@@ -104,15 +107,16 @@ def build_model(vectors, settings):
 
     model = Model([input1, input2], out)
 
-    model.compile(optimizer=settings["optimizer"],
-                  loss=loss,
-                  metrics=settings["metrics"](), # Call get_metrics
-                  experimental_run_tf_function=False,)
+    if compile:
+        model.compile(optimizer=settings["optimizer"],
+                      loss=loss,
+                      metrics=settings["metrics"](), # Call get_metrics
+                      experimental_run_tf_function=False,)
 
     return model
 
 
-def create_embedding(vectors, max_length, projected_dim, trainable=False):
+def create_embedding(vectors, emb_dim, vocab_size, max_length, projected_dim, trainable=False):
     return models.Sequential(
         [
             # layers.Embedding(
@@ -123,8 +127,8 @@ def create_embedding(vectors, max_length, projected_dim, trainable=False):
             #     trainable=trainable,
             #     mask_zero=True,
             # ),
-            TokenAndPositionEmbedding(vectors.shape[0],
-                                      vectors.shape[1],
+            TokenAndPositionEmbedding(vocab_size,
+                                      emb_dim,
                                       max_length,
                                       vectors,
                                       trainable),
diff --git a/bicleaner_ai/models.py b/bicleaner_ai/models.py
index f6fa6ff..c9f136c 100644
--- a/bicleaner_ai/models.py
+++ b/bicleaner_ai/models.py
@@ -201,7 +201,7 @@ class BaseModel(ModelInterface):
         ''' Returns a sentence generator instance according to the model input '''
         raise NotImplementedError("Subclass must define its sentence generator")
 
-    def build_model(self):
+    def build_model(self, compile=True):
         '''Returns a compiled Keras model instance'''
         raise NotImplementedError("Subclass must implement its model architecture")
 
@@ -259,8 +259,16 @@ class BaseModel(ModelInterface):
                 'MatthewsCorrCoef': MatthewsCorrCoef,
                 'TokenAndPositionEmbedding': TokenAndPositionEmbedding,
         }
-        self.model = load_model(self.dir+'/'+self.settings["model_file"],
-                                custom_objects=deps, compile=False)
+
+        # Try loading the whole model
+        # If it fails due to bad marshal (saved with different Python version)
+        # build a new model and load weights
+        try:
+            self.model = load_model(self.dir+'/'+self.settings["model_file"],
+                                    custom_objects=deps, compile=False)
+        except ValueError:
+            self.model = self.build_model(compile=False)
+            self.model.load_weights(self.dir+'/'+self.settings["model_file"])
 
     def train_vocab(self, monolingual, threads):
         '''Trains SentencePiece model and embeddings with Glove'''
@@ -377,8 +385,8 @@ class DecomposableAttention(BaseModel):
                     batch_size=batch_size,
                     maxlen=self.settings["maxlen"])
 
-    def build_model(self):
-        return decomposable_attention.build_model(self.wv, self.settings)
+    def build_model(self, compile=True):
+        return decomposable_attention.build_model(self.wv, self.settings, compile)
 
 class Transformer(BaseModel):
     '''Basic Transformer model'''
@@ -418,7 +426,7 @@ class Transformer(BaseModel):
                     maxlen=self.settings["maxlen"],
                     separator=self.settings["separator"])
 
-    def build_model(self):
+    def build_model(self, compile=True):
         settings = self.settings
         inputs = layers.Input(shape=(settings["maxlen"],), dtype='int32')
         embedding = TokenAndPositionEmbedding(self.wv,
@@ -442,9 +450,10 @@ class Transformer(BaseModel):
             outputs = layers.Dense(settings["n_classes"], activation='sigmoid')(x)
 
         model = tf.keras.Model(inputs=inputs, outputs=outputs)
-        model.compile(optimizer=settings["optimizer"],
-                      loss=settings["loss"],
-                      metrics=settings["metrics"]())
+        if compile:
+            model.compile(optimizer=settings["optimizer"],
+                          loss=settings["loss"],
+                          metrics=settings["metrics"]())
         return model