classification/basic_transformer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "basic_transformer.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "jAFyvoEV1WRB",
        "Ga6G28Qe1j0v",
        "RseynpJDHYtj",
        "J8vTrTw6c_OR"
      ],
      "authorship_tag": "ABX9TyP5VPOjDaBymJwBxJP5Gz8R",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/deterministic-algorithms-lab/NLP-Journey/blob/main/classification/basic_transformer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wMwOcI23zjM-"
      },
      "source": [
        "!pip install git+https://github.com/deepmind/dm-haiku\n",
        "!pip install transformers\n",
        "!pip install git+git://github.com/deepmind/optax.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cincO7Umzsu1"
      },
      "source": [
        "import haiku as hk\n",
        "import jax.numpy as jnp\n",
        "import jax\n",
        "\n",
        "from jax import jit\n",
        "from jax.random import PRNGKey\n",
        "import numpy as np"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2LD9FiTbzyjU"
      },
      "source": [
        "# Transformers-Classification Using pre-trained weights from RoBERTa"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FU2MeYIv0BHT"
      },
      "source": [
        "## Embedding Layers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I7w2UjDAzxcq"
      },
      "source": [
        "from transformers import RobertaModel\n",
        "\n",
        "class Embedding(hk.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "\n",
        "    def __call__(self, token_ids, training=False):\n",
        "        \"\"\"\n",
        "        token_ids: ints of shape (batch, n_seq)\n",
        "        \"\"\"\n",
        "        word_embeddings = self.config['pretrained']['embeddings/word_embeddings']\n",
        "        \n",
        "        flat_token_ids = jnp.reshape(token_ids, [-1])\n",
        "        \n",
        "        flat_token_embeddings = hk.Embed(vocab_size=word_embeddings.shape[0],\n",
        "                                         embed_dim=word_embeddings.shape[1],\n",
        "                                         w_init=hk.initializers.Constant(word_embeddings))(flat_token_ids)\n",
        "\n",
        "        token_embeddings = jnp.reshape(flat_token_embeddings, [token_ids.shape[0], -1, word_embeddings.shape[1]])\n",
        "        \n",
        "        embeddings = token_embeddings + PositionEmbeddings(self.config)()\n",
        "\n",
        "        embeddings = hk.LayerNorm(axis=-1,\n",
        "                                  create_scale=True,\n",
        "                                  create_offset=True,\n",
        "                                  scale_init=hk.initializers.Constant(self.config['pretrained']['embeddings/LayerNorm/gamma']),\n",
        "                                  offset_init=hk.initializers.Constant(self.config['pretrained']['embeddings/LayerNorm/beta']))(embeddings)\n",
        "        if training:\n",
        "            embeddings = hk.dropout(hk.next_rng_key(),\n",
        "                                    rate=self.config['embed_dropout_rate'],\n",
        "                                    x=embeddings)\n",
        "        \n",
        "        return embeddings"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kAoqTp6Uz7ta"
      },
      "source": [
        "class PositionEmbeddings(hk.Module):\n",
        "    \"\"\"\n",
        "    A position embedding of size [max_seq_leq, word_embedding_dim]\n",
        "    \"\"\"\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "        self.offset = 2\n",
        "\n",
        "    def __call__(self):\n",
        "        pretrained_position_embeddings = self.config['pretrained']['embeddings/position_embeddings']\n",
        "\n",
        "        position_weights = hk.get_parameter(\"position_embeddings\",\n",
        "                                            pretrained_position_embeddings.shape,\n",
        "                                            init=hk.initializers.Constant(pretrained_position_embeddings))\n",
        "        \n",
        "        start = self.offset\n",
        "        end = self.offset+self.config['max_length']\n",
        "        \n",
        "        return position_weights[start:end]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tZgxynPM0a5G"
      },
      "source": [
        "## Tokenizer and Utilities for Downloading and Extracting pre-trained weights"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Iu1HMppK0Fzh"
      },
      "source": [
        "from io import BytesIO\n",
        "from functools import lru_cache\n",
        "\n",
        "import joblib\n",
        "import requests\n",
        "\n",
        "from transformers import RobertaModel, RobertaTokenizer\n",
        "\n",
        "huggingface_roberta = RobertaModel.from_pretrained('roberta-base', output_hidden_states=True)\n",
        "\n",
        "huggingface_tokenizer = RobertaTokenizer.from_pretrained('roberta-base')\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kFw7QoSj0IRe"
      },
      "source": [
        "def postprocess_key(key):\n",
        "    key = key.replace('model/featurizer/bert/', '')\n",
        "    key = key.replace(':0', '')\n",
        "    key = key.replace('self/', '')\n",
        "    return key"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eEr4QWMj0L6O"
      },
      "source": [
        "@lru_cache()\n",
        "def get_pretrained_weights():\n",
        "    # We'll use the weight dictionary from the Roberta encoder at \n",
        "    # https://github.com/IndicoDataSolutions/finetune\n",
        "    remote_url = \"https://bendropbox.s3.amazonaws.com/roberta/roberta-model-sm-v2.jl\"\n",
        "    weights = joblib.load(BytesIO(requests.get(remote_url).content))\n",
        "\n",
        "    weights = {\n",
        "        postprocess_key(key): value\n",
        "        for key, value in weights.items()\n",
        "    }\n",
        "\n",
        "    input_embeddings = huggingface_roberta.get_input_embeddings()\n",
        "    weights['embeddings/word_embeddings'] = input_embeddings.weight.detach().numpy()\n",
        "\n",
        "    return weights\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IsfSZH7C0YY6"
      },
      "source": [
        "class Scope(object):\n",
        "    \"\"\"\n",
        "    A tiny utility to help make looking up into our dictionary cleaner.\n",
        "    There's no haiku magic here.\n",
        "    \"\"\"\n",
        "    def __init__(self, weights, prefix):\n",
        "        self.weights = weights\n",
        "        self.prefix = prefix\n",
        "\n",
        "    def __getitem__(self, key):\n",
        "        return self.weights[self.prefix + key]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BMhZBNCE0nH8"
      },
      "source": [
        "## Running the Embedding Layers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G5DPXgub0jIu"
      },
      "source": [
        "sample_text = \"This was a flower of evil.\"\n",
        "\n",
        "\n",
        "config = {'pretrained' : get_pretrained_weights(),\n",
        "          'max_length' : 512,\n",
        "          'embed_dropout_rate' : 0.1\n",
        "          }\n",
        "\n",
        "encoded = huggingface_tokenizer.batch_encode_plus([sample_text, sample_text],\n",
        "                                                  padding='max_length',\n",
        "                                                  max_length=config['max_length'])\n",
        "\n",
        "sample_tokens = encoded['input_ids']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xdEjN5KY0leO"
      },
      "source": [
        "\n",
        "\n",
        "def embed_fn(tokens, training=False) :\n",
        "    embedding = Embedding(config)(tokens)\n",
        "    return embedding\n",
        "\n",
        "rng = PRNGKey(42)\n",
        "embed = hk.transform(embed_fn, apply_rng=True)\n",
        "sample_tokens = np.asarray(sample_tokens)\n",
        "params = embed.init(rng, sample_tokens, training=False)\n",
        "embedded_tokens = jit(embed.apply)(params, rng, sample_tokens, training=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jNRxih2y0vI2"
      },
      "source": [
        "## Transformer Block"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tz6FnSL_00Md"
      },
      "source": [
        "class TransformerBlock(hk.Module):\n",
        "\n",
        "    def __init__(self, config, layer_num):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "        self.n = layer_num\n",
        "\n",
        "    def __call__(self, x, mask, training = False):\n",
        "\n",
        "        scope = Scope(\n",
        "            self.config['pretrained'], f'encoder/layer_{self.n}/'\n",
        "        )\n",
        "\n",
        "        attention_output = MultiHeadAttention(self.config,\n",
        "                                              self.n)(x, mask, training=training)\n",
        "        \n",
        "        residual = attention_output+x\n",
        "\n",
        "        attention_output = hk.LayerNorm(axis=-1,\n",
        "                                        create_scale=True,\n",
        "                                        create_offset=True,\n",
        "                                        scale_init=hk.initializers.Constant(scope['attention/output/LayerNorm/gamma']),\n",
        "                                        offset_init=hk.initializers.Constant(scope['attention/output/LayerNorm/beta']),)(residual)\n",
        "\n",
        "        mlp_output = TransformerMLP(self.config, self.n)(attention_output, training=training)\n",
        "\n",
        "        output_residual = mlp_output+attention_output\n",
        "\n",
        "        layer_output = hk.LayerNorm(axis=-1,\n",
        "                                    create_scale=True,\n",
        "                                    create_offset=True,\n",
        "                                    scale_init=hk.initializers.Constant(scope['output/LayerNorm/gamma']),\n",
        "                                    offset_init=hk.initializers.Constant(scope['output/LayerNorm/beta']))(output_residual)\n",
        "        \n",
        "        return layer_output"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bD18sx8U01d8"
      },
      "source": [
        "## Multi-Head Attention"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cYPItvV700_9"
      },
      "source": [
        "class MultiHeadAttention(hk.Module):\n",
        "    def __init__(self, config, layer_num):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "        self.n = layer_num\n",
        "\n",
        "    def _split_into_heads(self, x):\n",
        "        return jnp.reshape(x, [x.shape[0], x.shape[1], self.config['n_heads'], x.shape[2]//self.config['n_heads']])\n",
        "\n",
        "    def __call__(self, x, mask, training=False):\n",
        "        \n",
        "        scope = Scope(self.config['pretrained'], f'encoder/layer_{self.n}/attention/')\n",
        "\n",
        "        queries = hk.Linear(output_size=self.config['hidden_size'],\n",
        "                            w_init=hk.initializers.Constant(scope['query/kernel']),\n",
        "                            b_init=hk.initializers.Constant(scope['query/bias']))(x)\n",
        "        \n",
        "        keys = hk.Linear(output_size=self.config['hidden_size'],\n",
        "                         w_init=hk.initializers.Constant(scope['key/kernel']),\n",
        "                         b_init=hk.initializers.Constant(scope['key/bias']))(x)\n",
        "        \n",
        "        values = hk.Linear(output_size=self.config['hidden_size'],\n",
        "                           w_init=hk.initializers.Constant(scope['value/kernel']),\n",
        "                           b_init=hk.initializers.Constant(scope['value/bias']))(x)\n",
        "        \n",
        "        queries = self._split_into_heads(queries)\n",
        "        keys = self._split_into_heads(keys)\n",
        "        values = self._split_into_heads(values)\n",
        "\n",
        "        attention_logits = jnp.einsum('bsnh,btnh->bnst', queries, keys)\n",
        "        attention_logits /= np.sqrt(queries.shape[-1])\n",
        "\n",
        "        attention_logits += jnp.reshape(mask*-2**32, [mask.shape[0],1,1,mask.shape[1]])\n",
        "        attention_weights = jax.nn.softmax(attention_logits, axis=-1)\n",
        "        per_head_attention_output = jnp.einsum('btnh,bnst->bsnh', values, attention_weights)\n",
        "        attention_output = jnp.reshape(per_head_attention_output, [per_head_attention_output.shape[0], per_head_attention_output.shape[1], -1])\n",
        "\n",
        "        attention_output = hk.Linear(output_size=self.config['hidden_size'],\n",
        "                                     w_init=hk.initializers.Constant(scope['output/dense/kernel']),\n",
        "                                     b_init=hk.initializers.Constant(scope['output/dense/bias']))(attention_output)\n",
        "        \n",
        "        if training:\n",
        "            attention_output = hk.dropout(rng=hk.next_rng_key(),\n",
        "                                          rate=self.config['attention_drop_rate'],\n",
        "                                          x=attention_output)\n",
        "        \n",
        "        return attention_output"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xH1rORmz08oU"
      },
      "source": [
        "## Transformer MLP"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ehX4oUHY1ALU"
      },
      "source": [
        "def gelu(x):\n",
        "    return x*0.5*(1.0+jax.scipy.special.erf(x / jnp.sqrt(2.0)))\n",
        "\n",
        "class TransformerMLP(hk.Module):\n",
        "\n",
        "    def __init__(self, config, layer_num):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "        self.n = layer_num\n",
        "\n",
        "    def __call__(self, x, training=False):\n",
        "\n",
        "        scope = Scope(self.config['pretrained'], f'encoder/layer_{self.n}/')\n",
        "\n",
        "        intermediate_output = hk.Linear(output_size=self.config['intermediate_size'],\n",
        "                                        w_init=hk.initializers.Constant(scope['intermediate/dense/kernel']),\n",
        "                                        b_init=hk.initializers.Constant(scope['intermediate/dense/bias']))(x)\n",
        "\n",
        "        intermediate_output = gelu(intermediate_output)\n",
        "\n",
        "        output = hk.Linear(output_size=self.config['hidden_size'],\n",
        "                           w_init=hk.initializers.Constant(scope['output/dense/kernel']),\n",
        "                           b_init=hk.initializers.Constant(scope['output/dense/bias']))(intermediate_output)\n",
        "        \n",
        "        if training:\n",
        "            output = hk.dropout(rng=hk.next_rng_key(),\n",
        "                                rate=self.config['fully_connected_drop_rate'],\n",
        "                                x=output)\n",
        "        \n",
        "        return output"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xblfTQfw1Jh8"
      },
      "source": [
        "## Confg and Getting Features from the model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QFkccsZI1DpG"
      },
      "source": [
        "class RobertaFeaturizer(hk.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "\n",
        "    def __call__(self, token_ids, training=False):\n",
        "        x = Embedding(self.config)(token_ids, training=training)\n",
        "        mask = (token_ids==self.config['mask_id']).astype(jnp.float32)\n",
        "        for layer_num in range(self.config['n_layers']):\n",
        "            x = TransformerBlock(config, layer_num=layer_num)(x,mask,training)\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6CYk3wgQ1Ihs",
        "outputId": "99ab8e37-64b2-4f47-f640-a0a8d5f0ba0e"
      },
      "source": [
        "config = {\n",
        "          'pretrained' : config['pretrained'], \n",
        "          'max_length' : config['max_length'], \n",
        "          'embed_dropout_rate' : 0.1,\n",
        "          'fully_connected_drop_rate' : 0.1,\n",
        "          'attention_drop_rate' : 0.1,\n",
        "          'hidden_size' : 768,\n",
        "          'intermediate_size' : 3072,\n",
        "          'n_heads' : 12,\n",
        "          'n_layers' : 12,\n",
        "          'mask_id' : 1,\n",
        "          'weight_stddev' : 0.02,\n",
        "          \n",
        "          'n_classes' : 2,\n",
        "          'classifier_drop_rate' : 0.1,\n",
        "          'learning_rate' : 1e-5,\n",
        "          'max_grad_norm' : 1.0,\n",
        "          'l2' : 0.1,\n",
        "          'n_epochs' : 5,\n",
        "          'batch_size' : 4\n",
        "          }\n",
        "\n",
        "def featurizer_fn(tokens, training=False):\n",
        "    contextual_embeddings = RobertaFeaturizer(config)(tokens, training=training)\n",
        "    return contextual_embeddings\n",
        "\n",
        "rng = PRNGKey(42)\n",
        "roberta = hk.transform(featurizer_fn)\n",
        "sample_tokens = np.asarray(sample_tokens)\n",
        "params = roberta.init(rng, sample_tokens, training=False)\n",
        "contextual_embeddings = jit(roberta.apply)(params, rng, sample_tokens)\n",
        "print(contextual_embeddings.shape)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "(2, 512, 768)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WtJaXolI-IL2"
      },
      "source": [
        "## Getting Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1OQUUQ9B-MEu"
      },
      "source": [
        "import tensorflow_datasets as tfds\n",
        "\n",
        "def load_dataset(split, training, batch_size, n_epochs=1, n_examples=None):\n",
        "    ds = tfds.load(\"imdb_reviews\", \n",
        "                   split=f\"{split}[:{n_examples}]\").cache().repeat(n_epochs)\n",
        "    \n",
        "    if training:\n",
        "        ds = ds.shuffle(10*batch_size, seed=0)\n",
        "    \n",
        "    ds = ds.batch(batch_size)\n",
        "\n",
        "    return tfds.as_numpy(ds)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vb7W-Jf-_KfI"
      },
      "source": [
        "n_examples = 25000\n",
        "train = load_dataset('train', training=True, batch_size=4, n_epochs=config['n_epochs'],n_examples=n_examples)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lrd1yBJwAsEt"
      },
      "source": [
        "def encode_batch(batch_text):\n",
        "    batch_text = [\n",
        "                  text[:512].decode('utf-8') if isinstance(text, bytes) else text[:512]\n",
        "                  for text in batch_text\n",
        "    ]\n",
        "    \n",
        "    token_ids = huggingface_tokenizer.batch_encode_plus(batch_text,\n",
        "                                                        padding='max_length',\n",
        "                                                        max_length=config['max_length'],\n",
        "                                                        )['input_ids']\n",
        "    \n",
        "    return np.asarray(token_ids)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jAFyvoEV1WRB"
      },
      "source": [
        "## The classifier"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ce1B55SL1R2z"
      },
      "source": [
        "class RobertaClassifier(hk.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "\n",
        "    def __call__(self, token_ids, training=False):\n",
        "        sequence_features = RobertaFeaturizer(self.config)(token_ids=token_ids, training=training)\n",
        "\n",
        "        clf_state = sequence_features[:,0,:]\n",
        "\n",
        "        if training:\n",
        "            clf_state = hk.dropout(rng=hk.next_rng_key(),\n",
        "                                   rate=self.config['classifier_drop_rate'],\n",
        "                                   x=clf_state)\n",
        "        \n",
        "        clf_logits = hk.Linear(output_size=self.config['n_classes'],\n",
        "                               w_init=hk.initializers.TruncatedNormal(self.config['weight_stddev']))(clf_state)\n",
        "\n",
        "        return clf_logits"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ga6G28Qe1j0v"
      },
      "source": [
        "## Running the Classifier"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K7bZRSRH1o_R"
      },
      "source": [
        "def roberta_classification_fn(batch_token_ids, training):\n",
        "    logits = RobertaClassifier(config)(batch_token_ids, training=training)\n",
        "    return logits\n",
        "\n",
        "rng = jax.random.PRNGKey(42)\n",
        "roberta_classifier = hk.transform(roberta_classification_fn)                    \n",
        "\n",
        "params = roberta_classifier.init(rng, \n",
        "                                 batch_token_ids=encode_batch(['sample sentence', 'Another one!']),\n",
        "                                 training=True)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0isUhSMSD-ZR"
      },
      "source": [
        "```roberta_classifier.init()``` and ```roberta_classifier.apply()``` are pure functions now. So, can be composed to gether and used with other functions. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lcbdVKceDlyl"
      },
      "source": [
        "def loss(params, rng, batch_token_ids, batch_labels):\n",
        "    logits = roberta_classifier.apply(params, rng, batch_token_ids, training=True)\n",
        "    labels = hk.one_hot(batch_labels, config['n_classes'])\n",
        "    softmax_xent = -jnp.sum(labels*jax.nn.log_softmax(logits))\n",
        "    softmax_xent /= labels.shape[0]\n",
        "    return softmax_xent\n",
        "\n",
        "@jax.jit\n",
        "def accuracy(params, rng, batch_token_ids, batch_labels):\n",
        "    logits = roberta_classifier.apply(params, rng, batch_token_ids, training=False)\n",
        "    return jnp.mean(jnp.argmax(logits, axis=-1)==batch_labels)\n",
        "\n",
        "@jax.jit\n",
        "def update(params, rng, opt_state, batch_token_ids, batch_labels):\n",
        "    batch_loss, grad = jax.value_and_grad(loss)(params, rng, batch_token_ids, batch_labels)\n",
        "    updates, opt_state = opt.update(grad, opt_state)\n",
        "    new_params = optax.apply_updates(params, updates)\n",
        "    return new_params, opt_state, batch_loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RseynpJDHYtj"
      },
      "source": [
        "## Defining Learning rate scheduler and Optimizer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xjHY4tKiHX8A"
      },
      "source": [
        "import optax"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OJJ1-bzCYYru"
      },
      "source": [
        "The below way of defining a functionality allows you to tie together namespaces with functions.(Or \"wrap\" a function in a namespace consisting of variables defined in the outer function).\n",
        "\n",
        "Here, ```warmup_percentage``` and ```total_steps``` act as if they were variables in a class with a function ```lr_schedule()```. The ```lr_schedule()``` function can access them, freely. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lvGKeVarIE2K"
      },
      "source": [
        "def make_lr_schedule(warmup_percentage, total_steps):\n",
        "    \n",
        "    def lr_schedule(step):\n",
        "        percent_complete = step/total_steps\n",
        "        \n",
        "        #0 or 1 based on whether we are before peak\n",
        "        before_peak = jax.lax.convert_element_type((percent_complete<=warmup_percentage),\n",
        "                                                   np.float32)\n",
        "        #Factor for scaling learning rate\n",
        "        scale = ( before_peak*(percent_complete/warmup_percentage)\n",
        "                + (1-before_peak) ) * (1-percent_complete)\n",
        "        \n",
        "        return scale\n",
        "    \n",
        "    return lr_schedule"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HI5QJ_lLXeg2"
      },
      "source": [
        "total_steps = config['n_epochs']*(n_examples//config['batch_size'])\n",
        "\n",
        "lr_schedule = make_lr_schedule(warmup_percentage=0.1, total_steps=total_steps)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MclJF-XRX1Mi"
      },
      "source": [
        "opt = optax.chain(\n",
        "    optax.clip_by_global_norm(config['max_grad_norm']),\n",
        "    optax.adam(learning_rate=config['learning_rate']),\n",
        "    optax.scale_by_schedule(lr_schedule),\n",
        ")\n",
        "opt_state = opt.init(params)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J8vTrTw6c_OR"
      },
      "source": [
        "## Utility for Measuring Performance"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VZ-UokInc-jh"
      },
      "source": [
        "def measure_current_performance(params, n_examples=None, splits=('train', 'test')):\n",
        "    if 'train' in splits:\n",
        "        train_eval = load_dataset('train', training=False, batch_size=25, n_examples=n_examples)\n",
        "\n",
        "        train_accuracy = np.mean([accuracy(params, rng, \n",
        "                                          encode_batch(train_eval_batch['text']),\n",
        "                                          train_eval_batch['label'])\n",
        "                                for train_eval_batch in train_eval])\n",
        "        \n",
        "        print(f\"\\t Train validation acc: {train_accuracy:.3f}\")\n",
        "\n",
        "    if 'test' in splits:\n",
        "        test_eval = load_dataset('test', training=False, batch_size=25, n_examples=n_examples)\n",
        "\n",
        "        test_accuracy  = np.mean([accuracy(params, rng, \n",
        "                                           encode_batch(test_eval_batch['text']),\n",
        "                                           test_eval_batch['label'])\n",
        "                                  for test_eval_batch in test_eval])\n",
        "    \n",
        "    print(f\"\\t Test validation accuracy: {test_accuracy:.3f}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mM_37Nq_fArJ"
      },
      "source": [
        "## Training Loop"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ygZLLAAJnmPN"
      },
      "source": [
        "### For running on a different dataset : "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u3HPLdlnrZTB"
      },
      "source": [
        "**In the cell below :**\n",
        "\n",
        "* Change Line 1 to enumerate any data set returning batches of actual text(can have emojis too), with their integer labels. For example, ```train_batch['text']``` can be a list(or any other iterable) ```['My name is Jeevesh.', 'I live at your house.']``` with ```train_batch['labels']``` as another list ```[1,2]```.\n",
        "\n",
        "* Change ```n_classes``` in config.\n",
        "\n",
        "* Change tokenizer/provide vocabulary to add new tokens for additional languages, using ```huggingface_tokenizer.add_tokens(<list of new tokens>)``` .\n",
        "\n",
        "* Rest remains same."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7XJXY72DfCd6"
      },
      "source": [
        "for step, train_batch in enumerate(train):\n",
        "    \n",
        "    if step%100==0:\n",
        "        print(f'[Step {step}]')\n",
        "    if step%1000==0 and step!=0:\n",
        "        measure_current_performance(params, n_examples=100)\n",
        "    print(\"Here\")\n",
        "    batch_token_ids = encode_batch(train_batch['text'])\n",
        "    batch_labels = train_batch['label']\n",
        "    params, opt_state, batch_loss = update(params, rng, opt_state, batch_token_ids, batch_labels)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}