import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_datasets as tfds


train_data, validation_data, test_data = tfds.load(name='imdb_reviews',
                                                   split=('train[:60%]', 'train[60%:]', 'test'),
                                                   as_supervised=True)

embedding = "https://tfhub.dev/google/tf2-preview/gnews-swivel-20dim/1"
hub_layer = hub.KerasLayer(embedding, input_shape=[],
                           dtype=tf.string, trainable=True)

model = tf.keras.Sequential()

model.add(hub_layer)
model.add(tf.keras.layers.Dense(16, activation='relu'))
model.add(tf.keras.layers.Dense(1))

model.summary()

ES = tf.keras.callbacks.EarlyStopping(monitor='val_accuracy', patience=5, restore_best_weights=True)

model.compile(optimizer='adam',
              loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
              metrics=['accuracy'])

history = model.fit(train_data.shuffle(10000).batch(512),
                    epochs=100,
                    validation_data=validation_data.batch(512),
                    callbacks=[ES])

results = model.evaluate(test_data.batch(512), verbose=2)

for name, value, in zip(model.metrics_names, results):
    print("%s: %.3f" % (name, value))

