I am doing some NLP exercises. I am unable to use a saved model with a hub.KerasLayer layer :(

Here are the steps I done:

1. Create the model - OK

2. Compile the model - OK

3. Fit/Train the model - OK

import tensorflow as tf
from tensorflow.keras import layers
import tensorflow_hub as hub

# Create a Keras layer using the USE pretrained layer from tensorflow hub
sentence_encoder_layer = hub.KerasLayer("https://www.kaggle.com/models/google/universal-sentence-encoder/TensorFlow2/universal-sentence-encoder/2",
                                        input_shape=[],
                                        dtype=tf.string,
                                        trainable=False,
                                        name="USE")

# Create model using the Sequential API
model_6 = tf.keras.Sequential([
  # sentence_encoder_layer, # take in sentences and then encode them into an embedding
  layers.Lambda(lambda x: sentence_encoder_layer(x)),
  layers.Dense(64, activation="relu"),
  layers.Dense(1, activation="sigmoid"),
], name="model_6_USE")


# Compile the model
model_6.compile(loss="binary_crossentropy",
                optimizer=tf.keras.optimizers.Adam(),
                metrics=["accuracy"])

# Train and fit the model
history_6 = model_6.fit(train_sentences,
                        train_labels,
                        epochs=5,
                        validation_data=(val_sentences, val_labels))

4. Make model predictions - OK

# Make model predictions
model_6_pred_prods = model_6.predict(val_sentences) 
model_6_pred_prods[:10]

5. Save the model - OK

# Save model_6 to native Keras format
model_6.save("model_6_saved.keras")

6. Load the model - OK

# Re-initialize the hub keraslayer
sentence_encoder_layer = hub.KerasLayer("https://www.kaggle.com/models/google/universal-sentence-encoder/TensorFlow2/universal-sentence-encoder/2",
                                        input_shape=[],
                                        dtype=tf.string,
                                        trainable=False,
                                        name="USE")

# Load the model
loaded_model_6 = tf.keras.models.load_model("model_6_saved.keras",
                                            custom_objects={'KerasLayer':hub.KerasLayer,
 'tf':tf,
 'sentence_encoder_layer': sentence_encoder_layer,
 },
                                            safe_mode=False)
                                                            
# Verify the loaded model
print(loaded_model_6.summary())

7. Evaluate the loaded model - FAIL.

# Evaluate the loaded_model
loaded_model_pred_probs = loaded_model_6.evaluate(val_sentences, val_labels)

See the error below.. it says that the hub layer was not defined eventhough I had re-initialized in the same code cell... It was also defined when I loaded the model...

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
/tmp/ipykernel_2313/2915160527.py in ()
      1 # Evaluate the loaded_model
----> 2 loaded_model_pred_probs = loaded_model_6.evaluate(val_sentences, val_labels)

1 frames
/usr/local/lib/python3.12/dist-packages/keras/src/utils/python_utils.py in (x)
     13 model_6 = tf.keras.Sequential([
     14   # sentence_encoder_layer, # take in sentences and then encode them into an embedding
---> 15   layers.Lambda(lambda x: sentence_encoder_layer(x)),
     16   layers.Dense(64, activation="relu"),
     17   layers.Dense(1, activation="sigmoid"),

NameError: Exception encountered when calling Lambda.call().

name 'sentence_encoder_layer' is not defined

Arguments received by Lambda.call():
  • inputs=tf.Tensor(shape=(None,), dtype=string)
  • mask=None
  • training=False