<jupyter_start><jupyter_text>This notebook is a part of the project [Generating-categories-from-arXiv-paper-titles](https://github.com/sayakpaul/Generating-categories-from-arXiv-paper-titles). It shows how to use Cloud TPUs to train `tf.keras` models. If you want to be able to do the same take a look at these [instructions](http://bit.ly/keras-tpu-tf21) prepared by [<NAME>](https://twitter.com/martin_gorner). Also, an insane amount of thanks to the [TensorFlow Research Cloud](https://www.tensorflow.org/tfrc) program for providing me the access to v3 Cloud TPUs.  Data gathering<jupyter_code># Imports
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import *
from tensorflow.keras.callbacks import *
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.utils import class_weight
from ast import literal_eval
from typing import Union
from utils import utils
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
import time
tf.__version__<jupyter_output><empty_output><jupyter_text>Data gathering and preprocessing<jupyter_code># Data gathering
def load_data(filename:str) -> np.ndarray:
    content = np.load(filename, allow_pickle=True)
    return content
X_train, y_train = load_data('data/X_train.npy'), load_data('data/y_train.npy')
X_test, y_test = load_data('data/X_test.npy'), load_data('data/y_test.npy')

X_train.shape, X_test.shape
clean_title = np.vectorize(utils.clean_title)
X_train = clean_title(X_train)
X_test = clean_title(X_test)
# Preview
X_train[:10]
tokenizer = Tokenizer(num_words=3000, lower=True)
tokenizer.fit_on_texts(X_train)
def get_features(text_sequence: np.ndarray) -> np.ndarray:
    sequences = tokenizer.texts_to_sequences(text_sequence)
    return pad_sequences(sequences, maxlen=300)

train_features = get_features(X_train)
test_features = get_features(X_test)
train_features.shape, test_features.shape
y_train[:10]
type(train_features), type(y_train)
# Label binarization
list_preprocessed = [literal_eval(i) for i in y_train]
mlb = MultiLabelBinarizer()
y_train_binarized = mlb.fit_transform(list_preprocessed)
mlb.classes_
y_train_binarized[:10]
y_train_binarized[0].shape
y_test_binarized = mlb.transform([literal_eval(i) for i in y_test])
y_test_binarized[:10]
class_weight = class_weight.compute_sample_weight('balanced', y_train)
class_weight<jupyter_output><empty_output><jupyter_text>Configure cloud TPUs<jupyter_code># Detect hardware, return appropriate distribution strategy
try:
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
    print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
except ValueError:
    tpu = None

if tpu:
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    strategy = tf.distribute.experimental.TPUStrategy(tpu)
else:
    strategy = tf.distribute.get_strategy()

print("REPLICAS: ", strategy.num_replicas_in_sync)<jupyter_output>Running on TPU  ['192.168.51.2:8470']
INFO:tensorflow:Initializing the TPU system: tpu-vm<jupyter_text>Preparation for model training<jupyter_code># Helper function to return a CNN-based model
def get_a_cnn_model() -> tf.keras.models.Sequential:
    model = Sequential()
    model.add(Embedding(3000, 30, 
        input_length=300))
    model.add(Dropout(0.1))
    model.add(Conv1D(300, 3, 
        padding='valid', activation='relu', strides=1))
    model.add(GlobalMaxPool1D())
    model.add(Dense(32, activation='sigmoid'))
    
    return model
def compile_model_w_strat(model:tf.keras.models.Sequential, function_name) -> tf.keras.models.Sequential:
    with strategy.scope():
        model = function_name()
        model.compile(optimizer='adam', 
                        loss='binary_crossentropy', 
                        metrics=['categorical_accuracy'])
        
        return model
compiled_model = compile_model_w_strat(model, get_a_cnn_model)
compiled_model.summary()<jupyter_output>Model: "sequential_5"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
embedding_5 (Embedding)      (None, 300, 30)           90000     
_________________________________________________________________
dropout_5 (Dropout)          (None, 300, 30)           0         
_________________________________________________________________
conv1d_5 (Conv1D)            (None, 298, 300)          27300     
_________________________________________________________________
global_max_pooling1d_5 (Glob (None, 300)               0         
_________________________________________________________________
dense_5 (Dense)              (None, 32)                9632      
=================================================================
Total params: 126,932
Trainable params: 126,932
Non-trainable params: 0
________________________________________________[...]<jupyter_text>And we are good to go<jupyter_code>compiled_model.fit(train_features, y_train_binarized,
                        class_weight=class_weight,
                        steps_per_epoch=len(train_features)//(8 * strategy.num_replicas_in_sync),
                        validation_data=(test_features, y_test_binarized),
                        validation_steps=len(test_features)//(8 * strategy.num_replicas_in_sync),
                        epochs=15,
                        batch_size=8 * strategy.num_replicas_in_sync)
def get_a_sequential_model():
    model = Sequential()
    model.add(Embedding(3000, 20, input_length=300))
    model.add(Bidirectional(LSTM(10, activation="sigmoid")))
    model.add(Dense(32, activation='sigmoid'))
    
    return model
compiled_model = compile_model_w_strat(model, get_a_sequential_model)
compiled_model.fit(train_features, y_train_binarized,
                        class_weight=class_weight,
                        steps_per_epoch=len(train_features)//(8 * strategy.num_replicas_in_sync),
                        validation_data=(test_features, y_test_binarized),
                        validation_steps=len(test_features)//(8 * strategy.num_replicas_in_sync),
                        epochs=15,
                        batch_size=8 * strategy.num_replicas_in_sync)<jupyter_output>Train on 26152 samples, validate on 6538 samples
Epoch 1/15
26112/26152 [============================>.] - ETA: 0s - loss: 0.2301 - categorical_accuracy: 0.2482 - val_loss: 0.1260 - val_categorical_accuracy: 0.3073Epoch 2/15
26088/26152 [============================>.] - ETA: 0s - loss: 0.1229 - categorical_accuracy: 0.3073 - val_loss: 0.1210 - val_categorical_accuracy: 0.3073Epoch 3/15
26088/26152 [============================>.] - ETA: 0s - loss: 0.1207 - categorical_accuracy: 0.3072 - val_loss: 0.1200 - val_categorical_accuracy: 0.3073Epoch 4/15
26088/26152 [============================>.] - ETA: 0s - loss: 0.1197 - categorical_accuracy: 0.3098 - val_loss: 0.1185 - val_categorical_accuracy: 0.3232Epoch 5/15
26088/26152 [============================>.] - ETA: 0s - loss: 0.1168 - categorical_accuracy: 0.3489 - val_loss: 0.1154 - val_categorical_accuracy: 0.3684Epoch 6/15
26088/26152 [============================>.] - ETA: 0s - loss: 0.1124 - categorical_accuracy: 0.3903 - val_loss: 0.[...]
{'fraction_non_alphanumeric': 0.13482238966630786, 'fraction_numerical': 0.045479009687836386, 'mean_word_length': 4.271631205673759, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 29, 'https://': 3, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '311987', 'license': 'MIT', 'n_tokens_mistral': 2567, 'n_tokens_neox': 2304, 'n_words': 447, 'path': 'TPU_Experimentation.ipynb', 'repo_name': 'sayakpaul/Generating-categories-from-arXiv-paper-titles', 'size': '7439'}