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

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1B3eC2inD2i4fcYcydkYwx9up6ixHZmai

#### Setup
"""

# A dependency of the preprocessing for BERT inputs
!pip install -q -U tensorflow-text
!pip install -q tf-models-official

import os
import shutil

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text as text
from official.nlp import optimization  # to create AdamW optimizer

import matplotlib.pyplot as plt

tf.get_logger().setLevel('ERROR')

"""#### Download the IMDB dataset"""

url = 'https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz'

imdb_dataset = tf.keras.utils.get_file(
    fname="imdb.tar.gz",
    origin=url,
    untar=True,
    cache_dir=".",
    cache_subdir="",
)

train_dir = "/content/aclImdb/train/"
test_dir = "/content/aclImdb/test/"

BATCH_SIZE = 64
MAX_LENGTH = 256


# remove additiona dir: unsup in train
shutil.rmtree(train_dir + "unsup/")

train_dataset = tf.keras.utils.text_dataset_from_directory(
    directory=train_dir, 
    labels='inferred', 
    label_mode='int',
    batch_size=BATCH_SIZE, 
    max_length=MAX_LENGTH, 
    shuffle=True, 
    seed=42,
    validation_split=0.2, 
    subset="training", 
) 

class_name = train_dataset.class_names

valid_dataset = tf.keras.utils.text_dataset_from_directory(
    directory=train_dir, 
    labels='inferred', 
    label_mode='int',
    batch_size=BATCH_SIZE, 
    max_length=MAX_LENGTH, 
    shuffle=True, 
    seed=42,
    validation_split=0.2, 
    subset="validation", 
) 

test_dataset = tf.keras.utils.text_dataset_from_directory(
    directory=test_dir, 
    labels='inferred', 
    label_mode='int',
    batch_size=BATCH_SIZE, 
    max_length=MAX_LENGTH, 
    shuffle=True, 
    seed=42,
) 


train_dataset = train_dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)
valid_dataset = valid_dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)
test_dataset = test_dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)

class_name

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

"""#### Loading models from TensorFlow Hub"""

bert_model_url = "https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-4_H-512_A-8/2"
bert_tokenizer_url = "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3"

bert_model = hub.KerasLayer(bert_model_url)
tokenizer = hub.KerasLayer(bert_tokenizer_url)

tokenizer

sample = tf.constant(["This movie is awesome. Love it !"])

processed_sample = tokenizer(sample)

print(f"Processed sample keys: {processed_sample.keys()}")

print(f"input_mask: {processed_sample['input_mask']}")
print(f"input_mask size: {processed_sample['input_mask'].shape}")
print(f"input_word_ids: {processed_sample['input_word_ids']}")
print(f"input_type_ids: {processed_sample['input_type_ids']}")

bert_output = bert_model(processed_sample)

print(f"Outputs: {bert_output.keys()}")
print(f"pooled_output: {bert_output['pooled_output'].shape}")
print(f"sequence_output: {bert_output['sequence_output'].shape}")

"""#### Define your model"""

def build_model():
    bert_model = hub.KerasLayer(bert_model_url, trainable=True, name="bert_model")
    tokenizer = hub.KerasLayer(bert_tokenizer_url, name='tokenizer')
    
    text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name="text")
    tokenized = tokenizer(text_input)
    outputs = bert_model(tokenized)
    pooled_output = outputs['pooled_output']
    output = tf.keras.layers.Dense(1, activation=None, name="classifier")(pooled_output)
    model = tf.keras.Model(text_input, output)

    return model

model = build_model()

model_output = model(sample)

print(tf.sigmoid(model_output))

tf.keras.utils.plot_model(model)

model.summary()

"""#### Model training"""

loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam(learning_rate=0.00003)
metrics = tf.keras.metrics.BinaryAccuracy()

model.compile(loss=loss, optimizer=optimizer, metrics=metrics)

history = model.fit(
    x=train_dataset,
    validation_data=valid_dataset,
    epochs=5,
)

"""#### Evaluate the model"""

model.evaluate(test_dataset)

history_dict = history.history
print(history_dict.keys())

acc = history_dict['binary_accuracy']
val_acc = history_dict['val_binary_accuracy']
loss = history_dict['loss']
val_loss = history_dict['val_loss']

epochs = range(1, len(acc) + 1)
fig = plt.figure(figsize=(10, 6))
fig.tight_layout()

plt.subplot(2, 1, 1)
# r is for "solid red line"
plt.plot(epochs, loss, 'r', label='Training loss')
# b is for "solid blue line"
plt.plot(epochs, val_loss, 'b', label='Validation loss')
plt.title('Training and validation loss')
# plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()

plt.subplot(2, 1, 2)
plt.plot(epochs, acc, 'r', label='Training acc')
plt.plot(epochs, val_acc, 'b', label='Validation acc')
plt.title('Training and validation accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend(loc='lower right')



"""#### Export for inference"""

dataset_name = 'imdb'
saved_model_path = './{}_bert'.format(dataset_name.replace('/', '_'))

model.save(saved_model_path, include_optimizer=False)

reloaded_model = tf.saved_model.load(saved_model_path)

def print_my_examples(inputs, results):
  result_for_printing = \
    [f'input: {inputs[i]:<30} : score: {results[i][0]:.6f}'
                         for i in range(len(inputs))]
  print(*result_for_printing, sep='\n')
  print()


examples = [
    'this is such an amazing movie!',  # this is the same sentence tried earlier
    'The movie was great!',
    'The movie was meh.',
    'The movie was okish.',
    'The movie was terrible...'
]

reloaded_results = tf.sigmoid(reloaded_model(tf.constant(examples)))
original_results = tf.sigmoid(model(tf.constant(examples)))

print('Results from the saved model:')
print_my_examples(examples, reloaded_results)
print('Results from the model in memory:')
print_my_examples(examples, original_results)

