{"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":"### Intro to JAX\n[JAX](https://github.com/google/jax) is a framework which is used for high-performance numerical computing and machine learning research developed at [Google Research](https://research.google/) teams. It allows you to build Python applications with a NumPy-consistent API that specializes in differentiating, vectorizing, parallelizing, and compiling to GPU/TPU Just-In-Time. JAX was designed with performance and speed as a first priority, and is natively compatible with common machine learning accelerators such as [GPUs](https://www.kaggle.com/docs/efficient-gpu-usage) and [TPUs](https://www.kaggle.com/docs/tpu). Large ML models can take ages to train -- you might be interested in using JAX for applications where speed and performance are particularly important!\n### When to use JAX vs TensorFlow?\n[TensorFlow](https://www.tensorflow.org/guide) is a fantastic product, with a rich and fully-featured ecosystem, capable of supporting most every use case a machine learning practitioner might have (e.g. [TFLite](https://www.tensorflow.org/lite) for on-device inference computing, [TFHub](https://tfhub.dev/) for sharing pre-trained models, and many additional specialized applications as well). This type of broad mandate both contrasts and compliments JAX's philosophy, which is more narrowly focused on speed and performance.  We recommend using JAX in situations where you do want to maximize speed and performance but you do not require any of the long tail of features and additional functionalities that only the [TensorFlow ecosystem](https://www.tensorflow.org/learn) can provide.\n### Intro to the FLAX\nJust like [JAX](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html) focuses on speed, other members of the JAX ecosystem are encouraged to specialize as well.  For example, [Flax](https://flax.readthedocs.io/en/latest/) focuses on neural networks and [jgraph](https://github.com/deepmind/jraph) focuses on graph networks.  \n\n[Flax](https://flax.readthedocs.io/en/latest/) is a JAX-based neural network library that was initially developed by  Google Research's Brain Team (in close collaboration with the JAX team) but is now open source.  If you want to train machine learning models on GPUs and TPUs at an accelerated speed, or if you have an ML project that might benefit from bringing together both [Autograd](https://github.com/hips/autograd) and [XLA](https://www.tensorflow.org/xla), consider using [Flax](https://flax.readthedocs.io/en/latest/) for your next project! [Flax](https://flax.readthedocs.io/en/latest/) is especially well-suited for projects that use large language models, and is a popular choice for cutting-edge [machine learning research](https://arxiv.org/search/?query=JAX&searchtype=all&abstracts=show&order=-announced_date_first&size=50).\n\n### Disclaimer:\n**We recommend using [GPUs](https://www.kaggle.com/docs/efficient-gpu-usage) when working with JAX on Kaggle.** These notebooks are compatible with the v3-8 [TPUs](https://www.kaggle.com/docs/tpu) that are provided for free in [Kaggle Notebooks](https://www.kaggle.com/code/new), but JAX was optimized for the newly updated [TPU VM](https://cloud.google.com/blog/products/compute/introducing-cloud-tpu-vms) architecture which is not yet available on Kaggle.\n","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.execute_input":"2022-01-02T05:30:15.814483Z","iopub.status.busy":"2022-01-02T05:30:15.814217Z","iopub.status.idle":"2022-01-02T05:30:15.818748Z","shell.execute_reply":"2022-01-02T05:30:15.817716Z","shell.execute_reply.started":"2022-01-02T05:30:15.814451Z"},"papermill":{"duration":0.052493,"end_time":"2022-03-03T16:35:00.136503","exception":false,"start_time":"2022-03-03T16:35:00.084010","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Imports","metadata":{"papermill":{"duration":0.049039,"end_time":"2022-03-03T16:35:00.237029","exception":false,"start_time":"2022-03-03T16:35:00.187990","status":"completed"},"tags":[]}},{"cell_type":"code","source":"#Uncomment and Run when only accelerator is TPU\n#%%capture\n#!conda install -y -c conda-forge jax jaxlib flax optax datasets transformers\n#!conda install -y importlib-metadata","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:35:00.440274Z","iopub.status.busy":"2022-03-03T16:35:00.437872Z","iopub.status.idle":"2022-03-03T16:35:00.451187Z","shell.execute_reply":"2022-03-03T16:35:00.452005Z","shell.execute_reply.started":"2022-02-16T20:44:09.318412Z"},"papermill":{"duration":0.137667,"end_time":"2022-03-03T16:35:00.452367","exception":false,"start_time":"2022-03-03T16:35:00.314700","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install datasets transformers","metadata":{"_kg_hide-output":true,"execution":{"iopub.execute_input":"2022-03-03T16:35:00.637425Z","iopub.status.busy":"2022-03-03T16:35:00.635667Z","iopub.status.idle":"2022-03-03T16:35:14.670179Z","shell.execute_reply":"2022-03-03T16:35:14.668872Z","shell.execute_reply.started":"2022-03-03T13:58:59.537035Z"},"papermill":{"duration":14.13431,"end_time":"2022-03-03T16:35:14.670356","exception":false,"start_time":"2022-03-03T16:35:00.536046","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Importing all the libraries necessary for the project\nimport os\nimport time\nimport jax\nimport flax\nimport optax\nimport datasets\nimport pandas as pd \nimport numpy as np\nfrom jax import jit\nimport jax.numpy as jnp\nimport tensorflow as tf\nfrom flax.training import train_state\nfrom itertools import chain\nfrom tqdm.notebook import tqdm\nfrom typing import Callable\nfrom flax import traverse_util\nfrom datasets import load_dataset, load_metric ,Dataset,list_metrics,load_from_disk\nfrom flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key\nfrom transformers import FlaxAutoModelForSequenceClassification, AutoConfig, AutoTokenizer, BertTokenizer\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n# to suppress warnings caused by cuda version\nos.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:35:14.798589Z","iopub.status.busy":"2022-03-03T16:35:14.797595Z","iopub.status.idle":"2022-03-03T16:35:24.770069Z","shell.execute_reply":"2022-03-03T16:35:24.769280Z","shell.execute_reply.started":"2022-03-03T13:59:10.522560Z"},"papermill":{"duration":10.039828,"end_time":"2022-03-03T16:35:24.770273","exception":false,"start_time":"2022-03-03T16:35:14.730445","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### TPU detection and configuration\n**We recommend using [GPUs](https://www.kaggle.com/docs/efficient-gpu-usage) when working with JAX on Kaggle.** These notebooks are compatible with the v3-8 [TPUs](https://www.kaggle.com/docs/tpu) that are provided for free in [Kaggle Notebooks](https://www.kaggle.com/code/new), but JAX was optimized for the newly updated [TPU VM](https://cloud.google.com/blog/products/compute/introducing-cloud-tpu-vms) architecture which is not yet available on Kaggle.\n","metadata":{"papermill":{"duration":0.058995,"end_time":"2022-03-03T16:35:24.890272","exception":false,"start_time":"2022-03-03T16:35:24.831277","status":"completed"},"tags":[]}},{"cell_type":"code","source":"if '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.')","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:35:25.016933Z","iopub.status.busy":"2022-03-03T16:35:25.015968Z","iopub.status.idle":"2022-03-03T16:35:25.019909Z","shell.execute_reply":"2022-03-03T16:35:25.020763Z","shell.execute_reply.started":"2022-03-03T13:59:33.685423Z"},"papermill":{"duration":0.070295,"end_time":"2022-03-03T16:35:25.021034","exception":false,"start_time":"2022-03-03T16:35:24.950739","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"jax.local_devices()","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:35:25.143494Z","iopub.status.busy":"2022-03-03T16:35:25.142736Z","iopub.status.idle":"2022-03-03T16:35:25.472720Z","shell.execute_reply":"2022-03-03T16:35:25.473263Z","shell.execute_reply.started":"2022-03-03T13:59:45.928166Z"},"papermill":{"duration":0.391722,"end_time":"2022-03-03T16:35:25.473438","exception":false,"start_time":"2022-03-03T16:35:25.081716","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load data and preprocess the data\nLoading train data csv file using Huggingface's [`load_dataset`](https://huggingface.co/docs/datasets/loading_datasets.html) function from Dataset class","metadata":{"papermill":{"duration":0.058117,"end_time":"2022-03-03T16:35:25.589500","exception":false,"start_time":"2022-03-03T16:35:25.531383","status":"completed"},"tags":[]}},{"cell_type":"code","source":"raw_train = load_dataset(\"csv\", data_files={'train': ['../input/jigsaw-multilingual-toxic-comment-classification/jigsaw-toxic-comment-train.csv']})","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:35:25.717714Z","iopub.status.busy":"2022-03-03T16:35:25.711801Z","iopub.status.idle":"2022-03-03T16:35:29.138336Z","shell.execute_reply":"2022-03-03T16:35:29.136391Z","shell.execute_reply.started":"2022-03-03T13:59:57.780539Z"},"papermill":{"duration":3.490245,"end_time":"2022-03-03T16:35:29.138475","exception":false,"start_time":"2022-03-03T16:35:25.648230","status":"completed"},"tags":[],"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Spliting the dataset into train and eval sets","metadata":{"papermill":{"duration":0.060536,"end_time":"2022-03-03T16:35:29.260671","exception":false,"start_time":"2022-03-03T16:35:29.200135","status":"completed"},"tags":[]}},{"cell_type":"code","source":"raw_train = raw_train[\"train\"].train_test_split(0.2)","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:35:29.392287Z","iopub.status.busy":"2022-03-03T16:35:29.391151Z","iopub.status.idle":"2022-03-03T16:35:29.428435Z","shell.execute_reply":"2022-03-03T16:35:29.429337Z","shell.execute_reply.started":"2022-02-16T04:06:58.440328Z"},"papermill":{"duration":0.107311,"end_time":"2022-03-03T16:35:29.429503","exception":false,"start_time":"2022-03-03T16:35:29.322192","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"raw_train","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:35:29.559726Z","iopub.status.busy":"2022-03-03T16:35:29.555982Z","iopub.status.idle":"2022-03-03T16:35:29.563559Z","shell.execute_reply":"2022-03-03T16:35:29.564155Z","shell.execute_reply.started":"2022-02-16T04:07:11.708019Z"},"papermill":{"duration":0.073108,"end_time":"2022-03-03T16:35:29.564324","exception":false,"start_time":"2022-03-03T16:35:29.491216","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Loading the model checkpoint and tokenizer","metadata":{"papermill":{"duration":0.06175,"end_time":"2022-03-03T16:35:29.687887","exception":false,"start_time":"2022-03-03T16:35:29.626137","status":"completed"},"tags":[]}},{"cell_type":"code","source":"model_checkpoint = \"bert-base-cased\" \ntokenizer = BertTokenizer.from_pretrained(model_checkpoint,use_fast=True)","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:35:29.820188Z","iopub.status.busy":"2022-03-03T16:35:29.819503Z","iopub.status.idle":"2022-03-03T16:35:35.973706Z","shell.execute_reply":"2022-03-03T16:35:35.973076Z","shell.execute_reply.started":"2022-02-16T04:07:35.000718Z"},"papermill":{"duration":6.223742,"end_time":"2022-03-03T16:35:35.973895","exception":false,"start_time":"2022-03-03T16:35:29.750153","status":"completed"},"tags":[],"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Pre-process the dataset\nNow, this function will preprocess the dataset by taking batch of data and returns the tokenized processed data","metadata":{"papermill":{"duration":0.064788,"end_time":"2022-03-03T16:35:36.104234","exception":false,"start_time":"2022-03-03T16:35:36.039446","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def preprocess_function(input_batch):\n    '''\n    INPUT - input batch from from original dataset\n    RETURNS preprocessed data\n    '''\n    texts = (input_batch[\"comment_text\"],)\n    processed = tokenizer(*texts, padding=\"max_length\", max_length=128, truncation=True)\n    processed[\"labels\"] = input_batch[\"toxic\"]\n    return processed","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:35:36.243895Z","iopub.status.busy":"2022-03-03T16:35:36.242635Z","iopub.status.idle":"2022-03-03T16:35:36.245355Z","shell.execute_reply":"2022-03-03T16:35:36.246097Z","shell.execute_reply.started":"2022-02-16T04:08:18.446424Z"},"papermill":{"duration":0.077083,"end_time":"2022-03-03T16:35:36.246261","exception":false,"start_time":"2022-03-03T16:35:36.169178","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tokenized_dataset = raw_train.map(preprocess_function, batched=True, remove_columns=raw_train[\"train\"].column_names)","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:35:36.426279Z","iopub.status.busy":"2022-03-03T16:35:36.385136Z","iopub.status.idle":"2022-03-03T16:44:34.327616Z","shell.execute_reply":"2022-03-03T16:44:34.328217Z","shell.execute_reply.started":"2022-02-16T04:09:31.093933Z"},"papermill":{"duration":538.015733,"end_time":"2022-03-03T16:44:34.328431","exception":false,"start_time":"2022-03-03T16:35:36.312698","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tokenized_dataset","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:44:34.567822Z","iopub.status.busy":"2022-03-03T16:44:34.566006Z","iopub.status.idle":"2022-03-03T16:44:34.570554Z","shell.execute_reply":"2022-03-03T16:44:34.566872Z","shell.execute_reply.started":"2022-02-16T04:18:16.897717Z"},"papermill":{"duration":0.124386,"end_time":"2022-03-03T16:44:34.570809","exception":false,"start_time":"2022-03-03T16:44:34.446423","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = tokenized_dataset[\"train\"]\nvalidation_dataset = tokenized_dataset[\"test\"]","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:44:34.798001Z","iopub.status.busy":"2022-03-03T16:44:34.796779Z","iopub.status.idle":"2022-03-03T16:44:34.799905Z","shell.execute_reply":"2022-03-03T16:44:34.798904Z","shell.execute_reply.started":"2022-02-16T04:18:19.610415Z"},"papermill":{"duration":0.119842,"end_time":"2022-03-03T16:44:34.800165","exception":false,"start_time":"2022-03-03T16:44:34.680323","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:44:35.025701Z","iopub.status.busy":"2022-03-03T16:44:35.024667Z","iopub.status.idle":"2022-03-03T16:44:35.030508Z","shell.execute_reply":"2022-03-03T16:44:35.031313Z","shell.execute_reply.started":"2022-02-16T04:18:21.918072Z"},"papermill":{"duration":0.123908,"end_time":"2022-03-03T16:44:35.031552","exception":false,"start_time":"2022-03-03T16:44:34.907644","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"validation_dataset","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:44:35.210195Z","iopub.status.busy":"2022-03-03T16:44:35.206498Z","iopub.status.idle":"2022-03-03T16:44:35.213095Z","shell.execute_reply":"2022-03-03T16:44:35.213628Z","shell.execute_reply.started":"2022-02-16T04:18:32.450312Z"},"papermill":{"duration":0.077618,"end_time":"2022-03-03T16:44:35.213806","exception":false,"start_time":"2022-03-03T16:44:35.136188","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Listing and selecting the metrics\nSelecting the evaluation metrics using HuggingFace's [`load_metrics`](https://huggingface.co/docs/datasets/loading_metrics.html) ","metadata":{"papermill":{"duration":0.066028,"end_time":"2022-03-03T16:44:35.345962","exception":false,"start_time":"2022-03-03T16:44:35.279934","status":"completed"},"tags":[]}},{"cell_type":"code","source":"metrics_list = list_metrics()\nmetrics_list","metadata":{"_kg_hide-output":true,"execution":{"iopub.execute_input":"2022-03-03T16:44:35.490078Z","iopub.status.busy":"2022-03-03T16:44:35.489333Z","iopub.status.idle":"2022-03-03T16:44:36.111106Z","shell.execute_reply":"2022-03-03T16:44:36.110520Z","shell.execute_reply.started":"2022-03-03T14:00:04.427217Z"},"papermill":{"duration":0.696001,"end_time":"2022-03-03T16:44:36.111258","exception":false,"start_time":"2022-03-03T16:44:35.415257","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"metric = load_metric('f1')\nmetric","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:44:36.251181Z","iopub.status.busy":"2022-03-03T16:44:36.250454Z","iopub.status.idle":"2022-03-03T16:44:37.379231Z","shell.execute_reply":"2022-03-03T16:44:37.378392Z","shell.execute_reply.started":"2022-03-03T14:01:21.796429Z"},"papermill":{"duration":1.201732,"end_time":"2022-03-03T16:44:37.379380","exception":false,"start_time":"2022-03-03T16:44:36.177648","status":"completed"},"tags":[],"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model config\nDefining all the model config parameters below","metadata":{"papermill":{"duration":0.06648,"end_time":"2022-03-03T16:44:37.514572","exception":false,"start_time":"2022-03-03T16:44:37.448092","status":"completed"},"tags":[]}},{"cell_type":"code","source":"num_labels = 2 # 0-1\nseed = 0\nnum_train_epochs = 5\nlearning_rate = 2e-5\nper_device_batch_size = 32\nweight_decay=1e-2","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:44:37.656790Z","iopub.status.busy":"2022-03-03T16:44:37.654673Z","iopub.status.idle":"2022-03-03T16:44:37.657548Z","shell.execute_reply":"2022-03-03T16:44:37.658108Z","shell.execute_reply.started":"2022-02-16T20:49:51.91484Z"},"papermill":{"duration":0.07606,"end_time":"2022-03-03T16:44:37.658260","exception":false,"start_time":"2022-03-03T16:44:37.582200","status":"completed"},"tags":[]},"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.execute_input":"2022-03-03T16:44:37.801658Z","iopub.status.busy":"2022-03-03T16:44:37.801005Z","iopub.status.idle":"2022-03-03T16:44:37.806170Z","shell.execute_reply":"2022-03-03T16:44:37.805498Z","shell.execute_reply.started":"2022-02-16T20:49:54.934173Z"},"papermill":{"duration":0.078933,"end_time":"2022-03-03T16:44:37.806310","exception":false,"start_time":"2022-03-03T16:44:37.727377","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Loading the config and the pre-trained model using HuggingFace's from_pretrained \nconfig = AutoConfig.from_pretrained(model_checkpoint, num_labels=num_labels)\nmodel = FlaxAutoModelForSequenceClassification.from_pretrained(model_checkpoint, config=config, seed=seed)","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:44:37.950527Z","iopub.status.busy":"2022-03-03T16:44:37.944826Z","iopub.status.idle":"2022-03-03T16:45:15.782790Z","shell.execute_reply":"2022-03-03T16:45:15.782161Z","shell.execute_reply.started":"2022-02-16T04:19:58.589238Z"},"papermill":{"duration":37.90983,"end_time":"2022-03-03T16:45:15.782938","exception":false,"start_time":"2022-03-03T16:44:37.873108","status":"completed"},"tags":[],"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"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.execute_input":"2022-03-03T16:45:15.936969Z","iopub.status.busy":"2022-03-03T16:45:15.928890Z","iopub.status.idle":"2022-03-03T16:45:16.436014Z","shell.execute_reply":"2022-03-03T16:45:16.437500Z","shell.execute_reply.started":"2022-02-16T04:20:35.042068Z"},"papermill":{"duration":0.583189,"end_time":"2022-03-03T16:45:16.437853","exception":false,"start_time":"2022-03-03T16:45:15.854664","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train state\n","metadata":{"papermill":{"duration":0.140519,"end_time":"2022-03-03T16:45:16.705877","exception":false,"start_time":"2022-03-03T16:45:16.565358","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class TrainState(train_state.TrainState):\n    '''\n    Derived TrainState class that saves the forward pass of the model as an eval function and a loss function\n    '''\n    logits_function: Callable = flax.struct.field(pytree_node=False)\n    loss_function: Callable = flax.struct.field(pytree_node=False)","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:45:16.957945Z","iopub.status.busy":"2022-03-03T16:45:16.955636Z","iopub.status.idle":"2022-03-03T16:45:16.958673Z","shell.execute_reply":"2022-03-03T16:45:16.959256Z","shell.execute_reply.started":"2022-02-16T04:20:44.057816Z"},"papermill":{"duration":0.12909,"end_time":"2022-03-03T16:45:16.959409","exception":false,"start_time":"2022-03-03T16:45:16.830319","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decay_mask_fn(params):\n    '''\n    This function's task is to make sure that weight decay is not applies to any bias or Layernorm weights\n    '''\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.execute_input":"2022-03-03T16:45:17.109440Z","iopub.status.busy":"2022-03-03T16:45:17.108262Z","iopub.status.idle":"2022-03-03T16:45:17.110829Z","shell.execute_reply":"2022-03-03T16:45:17.111401Z","shell.execute_reply.started":"2022-02-16T04:20:53.732169Z"},"papermill":{"duration":0.080972,"end_time":"2022-03-03T16:45:17.111608","exception":false,"start_time":"2022-03-03T16:45:17.030636","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Adam optimizer function using optax.adamw\ndef 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.execute_input":"2022-03-03T16:45:17.260576Z","iopub.status.busy":"2022-03-03T16:45:17.259686Z","iopub.status.idle":"2022-03-03T16:45:17.263259Z","shell.execute_reply":"2022-03-03T16:45:17.263855Z","shell.execute_reply.started":"2022-02-16T04:21:23.490741Z"},"papermill":{"duration":0.081504,"end_time":"2022-03-03T16:45:17.264016","exception":false,"start_time":"2022-03-03T16:45:17.182512","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"adamw = adamw(weight_decay)","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:45:17.416214Z","iopub.status.busy":"2022-03-03T16:45:17.414022Z","iopub.status.idle":"2022-03-03T16:45:17.417564Z","shell.execute_reply":"2022-03-03T16:45:17.418130Z","shell.execute_reply.started":"2022-02-16T04:21:24.023497Z"},"papermill":{"duration":0.081681,"end_time":"2022-03-03T16:45:17.418299","exception":false,"start_time":"2022-03-03T16:45:17.336618","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## Defining the loss and the evaluation function\n@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.execute_input":"2022-03-03T16:45:17.567633Z","iopub.status.busy":"2022-03-03T16:45:17.566878Z","iopub.status.idle":"2022-03-03T16:45:17.570658Z","shell.execute_reply":"2022-03-03T16:45:17.570126Z","shell.execute_reply.started":"2022-02-16T04:22:03.156557Z"},"papermill":{"duration":0.082091,"end_time":"2022-03-03T16:45:17.570846","exception":false,"start_time":"2022-03-03T16:45:17.488755","status":"completed"},"tags":[]},"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.execute_input":"2022-03-03T16:45:17.723771Z","iopub.status.busy":"2022-03-03T16:45:17.717261Z","iopub.status.idle":"2022-03-03T16:45:18.304979Z","shell.execute_reply":"2022-03-03T16:45:18.304001Z","shell.execute_reply.started":"2022-02-16T04:22:20.932544Z"},"papermill":{"duration":0.664078,"end_time":"2022-03-03T16:45:18.305132","exception":false,"start_time":"2022-03-03T16:45:17.641054","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train and evaluate steps","metadata":{"papermill":{"duration":0.069359,"end_time":"2022-03-03T16:45:18.448810","exception":false,"start_time":"2022-03-03T16:45:18.379451","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def 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.execute_input":"2022-03-03T16:45:18.596561Z","iopub.status.busy":"2022-03-03T16:45:18.595847Z","iopub.status.idle":"2022-03-03T16:45:18.600084Z","shell.execute_reply":"2022-03-03T16:45:18.600585Z","shell.execute_reply.started":"2022-02-16T04:22:37.388023Z"},"papermill":{"duration":0.082655,"end_time":"2022-03-03T16:45:18.600766","exception":false,"start_time":"2022-03-03T16:45:18.518111","status":"completed"},"tags":[]},"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\n","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:45:18.746822Z","iopub.status.busy":"2022-03-03T16:45:18.745779Z","iopub.status.idle":"2022-03-03T16:45:18.748957Z","shell.execute_reply":"2022-03-03T16:45:18.749523Z","shell.execute_reply.started":"2022-02-16T04:22:49.243077Z"},"papermill":{"duration":0.079757,"end_time":"2022-03-03T16:45:18.749671","exception":false,"start_time":"2022-03-03T16:45:18.669914","status":"completed"},"tags":[]},"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.execute_input":"2022-03-03T16:45:18.896735Z","iopub.status.busy":"2022-03-03T16:45:18.895709Z","iopub.status.idle":"2022-03-03T16:45:18.899122Z","shell.execute_reply":"2022-03-03T16:45:18.898518Z","shell.execute_reply.started":"2022-02-16T04:22:58.56507Z"},"papermill":{"duration":0.078878,"end_time":"2022-03-03T16:45:18.899269","exception":false,"start_time":"2022-03-03T16:45:18.820391","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"parallel_eval_step = jax.pmap(eval_step, axis_name=\"batch\")","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:45:19.049060Z","iopub.status.busy":"2022-03-03T16:45:19.046817Z","iopub.status.idle":"2022-03-03T16:45:19.049875Z","shell.execute_reply":"2022-03-03T16:45:19.050361Z","shell.execute_reply.started":"2022-02-16T04:23:06.747206Z"},"papermill":{"duration":0.080585,"end_time":"2022-03-03T16:45:19.050543","exception":false,"start_time":"2022-03-03T16:45:18.969958","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data loader","metadata":{"papermill":{"duration":0.070317,"end_time":"2022-03-03T16:45:19.191422","exception":false,"start_time":"2022-03-03T16:45:19.121105","status":"completed"},"tags":[]}},{"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.execute_input":"2022-03-03T16:45:19.342377Z","iopub.status.busy":"2022-03-03T16:45:19.341329Z","iopub.status.idle":"2022-03-03T16:45:19.344963Z","shell.execute_reply":"2022-03-03T16:45:19.344353Z","shell.execute_reply.started":"2022-02-16T04:23:22.0659Z"},"papermill":{"duration":0.082537,"end_time":"2022-03-03T16:45:19.345105","exception":false,"start_time":"2022-03-03T16:45:19.262568","status":"completed"},"tags":[]},"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.execute_input":"2022-03-03T16:45:19.494284Z","iopub.status.busy":"2022-03-03T16:45:19.493190Z","iopub.status.idle":"2022-03-03T16:45:19.495594Z","shell.execute_reply":"2022-03-03T16:45:19.496294Z","shell.execute_reply.started":"2022-02-16T04:23:34.0002Z"},"papermill":{"duration":0.0798,"end_time":"2022-03-03T16:45:19.496469","exception":false,"start_time":"2022-03-03T16:45:19.416669","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"state = flax.jax_utils.replicate(state)","metadata":{"execution":{"iopub.execute_input":"2022-03-03T16:45:19.665838Z","iopub.status.busy":"2022-03-03T16:45:19.655298Z","iopub.status.idle":"2022-03-03T16:45:19.757747Z","shell.execute_reply":"2022-03-03T16:45:19.757215Z","shell.execute_reply.started":"2022-02-16T04:23:43.897054Z"},"papermill":{"duration":0.191321,"end_time":"2022-03-03T16:45:19.757968","exception":false,"start_time":"2022-03-03T16:45:19.566647","status":"completed"},"tags":[]},"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.execute_input":"2022-03-03T16:45:19.906498Z","iopub.status.busy":"2022-03-03T16:45:19.905379Z","iopub.status.idle":"2022-03-03T16:45:20.034099Z","shell.execute_reply":"2022-03-03T16:45:20.032907Z","shell.execute_reply.started":"2022-02-16T04:23:50.783891Z"},"papermill":{"duration":0.20561,"end_time":"2022-03-03T16:45:20.034241","exception":false,"start_time":"2022-03-03T16:45:19.828631","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training ","metadata":{"papermill":{"duration":0.070445,"end_time":"2022-03-03T16:45:20.175599","exception":false,"start_time":"2022-03-03T16:45:20.105154","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Now, we'll define the training loop and train the pre-trained model\nstart = time.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}\")\n    \nprint(\"Total time: \", time.time() - start, \"seconds\")","metadata":{"_kg_hide-output":false,"execution":{"iopub.execute_input":"2022-03-03T16:45:20.359510Z","iopub.status.busy":"2022-03-03T16:45:20.358830Z","iopub.status.idle":"2022-03-03T20:34:52.485163Z","shell.execute_reply":"2022-03-03T20:34:52.484503Z","shell.execute_reply.started":"2022-02-14T16:30:32.800911Z"},"papermill":{"duration":13772.237367,"end_time":"2022-03-03T20:34:52.485325","exception":false,"start_time":"2022-03-03T16:45:20.247958","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **Conclusion**\nHere in this notebook, we've illustrated how [JAX](https://github.com/google/jax) and [FLAX](https://flax.readthedocs.io/en/latest/) can be used to train the pre-trained neural network for the text classification dataset, with the F1 score of more than 80%. To see more examples of how to use [JAX](https://github.com/google/jax) and [FLAX](https://flax.readthedocs.io/en/latest/) with different data formats, please see this discussion post.  \n\nNow, it's your turn to  create some amazing notebooks using [JAX](https://github.com/google/jax) and [FLAX](https://flax.readthedocs.io/en/latest/). \n\n### **Useful resources which helped me:**\n\n* https://flax.readthedocs.io/en/latest/index.html\n* https://github.com/google/flax/tree/main/examples\n* https://www.kaggle.com/heyytanay/sentiment-clf-jax-flax-on-tpus-w-b/notebook\n* https://www.kaggle.com/asvskartheek/bert-tpus-jax-huggingface/notebook\n* https://huggingface.co/docs/datasets/package_reference/main_classes.html#dataset\n* https://colab.sandbox.google.com/github/huggingface/notebooks/blob/master/examples/text_classification_flax.ipynb#scrollTo=Mn1GdGpipfWK","metadata":{"papermill":{"duration":0.079618,"end_time":"2022-03-03T20:34:52.645839","exception":false,"start_time":"2022-03-03T20:34:52.566221","status":"completed"},"tags":[]}}]}