{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"\n# Predict Book Review Rating using Jax/Flax ","metadata":{}},{"cell_type":"markdown","source":"# Introduction\n\nThis notebook is taking [goodreads competition dataset](https://www.kaggle.com/c/goodreads-books-reviews-290312/overview/description) consisting reviews from the Goodreads book review website, and a variety of attributes describing the items. and predicting review rating which ranges from 0 to 5 using [flax framework](https://flax.readthedocs.io/en/latest/overview.html#flax).\n\n# About Dataset\n\nThe dataset contains more than **1.3M book reviews** about **25,475 books** and **18,892 users** , which is a review subset for spoiler detection, where each book/user has at least one associated spoiler review.\n\n**Files** - \n* goodreads_train.csv - the training set\n* goodreads_test.csv - the test set\n* goodreads_sample_submission.csv - sample submission file\n\nHere are the columns of the dataset -\n\n* user_id - Id of user\n* book_id - Id of Book\n* review_id - Id of review\n* rating - rating from 0 to 5\n* review_text - review text\n* date_added - date added\n* date_updated - date updated\n* read_at - read at\n* started_at - started at\n* n_votes - no. of votes\n* n_comments - no. of comments\n\nNote - I also included tokenized training and testing data which can be use directly to save time.\n\nyou can access it here - https://www.kaggle.com/yashvi/goodreads-tokenized-data\n\n# How to use this notebook \nYou can **copy & edit** this notebook , which you can find on **top-right** corner of kaggle notebook viewer.\n","metadata":{}},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"# Import Libraries\nimport os\nimport jax\nimport optax\nimport flax\nimport warnings\nimport datasets\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nfrom wordcloud import WordCloud, STOPWORDS\nfrom jax import jit\nimport jax.numpy as jnp\nfrom itertools import chain\nfrom tqdm.notebook import tqdm\nfrom typing import Callable\nfrom flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key\nfrom flax.training import train_state\nfrom flax import traverse_util\nfrom datasets import load_dataset, load_metric ,Dataset,list_metrics,load_from_disk\nfrom transformers import AutoTokenizer,BertTokenizer,FlaxAutoModelForSequenceClassification,AutoConfig\nfrom sklearn.model_selection import train_test_split\n\nwarnings.filterwarnings(\"ignore\")\nkey = jax.random.PRNGKey(0)\n","metadata":{"_kg_hide-input":true,"_kg_hide-output":false,"execution":{"iopub.status.busy":"2022-01-21T14:08:18.638220Z","iopub.execute_input":"2022-01-21T14:08:18.639228Z","iopub.status.idle":"2022-01-21T14:08:18.651238Z","shell.execute_reply.started":"2022-01-21T14:08:18.639190Z","shell.execute_reply":"2022-01-21T14:08:18.650332Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('JAX is running on', jax.lib.xla_bridge.get_backend().platform)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:08:18.653323Z","iopub.execute_input":"2022-01-21T14:08:18.653875Z","iopub.status.idle":"2022-01-21T14:08:18.662073Z","shell.execute_reply.started":"2022-01-21T14:08:18.653835Z","shell.execute_reply":"2022-01-21T14:08:18.661164Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Set DEBUG_MODE = True while experimenting \nDEBUG_MODE=False","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:08:18.663405Z","iopub.execute_input":"2022-01-21T14:08:18.663660Z","iopub.status.idle":"2022-01-21T14:08:18.669709Z","shell.execute_reply.started":"2022-01-21T14:08:18.663626Z","shell.execute_reply":"2022-01-21T14:08:18.668912Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#some configurations \nif 'TPU_NAME' in os.environ:\n    import requests\n    if 'TPU_DRIVER_MODE' not in globals():\n        url = 'http:' + os.environ['TPU_NAME'].split(':')[1] + ':8475/requestversion/tpu_driver_nightly'\n        resp = requests.post(url)\n        TPU_DRIVER_MODE = 1\n    from jax.config import config\n    config.FLAGS.jax_xla_backend = \"tpu_driver\"\n    config.FLAGS.jax_backend_target = os.environ['TPU_NAME']\n    print('Registered TPU:', config.FLAGS.jax_backend_target)\nelse:\n    print('No TPU detected. Can be changed under \"Runtime/Change runtime type\".')","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:08:18.671980Z","iopub.execute_input":"2022-01-21T14:08:18.672256Z","iopub.status.idle":"2022-01-21T14:08:18.679218Z","shell.execute_reply.started":"2022-01-21T14:08:18.672220Z","shell.execute_reply":"2022-01-21T14:08:18.678479Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# check jax local devices \njax.local_devices()","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:08:18.680505Z","iopub.execute_input":"2022-01-21T14:08:18.680915Z","iopub.status.idle":"2022-01-21T14:08:18.688981Z","shell.execute_reply.started":"2022-01-21T14:08:18.680876Z","shell.execute_reply":"2022-01-21T14:08:18.688091Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# load dataset\ndata = load_dataset(\"csv\", data_files={'train': ['../input/goodreads-books-reviews-290312/goodreads_train.csv']})\nraw_test = load_dataset('csv', data_files={'test': ['../input/goodreads-books-reviews-290312/goodreads_test.csv']})\n\n#setting sample size\nsample_size=15000\nif DEBUG_MODE==True:\n    data=data[\"train\"][:sample_size]\n    data = Dataset.from_dict(data)\n    data = data.train_test_split(0.20)\nelse:  \n    data = data[\"train\"].train_test_split(0.20)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:08:18.690346Z","iopub.execute_input":"2022-01-21T14:08:18.690816Z","iopub.status.idle":"2022-01-21T14:08:19.304854Z","shell.execute_reply.started":"2022-01-21T14:08:18.690778Z","shell.execute_reply":"2022-01-21T14:08:19.304168Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"raw_test","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-01-21T14:08:19.306296Z","iopub.execute_input":"2022-01-21T14:08:19.306679Z","iopub.status.idle":"2022-01-21T14:08:19.311880Z","shell.execute_reply.started":"2022-01-21T14:08:19.306643Z","shell.execute_reply":"2022-01-21T14:08:19.311237Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data","metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-01-21T14:08:19.313101Z","iopub.execute_input":"2022-01-21T14:08:19.313550Z","iopub.status.idle":"2022-01-21T14:08:19.322828Z","shell.execute_reply.started":"2022-01-21T14:08:19.313512Z","shell.execute_reply":"2022-01-21T14:08:19.322145Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Exploratory data analysis ","metadata":{}},{"cell_type":"code","source":"# reading for analysis\ntrain=pd.read_csv(\"../input/goodreads-books-reviews-290312/goodreads_train.csv\")","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-01-21T14:08:19.325569Z","iopub.execute_input":"2022-01-21T14:08:19.325768Z","iopub.status.idle":"2022-01-21T14:08:30.291162Z","shell.execute_reply.started":"2022-01-21T14:08:19.325736Z","shell.execute_reply":"2022-01-21T14:08:30.290410Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train.head(2)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-01-21T14:08:30.294948Z","iopub.execute_input":"2022-01-21T14:08:30.295156Z","iopub.status.idle":"2022-01-21T14:08:30.307559Z","shell.execute_reply.started":"2022-01-21T14:08:30.295126Z","shell.execute_reply":"2022-01-21T14:08:30.306861Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Target distribution","metadata":{}},{"cell_type":"code","source":"plt.figure(figsize=(8, 8))\nax = sns.countplot(x=\"rating\", data=train,palette=\"hls\")\nplt.title(\"Target(rating) distribution\", fontsize=20)\nplt.xticks(rotation=90, fontsize=15)\nplt.yticks(fontsize=13)\nplt.ylabel(\"Count\", fontsize=18)\nplt.xlabel(\"Rating\",fontsize=18)\nax.xaxis.label.set_color('red')\nax.yaxis.label.set_color('red')","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-01-21T14:08:30.308918Z","iopub.execute_input":"2022-01-21T14:08:30.309411Z","iopub.status.idle":"2022-01-21T14:08:30.556250Z","shell.execute_reply.started":"2022-01-21T14:08:30.309360Z","shell.execute_reply":"2022-01-21T14:08:30.555546Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"> Majority of rating is from 3 to 5 ","metadata":{}},{"cell_type":"markdown","source":"### Word clouds for different Rating score.\n","metadata":{}},{"cell_type":"code","source":"for rating_score in range(0,6):\n    reviews=train[train[\"rating\"]==rating_score][\"review_text\"]\n    wordcloud = WordCloud(width=400, height=330, max_words=150,colormap=\"Dark2\").generate(reviews.values[0])  \n    plt.figure(figsize=(10,8))\n    plt.imshow(wordcloud, interpolation='bilinear')\n    plt.title(\"Reviews with score {}\".format(str(rating_score)),fontsize=15)\n    plt.axis(\"off\")\n    plt.tight_layout()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-01-21T14:08:30.557563Z","iopub.execute_input":"2022-01-21T14:08:30.557959Z","iopub.status.idle":"2022-01-21T14:08:34.394829Z","shell.execute_reply.started":"2022-01-21T14:08:30.557920Z","shell.execute_reply":"2022-01-21T14:08:34.394195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Preprocessing\n\n**bert-base-cased** - Pretrained model on English language using a masked language modeling (MLM) objective.\n\nlearn more about it here - https://huggingface.co/bert-base-cased","metadata":{}},{"cell_type":"code","source":"model_checkpoint = \"bert-base-cased\" \ntokenizer = BertTokenizer.from_pretrained(model_checkpoint,use_fast=True)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:08:34.396072Z","iopub.execute_input":"2022-01-21T14:08:34.398545Z","iopub.status.idle":"2022-01-21T14:08:35.336055Z","shell.execute_reply.started":"2022-01-21T14:08:34.398502Z","shell.execute_reply":"2022-01-21T14:08:35.335324Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def preprocess_function(input_batch):\n    '''\n    This function will take batch of data and returns the tokenized processed data\n    INPUT - input batch from from original dataset\n    RETURNS preprocessed data\n    '''\n    texts = (input_batch[\"review_text\"],)\n    processed = tokenizer(*texts, padding=\"max_length\", max_length=128, truncation=True)\n    processed[\"labels\"] = input_batch[\"rating\"]\n    return processed","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:08:35.337245Z","iopub.execute_input":"2022-01-21T14:08:35.337504Z","iopub.status.idle":"2022-01-21T14:08:35.341914Z","shell.execute_reply.started":"2022-01-21T14:08:35.337469Z","shell.execute_reply":"2022-01-21T14:08:35.341243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Here I already tokenized data and created tokenized dataset , so that i don't have to tokenize it again everytime.\n\nlink to the tokenized dataset - https://www.kaggle.com/yashvi/goodreads-tokenized-data ","metadata":{}},{"cell_type":"code","source":"if DEBUG_MODE==True:\n    tokenized_dataset = data.map(preprocess_function, batched=True, remove_columns=data[\"train\"].column_names)\nelse:\n    if not os.path.exists(\"../input/goodreads-tokenized-data/train_tokenized_dataset\"):\n        tokenized_dataset = data.map(preprocess_function, batched=True, remove_columns=data[\"train\"].column_names)\n    else:\n        print(\"training dataset is already tokenized\")\n        tokenized_dataset=load_from_disk(\"../input/goodreads-tokenized-data/train_tokenized_dataset\")","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:08:35.343375Z","iopub.execute_input":"2022-01-21T14:08:35.343938Z","iopub.status.idle":"2022-01-21T14:09:49.249352Z","shell.execute_reply.started":"2022-01-21T14:08:35.343899Z","shell.execute_reply":"2022-01-21T14:09:49.248665Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = tokenized_dataset[\"train\"]\nvalidation_dataset = tokenized_dataset[\"test\"]","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:49.250502Z","iopub.execute_input":"2022-01-21T14:09:49.250982Z","iopub.status.idle":"2022-01-21T14:09:49.271410Z","shell.execute_reply.started":"2022-01-21T14:09:49.250943Z","shell.execute_reply":"2022-01-21T14:09:49.270683Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Generation","metadata":{}},{"cell_type":"code","source":"metrics_list = list_metrics()\nmetrics_list","metadata":{"_kg_hide-output":true,"_kg_hide-input":false,"execution":{"iopub.status.busy":"2022-01-21T14:09:49.273072Z","iopub.execute_input":"2022-01-21T14:09:49.273322Z","iopub.status.idle":"2022-01-21T14:09:49.603564Z","shell.execute_reply.started":"2022-01-21T14:09:49.273289Z","shell.execute_reply":"2022-01-21T14:09:49.602901Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"I am using 'F1-score' metric as given in [Competition evaluation section](https://www.kaggle.com/c/goodreads-books-reviews-290312/overview/evaluation) ","metadata":{}},{"cell_type":"code","source":"metric = load_metric('f1')\nmetric","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-01-21T14:09:49.604736Z","iopub.execute_input":"2022-01-21T14:09:49.604969Z","iopub.status.idle":"2022-01-21T14:09:49.818730Z","shell.execute_reply.started":"2022-01-21T14:09:49.604934Z","shell.execute_reply":"2022-01-21T14:09:49.818015Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Parameters \nnum_labels = 6 # 0-5\nseed = 0\nnum_train_epochs = 3\nlearning_rate = 2e-5\nper_device_batch_size = 32\nweight_decay=1e-2","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:49.820172Z","iopub.execute_input":"2022-01-21T14:09:49.820441Z","iopub.status.idle":"2022-01-21T14:09:49.825193Z","shell.execute_reply.started":"2022-01-21T14:09:49.820406Z","shell.execute_reply":"2022-01-21T14:09:49.824439Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"total_batch_size = per_device_batch_size * jax.local_device_count()\nprint(\"The overall batch size (both for training and eval) is\", total_batch_size)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:49.826781Z","iopub.execute_input":"2022-01-21T14:09:49.827073Z","iopub.status.idle":"2022-01-21T14:09:49.837081Z","shell.execute_reply.started":"2022-01-21T14:09:49.827034Z","shell.execute_reply":"2022-01-21T14:09:49.836302Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Below step is downloading the pretrained model,as we are doing **sentence classification** we can use **FlaxAutoModelForSequenceClassification** class.\n**from_pretrained** method will download and cache the model.","metadata":{}},{"cell_type":"code","source":"config = AutoConfig.from_pretrained(model_checkpoint, num_labels=num_labels)\nmodel = FlaxAutoModelForSequenceClassification.from_pretrained(model_checkpoint, config=config, seed=seed)","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-01-21T14:09:49.838449Z","iopub.execute_input":"2022-01-21T14:09:49.839288Z","iopub.status.idle":"2022-01-21T14:09:51.983981Z","shell.execute_reply.started":"2022-01-21T14:09:49.839249Z","shell.execute_reply":"2022-01-21T14:09:51.983202Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Evaluating num of train steps and defining learning_rate_function using [optax](https://optax.readthedocs.io/en/latest/) for jax. \n\nHere I am using **[cosine onecycle learning rate scheduler](https://optax.readthedocs.io/en/latest/api.html#optax.cosine_onecycle_schedule)** from optax library ","metadata":{}},{"cell_type":"code","source":"num_train_steps = len(train_dataset) // total_batch_size * num_train_epochs\nlearning_rate_function = optax.cosine_onecycle_schedule(transition_steps=num_train_steps, peak_value=learning_rate, pct_start=0.1)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:51.985196Z","iopub.execute_input":"2022-01-21T14:09:51.986537Z","iopub.status.idle":"2022-01-21T14:09:51.998963Z","shell.execute_reply.started":"2022-01-21T14:09:51.986493Z","shell.execute_reply":"2022-01-21T14:09:51.998085Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Flax provides a class [flax.training.train_state.TrainState](https://flax.readthedocs.io/en/latest/flax.training.html#train-state), which stores the model parameters, the loss function, the optimizer, and exposes an **apply_gradients** function to update the model's weight parameters.","metadata":{}},{"cell_type":"code","source":"class TrainState(train_state.TrainState):\n    logits_function: Callable = flax.struct.field(pytree_node=False)\n    loss_function: Callable = flax.struct.field(pytree_node=False)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.000230Z","iopub.execute_input":"2022-01-21T14:09:52.000786Z","iopub.status.idle":"2022-01-21T14:09:52.006484Z","shell.execute_reply.started":"2022-01-21T14:09:52.000748Z","shell.execute_reply":"2022-01-21T14:09:52.005441Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Here we are using Adam optimizer with weight decay, and again we are using [optax](https://optax.readthedocs.io/en/latest/)  library.\n\nHere is the interesting article on adam optimizer - https://www.fast.ai/2018/07/02/adam-weight-decay/","metadata":{}},{"cell_type":"code","source":"def decay_mask_fn(params):\n    flat_params = traverse_util.flatten_dict(params)\n    flat_mask = {path: (path[-1] != \"bias\" and path[-2:] != (\"LayerNorm\", \"scale\")) for path in flat_params}\n    return traverse_util.unflatten_dict(flat_mask)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.008283Z","iopub.execute_input":"2022-01-21T14:09:52.008623Z","iopub.status.idle":"2022-01-21T14:09:52.015663Z","shell.execute_reply.started":"2022-01-21T14:09:52.008590Z","shell.execute_reply":"2022-01-21T14:09:52.014947Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def adamw(weight_decay):\n    return optax.adamw(learning_rate=learning_rate_function, b1=0.9, b2=0.999, eps=1e-6, weight_decay=weight_decay,mask=decay_mask_fn)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.020636Z","iopub.execute_input":"2022-01-21T14:09:52.020946Z","iopub.status.idle":"2022-01-21T14:09:52.026688Z","shell.execute_reply.started":"2022-01-21T14:09:52.020916Z","shell.execute_reply":"2022-01-21T14:09:52.025943Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"adamw = adamw(weight_decay)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.027961Z","iopub.execute_input":"2022-01-21T14:09:52.028440Z","iopub.status.idle":"2022-01-21T14:09:52.034518Z","shell.execute_reply.started":"2022-01-21T14:09:52.028404Z","shell.execute_reply":"2022-01-21T14:09:52.033784Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now , computing the **softmax cross entropy** between sets of logits and labels using the [optax](https://optax.readthedocs.io/en/latest/) library.","metadata":{}},{"cell_type":"code","source":"@jit\ndef loss_function(logits, labels):\n    xentropy = optax.softmax_cross_entropy(logits, onehot(labels, num_classes=num_labels))\n    return jnp.mean(xentropy)\n \n@jit\ndef eval_function(logits):\n    return logits.argmax(-1)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.035804Z","iopub.execute_input":"2022-01-21T14:09:52.036108Z","iopub.status.idle":"2022-01-21T14:09:52.044116Z","shell.execute_reply.started":"2022-01-21T14:09:52.036074Z","shell.execute_reply":"2022-01-21T14:09:52.043230Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Instantiate a TrainState.\nstate = TrainState.create(\n    apply_fn=model.__call__,\n    params=model.params,\n    tx=adamw,\n    logits_function=eval_function,\n    loss_function=loss_function,\n)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.045197Z","iopub.execute_input":"2022-01-21T14:09:52.046120Z","iopub.status.idle":"2022-01-21T14:09:52.305001Z","shell.execute_reply.started":"2022-01-21T14:09:52.046084Z","shell.execute_reply":"2022-01-21T14:09:52.304265Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Define train step \ndef train_step(state, batch, dropout_rng):\n    # take targets\n    targets = batch.pop(\"labels\")\n    dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)\n    \n    #define loss function which runs the forward pass \n    def loss_function(params):\n        logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]\n        loss = state.loss_function(logits, targets)\n        return loss\n    \n    \n    grad_fn = jax.value_and_grad(loss_function) #differentiate the loss function\n    loss, grad = grad_fn(state.params) \n    grad = jax.lax.pmean(grad, \"batch\") #compute the mean gradient over all devices \n    new_state = state.apply_gradients(grads=grad) #applies the gradients to the weights.\n    metrics = jax.lax.pmean({'loss': loss, 'learning_rate': learning_rate_function(state.step)}, axis_name='batch')\n    \n    return new_state, metrics, new_dropout_rng","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.306455Z","iopub.execute_input":"2022-01-21T14:09:52.306708Z","iopub.status.idle":"2022-01-21T14:09:52.313784Z","shell.execute_reply.started":"2022-01-21T14:09:52.306675Z","shell.execute_reply":"2022-01-21T14:09:52.313095Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"parallel_train_step = jax.pmap(train_step, axis_name=\"batch\", donate_argnums=(0,)) # parallelized training over all TPU devices","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.315340Z","iopub.execute_input":"2022-01-21T14:09:52.315829Z","iopub.status.idle":"2022-01-21T14:09:52.322789Z","shell.execute_reply.started":"2022-01-21T14:09:52.315790Z","shell.execute_reply":"2022-01-21T14:09:52.322001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Define evaluation step\ndef eval_step(state, batch):\n    logits = state.apply_fn(**batch, params=state.params, train=False)[0] #stack the model's forward pass with the logits function\n    return state.logits_function(logits)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.324539Z","iopub.execute_input":"2022-01-21T14:09:52.324797Z","iopub.status.idle":"2022-01-21T14:09:52.331650Z","shell.execute_reply.started":"2022-01-21T14:09:52.324763Z","shell.execute_reply":"2022-01-21T14:09:52.330955Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"parallel_eval_step = jax.pmap(eval_step, axis_name=\"batch\")","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.333235Z","iopub.execute_input":"2022-01-21T14:09:52.333995Z","iopub.status.idle":"2022-01-21T14:09:52.340439Z","shell.execute_reply.started":"2022-01-21T14:09:52.333951Z","shell.execute_reply":"2022-01-21T14:09:52.339629Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Data loaders ","metadata":{}},{"cell_type":"code","source":"# Returns batch model input\n# 1. define random permutation \n# 2. randomized dataset is extracted and then it converted to a JAX array and sharded over all local TPU devices.\ndef train_data_loader(rng, dataset, batch_size):\n    steps_per_epoch = len(dataset) // batch_size\n    perms = jax.random.permutation(rng, len(dataset))\n    perms = perms[: steps_per_epoch * batch_size]  # Skip incomplete batch.\n    perms = perms.reshape((steps_per_epoch, batch_size))\n\n    for perm in perms:\n        batch = dataset[perm]\n        batch = {k: jnp.array(v) for k, v in batch.items()}\n        batch = shard(batch)\n        yield batch","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.341643Z","iopub.execute_input":"2022-01-21T14:09:52.342005Z","iopub.status.idle":"2022-01-21T14:09:52.350619Z","shell.execute_reply.started":"2022-01-21T14:09:52.341967Z","shell.execute_reply":"2022-01-21T14:09:52.349808Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# similar to train data loader \ndef eval_data_loader(dataset, batch_size): \n    for i in range(len(dataset) // batch_size):\n        batch = dataset[i * batch_size : (i + 1) * batch_size]\n        batch = {k: jnp.array(v) for k, v in batch.items()}\n        batch = shard(batch)\n        yield batch","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.352665Z","iopub.execute_input":"2022-01-21T14:09:52.352852Z","iopub.status.idle":"2022-01-21T14:09:52.361127Z","shell.execute_reply.started":"2022-01-21T14:09:52.352829Z","shell.execute_reply":"2022-01-21T14:09:52.360402Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# replicate/copy the weight parameters on each device, to passthem to our pmapped functions.\nstate = flax.jax_utils.replicate(state)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.362598Z","iopub.execute_input":"2022-01-21T14:09:52.363134Z","iopub.status.idle":"2022-01-21T14:09:52.453247Z","shell.execute_reply.started":"2022-01-21T14:09:52.363089Z","shell.execute_reply":"2022-01-21T14:09:52.452511Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# generating a seeded PRNGKey for the dropout layers and dataset shuffling.\nrng = jax.random.PRNGKey(seed)\ndropout_rngs = jax.random.split(rng, jax.local_device_count())","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.454727Z","iopub.execute_input":"2022-01-21T14:09:52.454964Z","iopub.status.idle":"2022-01-21T14:09:52.462910Z","shell.execute_reply.started":"2022-01-21T14:09:52.454931Z","shell.execute_reply":"2022-01-21T14:09:52.461986Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n# Full training loop\nfor i, epoch in enumerate(tqdm(range(1, num_train_epochs + 1), desc=f\"Epoch ...\", position=0, leave=True)):\n    rng, input_rng = jax.random.split(rng)\n\n    # train\n    with tqdm(total=len(train_dataset) // total_batch_size, desc=\"Training...\", leave=False) as progress_bar_train:\n        for batch in train_data_loader(input_rng, train_dataset, total_batch_size):\n            state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)\n            progress_bar_train.update(1)\n\n    # evaluate\n    with tqdm(total=len(validation_dataset) // total_batch_size, desc=\"Evaluating...\", leave=False) as progress_bar_eval:\n          for batch in eval_data_loader(validation_dataset, total_batch_size):\n                labels = batch.pop(\"labels\")\n                predictions = parallel_eval_step(state, batch)\n                metric.add_batch(predictions=chain(*predictions), references=chain(*labels))\n                progress_bar_eval.update(1)\n\n    eval_metric = metric.compute(average='macro')\n\n    loss = round(flax.jax_utils.unreplicate(train_metrics)['loss'].item(), 3)\n    eval_score = round(list(eval_metric.values())[0],3)\n    metric_name = list(eval_metric.keys())[0]\n\n    print(f\"{i+1}/{num_train_epochs} | Train loss: {loss} | Eval {metric_name}: {eval_score}\")","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:09:52.464624Z","iopub.execute_input":"2022-01-21T14:09:52.464975Z","iopub.status.idle":"2022-01-21T14:32:15.436459Z","shell.execute_reply.started":"2022-01-21T14:09:52.464940Z","shell.execute_reply":"2022-01-21T14:32:15.435651Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# ","metadata":{}},{"cell_type":"markdown","source":"# Test Generation","metadata":{}},{"cell_type":"code","source":"# preprocess test dataset\ndef preprocess_test_set(input_batch):\n    '''\n    This function will take batch of data and returns the tokenized processed data\n    INPUT - input batch from from original dataset\n    RETURNS processed data\n    '''\n    texts = (input_batch[\"review_text\"],)\n    processed = tokenizer(*texts, padding=\"max_length\", max_length=128, truncation=True)\n    return processed","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:32:15.437938Z","iopub.execute_input":"2022-01-21T14:32:15.438206Z","iopub.status.idle":"2022-01-21T14:32:15.443098Z","shell.execute_reply.started":"2022-01-21T14:32:15.438171Z","shell.execute_reply":"2022-01-21T14:32:15.442230Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# using tokenized dataset\ntokenized_test_dataset=None\nif not os.path.exists(\"../input/goodreads-tokenized-data/test_tokenized_dataset\"):\n    tokenized_test_dataset = raw_test.map(preprocess_test_set, batched=True, remove_columns=raw_test[\"test\"].column_names)\nelse:\n    print(\"training dataset is already tokenized\")\n    tokenized_test_dataset=load_from_disk(\"../input/goodreads-tokenized-data/test_tokenized_dataset\")\ntest_dataset = tokenized_test_dataset[\"test\"]","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:32:15.445222Z","iopub.execute_input":"2022-01-21T14:32:15.445739Z","iopub.status.idle":"2022-01-21T14:32:15.479814Z","shell.execute_reply.started":"2022-01-21T14:32:15.445701Z","shell.execute_reply":"2022-01-21T14:32:15.479022Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dataset","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-01-21T14:32:15.481013Z","iopub.execute_input":"2022-01-21T14:32:15.481550Z","iopub.status.idle":"2022-01-21T14:32:15.487182Z","shell.execute_reply.started":"2022-01-21T14:32:15.481507Z","shell.execute_reply":"2022-01-21T14:32:15.486454Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# similar to train dataloader, it takes dataset and batch_size\ndef test_data_loader(dataset,batch_size):\n    if len(dataset)<batch_size:\n        batch = dataset[:]\n        batch = {k: jnp.array(v) for k, v in batch.items()}\n        batch = shard(batch)\n        yield batch\n    else:\n        for i in range(len(dataset) // batch_size):\n            batch = dataset[i * batch_size : (i + 1) * batch_size]\n            batch = {k: jnp.array(v) for k, v in batch.items()}\n            batch = shard(batch)\n            yield batch\n        batch = dataset[(i+1) * batch_size:]\n        batch = {k: jnp.array(v) for k, v in batch.items()}\n        batch = shard(batch)\n        yield batch","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:32:15.488436Z","iopub.execute_input":"2022-01-21T14:32:15.488906Z","iopub.status.idle":"2022-01-21T14:32:15.498492Z","shell.execute_reply.started":"2022-01-21T14:32:15.488869Z","shell.execute_reply":"2022-01-21T14:32:15.497663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n# Running on test dataset , storing predictions \npreds=[]\nfor batch in test_data_loader(test_dataset, total_batch_size):\n    predictions = parallel_eval_step(state, batch)\n    preds.append(predictions[0])    ","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:32:15.501060Z","iopub.execute_input":"2022-01-21T14:32:15.501250Z","iopub.status.idle":"2022-01-21T14:32:15.509187Z","shell.execute_reply.started":"2022-01-21T14:32:15.501220Z","shell.execute_reply":"2022-01-21T14:32:15.508126Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# convert it into numpy array\npredictions=[]\nfor pred in preds:\n    predictions.extend(np.array(pred))","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:32:15.510116Z","iopub.status.idle":"2022-01-21T14:32:15.511668Z","shell.execute_reply.started":"2022-01-21T14:32:15.510433Z","shell.execute_reply":"2022-01-21T14:32:15.510458Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# create final submission file \nsample = pd.read_csv('../input/goodreads-books-reviews-290312/goodreads_sample_submission.csv')\nsample[\"rating\"]=predictions","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:32:15.513318Z","iopub.status.idle":"2022-01-21T14:32:15.514173Z","shell.execute_reply.started":"2022-01-21T14:32:15.513875Z","shell.execute_reply":"2022-01-21T14:32:15.513901Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample.head(2)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:32:15.516075Z","iopub.status.idle":"2022-01-21T14:32:15.516884Z","shell.execute_reply.started":"2022-01-21T14:32:15.516675Z","shell.execute_reply":"2022-01-21T14:32:15.516697Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample.to_csv('submission.csv',index=False)","metadata":{"execution":{"iopub.status.busy":"2022-01-21T14:32:15.518202Z","iopub.status.idle":"2022-01-21T14:32:15.518801Z","shell.execute_reply.started":"2022-01-21T14:32:15.518554Z","shell.execute_reply":"2022-01-21T14:32:15.518579Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Conclusion \nSo we successfully predicted rating score from the review text and submission.csv is created which you can find in Outputs (\"/kaggle/working\"). You can submit your submission file to the competition.It can be found in sidebar of notebook editor, under **Competitions** by clicking on the \"**Submit**\" Button or you can go to competition and click on the **submit predictions**. \n\nNote that submission format should be same as described [here](https://www.kaggle.com/c/goodreads-books-reviews-290312/overview/evaluation)\n\n\n","metadata":{}},{"cell_type":"markdown","source":"# References\n* https://flax.readthedocs.io/en/latest/notebooks/annotated_mnist.html\n* https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification_flax.ipynb#scrollTo=Mn1GdGpipfWK\n* https://www.kaggle.com/heyytanay/sentiment-clf-jax-flax-on-tpus-w-b\n* https://optax.readthedocs.io/\n* https://flax.readthedocs.io/\n* https://jax.readthedocs.io/\n* https://github.com/google/flax/tree/main/examples","metadata":{}}]}