KerasLayer with ELMo

[tokestermw]

Hi, is it possible to load the ELMo module using tensorflow hub and return the full ELMo embeddings? I would like to extract the "elmo" embeddings which returns contextual embeddings for each token. I need to make as_dict=True, but KerasLayer doesn't seem to allow it. I am also trying to avoid using Lambda layers because of issues with serialization.

Code to reproduce

tokens = layers.Input(shape=[None], dtype=tf.string)
seq_lens = layers.Input(shape=[], dtype=tf.int32)

elmo = hub.KerasLayer(
    "https://tfhub.dev/google/elmo/3",
    trainable=True,
    signature="tokens",
)

out = elmo({"tokens": tokens, "sequence_len": seq_lens})
model = keras.Model(inputs=[tokens, seq_lens], outputs=out)
model.compile("adam", loss="sparse_categorical_crossentropy")
model.summary()

Model summary

__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_2 (InputLayer)            [(None,)]            0                                            
__________________________________________________________________________________________________
input_1 (InputLayer)            [(None, None)]       0                                            
__________________________________________________________________________________________________
keras_layer (KerasLayer)        (None, 1024)         93600852    input_2[0][0]                    
                                                                 input_1[0][0]                    
==================================================================================================
Total params: 93,600,852
Trainable params: 0
Non-trainable params: 93,600,852

Package versions

tensorflow==1.15.0
tensorflow-estimator==1.15.1
tensorflow-hub==0.7.0

trainable=True doesn't work.

Separately, trainable=True doesn't seem to do much. I would expect 4 weights to be trainable, but none of them are.

Thanks

[tokestermw]

Ah, there is an output_key option 😅

elmo = hub.KerasLayer("https://tfhub.dev/google/elmo/3",
                      trainable=True,
                      signature="tokens",
                      output_key="elmo")

[Zachary4biz]

@tokestermw according to the official page: https://www.tensorflow.org/hub/api_docs/python/hub/KerasLayer I didn't see any information about the initial params you used (signature, output_key ..). May you tell me where did you find out this use case?

[tokestermw]

@Zachary4biz In the source code for KerasLayer: https://github.com/tensorflow/hub/blob/master/tensorflow_hub/keras_layer.py#L117