Using the a Universal Sentence Encoder Embedding Layer in Keras

Question

I am trying to load USE as an embedding layer in my model using Keras. I used two approaches. the first one is adapted from the code here as follows:

import tensorflow as tf
tf.config.experimental_run_functions_eagerly(True)

import tensorflow_hub as hub
from keras import backend as K
 
module_url = "../emb_models/use/universal-sentence-encoder-large-5"
embed = hub.load(module_url)

# For the keras Lambda
def UniversalEmbedding(x):
    results = embed(tf.squeeze(tf.cast(x, tf.string)))
    # results = embed(tf.squeeze(tf.cast(x, tf.string)))["outputs"] 
    # removed outputs as it gave an error "TypeError: Only integers, slices (`:`), ellipsis (`...`),
    # tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid indices, got 'outputs'"
    print(results)
    return K.concatenate([results])

# model
sentence_input = Input(shape=(1,), name='sentences', dtype="string")
sentence_embeds = Lambda(UniversalEmbedding, output_shape=(embed_size,))(sentence_input)

The model is successfully created but during the fitting (once it starts training) it gave the following error:

2020-12-01 10:45:12.307164: W tensorflow/core/framework/op_kernel.cc:1502] OP_REQUIRES failed at lookup_table_op.cc:809 : Failed precondition: Table not initialized.

The second approach is adapted from this issue as follows:

module_url = "../emb_models/use/universal-sentence-encoder-large-5"
use_embeddings_layer = hub.KerasLayer(module_url, trainable=False, dtype=tf.string)

inputs = tf.keras.layers.Input(shape=(None,), dtype='string'))
sentence_input = Input(shape=(1,), name="sentences", dtype="string") 
sentence_input = Lambda(lambda x: K.squeeze(x, axis=1), name='squeezed_input')(sentence_input)    
sentence_embed = use_embeddings_layer(sentence_input)

The model was not created and the following error occurred:

AttributeError: 'tuple' object has no attribute 'layer'

Any ideas?

Info: tensorflow-gpu == 1.14.0, keras==2.3.1, tensorflow-hub==0.8.0

Answer 1

This thread has a relevant example showing how to use hub.KerasLayer with USE. The example uses training=true but it should work with training=false (pure inference, no fine-tuning as well).

Also it might be good to try to use the latest version of TF (e.g. TF 2.5) in order to rule out any issue due to old TF version.