convert_to_keras.py

# To convert the model to Keras, simply run:
keras_layer = trax.AsKeras(sentiment_analysis_model)
# This will be a trax.trax2keras.AsKeras object
print(keras_layer)

# Run the Keras layer to verify it returns the same result.
example_input = list(data.tokenize(iter(["I loved the way that the actors were cast, also, It is clear that they've put a huge effort in post-production."]), vocab_file="en_8k.subword"))[0]
sentiment_activations = keras_layer(example_input[None, :])
print(f'Keras returned sentiment activations: {numpy.asarray(sentiment_activations)}')