# -*- coding: utf-8 -*-
"""IMDB Review Classification using RNN Tensorflow.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1ccrQr-6XAzNyW1WY0bkE0dIjBh4kPaTL
"""

import numpy as np
import tensorflow_datasets as tfds
import tensorflow as tf

tfds.disable_progress_bar()

import matplotlib.pyplot as plt


def plot_graphs(history, metric):
  plt.plot(history.history[metric])
  plt.plot(history.history['val_'+metric], '')
  plt.xlabel("Epochs")
  plt.ylabel(metric)
  plt.legend([metric, 'val_'+metric])

"""#### Setup input pipeline"""

dataset, info = tfds.load("imdb_reviews", with_info=True, as_supervised=True)
dataset

train_dataset, test_dataset = dataset['train'], dataset['test']

for text, label in train_dataset.take(1):
    print(f"{text.numpy()}")
    print(f"{label.numpy()}")

BUFFER_SIZE = 10000
BATCH_SIZE = 64

train_dataset = train_dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
test_dataset = test_dataset.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

"""#### Create text encoder"""

VOCAB_SIZE = 10000

encoder = tf.keras.layers.TextVectorization(max_tokens=VOCAB_SIZE)

encoder.adapt(train_dataset.map(lambda text, label: text))

vocab = np.array(encoder.get_vocabulary())
vocab[:10]

tokenised_example = encoder("This is a good movie.")
tokenised_example

"""#### Create the model"""

def build_model():
    model = tf.keras.Sequential([
        encoder,
        tf.keras.layers.Embedding(input_dim=encoder.vocabulary_size(), output_dim=64, mask_zero=True),
        tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64)),
        tf.keras.layers.Dense(64, activation="relu"),
        tf.keras.layers.Dense(1),
    ])

    return model

model = build_model()

for layer in model.layers:
    print(f"{layer.supports_masking}")

# predict on a sample text without padding.

sample_text = ('The movie was cool. The animation and the graphics '
               'were out of this world. I would recommend this movie.')
predictions = model.predict(np.array([sample_text]))
print(predictions[0])

# predict on a sample text with padding

padding = "the " * 2000
predictions = model.predict(np.array([sample_text, padding]))
print(predictions[0])

model.compile(
    loss = tf.keras.losses.BinaryCrossentropy(from_logits=True),
    optimizer = tf.keras.optimizers.Adam(1e-4),
    metrics = ["accuracy"]
)

"""#### Train the model"""

history = model.fit(
    train_dataset,
    validation_data = test_dataset,
    validation_steps=30,
    epochs=10
)

test_loss, test_acc = model.evaluate(test_dataset)

print('Test Loss:', test_loss)
print('Test Accuracy:', test_acc)

plt.figure(figsize=(16, 8))
plt.subplot(1, 2, 1)
plot_graphs(history, 'accuracy')
plt.ylim(None, 1)
plt.subplot(1, 2, 2)
plot_graphs(history, 'loss')
plt.ylim(0, None)

sample_text = ('The movie was cool. The animation and the graphics '
               'were out of this world. I would recommend this movie.')
predictions = model.predict(np.array([sample_text]))

predictions

"""#### Reference

1. [Text classification with an RNN](https://www.tensorflow.org/text/tutorials/text_classification_rnn)
"""