# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: tf23nlp
#     language: python
#     name: tf23nlp
# ---

# # !pip install tensorflow_datasets
import tensorflow as tf
import tensorflow_datasets as tfds
import numpy as np

# # Only if you have a GPU

tf.config.list_physical_devices('GPU')

# +
######## GPU CONFIGS FOR RTX 2070 ###############
## Please ignore if not training on GPU       ##
## this is important for running CuDNN on GPU ##

# check if GPU can be seen by TF
tf.config.list_physical_devices('GPU')
# enable this line if you wish to see whether GPU/CPU executes
# each command
#tf.debugging.set_log_device_placement(True)
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only use the first GPU
  try:
    tf.config.experimental.set_memory_growth(gpus[0], True)
    tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")
  except RuntimeError as e:
    # Visible devices must be set before GPUs have been initialized
    print(e)
###############################################
# -

# # Load Data

# see available tfds data sets
", ".join(tfds.list_builders())

# using TFDS dataset
# note: as_supervised converts dicts to tuples
imdb_train, ds_info = tfds.load(name="imdb_reviews", split="train", 
                                with_info=True, as_supervised=True)
imdb_test = tfds.load(name="imdb_reviews", split="test", 
                      as_supervised=True)

print(ds_info)

for example, label in imdb_train.take(1):
    print(example, '\n', label)

# +
# Use the default tokenizer settings
tokenizer = tfds.features.text.Tokenizer()

vocabulary_set = set()
MAX_TOKENS = 0

for example, label in imdb_train:
  some_tokens = tokenizer.tokenize(example.numpy())
  if MAX_TOKENS < len(some_tokens):
        MAX_TOKENS = len(some_tokens)
  vocabulary_set.update(some_tokens)

# +
imdb_encoder = tfds.features.text.TokenTextEncoder(vocabulary_set, 
                                                   tokenizer=tokenizer)
vocab_size = imdb_encoder.vocab_size

print(vocab_size, MAX_TOKENS)
# -

# Lets verify tokenization and encoding works
for example, label in imdb_train.take(1):
    print(example)
    encoded = imdb_encoder.encode(example.numpy())
    print(imdb_encoder.decode(encoded))

# +
# transformation fucntions to be used with the dataset
from tensorflow.keras.preprocessing import sequence

def encode_pad_transform(sample):
    encoded = imdb_encoder.encode(sample.numpy())
    pad = sequence.pad_sequences([encoded], padding='post', 
                                 maxlen=150)
    return np.array(pad[0], dtype=np.int64)  


def encode_tf_fn(sample, label):
    encoded = tf.py_function(encode_pad_transform, 
                                       inp=[sample], 
                                       Tout=(tf.int64))
    encoded.set_shape([None])
    label.set_shape([])
    return encoded, label


# -

# test the transformation on a small subset
subset = imdb_train.take(10)
tst = subset.map(encode_tf_fn)

for review, label in tst.take(1):
    print(review, label)
    print(imdb_encoder.decode(review))

#now tokenize/encode/pad all training
# and testing data
encoded_train = imdb_train.map(encode_tf_fn,
                               num_parallel_calls=tf.data.experimental.AUTOTUNE)
encoded_test = imdb_test.map(encode_tf_fn,
                             num_parallel_calls=tf.data.experimental.AUTOTUNE)

# # Preparing the model

# +
# Length of the vocabulary in chars
vocab_size = imdb_encoder.vocab_size # len(chars)

# The embedding dimension
embedding_dim = 64

# Number of RNN units
rnn_units = 64

#batch size
BATCH_SIZE=100


# -

def build_model_lstm(vocab_size, embedding_dim, rnn_units, batch_size):
  model = tf.keras.Sequential([
    tf.keras.layers.Embedding(vocab_size, embedding_dim, mask_zero=True,
                              batch_input_shape=[batch_size, None]),
    tf.keras.layers.LSTM(rnn_units),
    tf.keras.layers.Dense(1, activation='sigmoid')
  ])
  return model



# +
model = build_model_lstm(
  vocab_size = vocab_size,
  embedding_dim=embedding_dim,
  rnn_units=rnn_units,
  batch_size=BATCH_SIZE)

model.summary()
# -

model.compile(loss='binary_crossentropy', 
             optimizer='adam', 
             metrics=['accuracy', 'Precision', 'Recall'])

# # Model Training and Evaluation

# Prefetch for performance
encoded_train_batched = encoded_train.batch(BATCH_SIZE).prefetch(100)

model.fit(encoded_train_batched, epochs=10)

model.evaluate(encoded_test.batch(BATCH_SIZE))

# # BiLSTM Model

# +
# Length of the vocabulary in chars
vocab_size = imdb_encoder.vocab_size # len(chars)

# The embedding dimension
embedding_dim = 128

# Number of RNN units
rnn_units = 64

#batch size
BATCH_SIZE=50
# -

dropout=0.2
def build_model_bilstm(vocab_size, embedding_dim, rnn_units, batch_size):
  model = tf.keras.Sequential([
    tf.keras.layers.Embedding(vocab_size, embedding_dim, mask_zero=True,
                              batch_input_shape=[batch_size, None]),
    tf.keras.layers.Dropout(dropout),
    tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(rnn_units, return_sequences=True, dropout=0.25)),
    tf.keras.layers.Dropout(dropout),
    tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(rnn_units, dropout=0.25)),
    tf.keras.layers.Dropout(dropout),
    tf.keras.layers.Dense(1, activation='sigmoid')
  ])
  return model


# +
bilstm = build_model_bilstm(
  vocab_size = vocab_size,
  embedding_dim=embedding_dim,
  rnn_units=rnn_units,
  batch_size=BATCH_SIZE)

bilstm.summary()
# -

bilstm.compile(loss='binary_crossentropy', 
             optimizer='adam', 
             metrics=['accuracy', 'Precision', 'Recall'])

encoded_train_batched = encoded_train.batch(BATCH_SIZE).prefetch(100)

bilstm.fit(encoded_train_batched, epochs=5)

bilstm.evaluate(encoded_test.batch(BATCH_SIZE))



StarcoderdataJupyterScriptsDedupFiltered