rnnLearn.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyO9W/B+A1RRXuFgvHTTZ7Zp",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/LeoMcBills/ML_ChatGPT_fromScratch/blob/main/rnnLearn.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "scfd-y2L3yET"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import os\n",
        "import time"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "path_to_file = tf.keras.utils.get_file('shakespeare.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Fsd_cH364YEl",
        "outputId": "d6697bcd-61a2-4ca1-ab2e-8fc39a76b9ef"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading data from https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt\n",
            "1115394/1115394 [==============================] - 1s 1us/step\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Read, then decode for py2 compat.\n",
        "text = open(path_to_file, 'rb').read().decode(encoding='utf-8')\n",
        "# length of text is the number of characters in it\n",
        "print(f'Length of text: {len(text)} characters')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Wj_IhabZ4dsk",
        "outputId": "d131a22c-4897-4420-db11-8bec24632e0c"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Length of text: 1115394 characters\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Take a look at the first 250 characters in text\n",
        "print(text[:250])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2nVWa5CA4jBX",
        "outputId": "1080f3d3-1e95-4db9-caae-a03ade51893c"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "First Citizen:\n",
            "Before we proceed any further, hear me speak.\n",
            "\n",
            "All:\n",
            "Speak, speak.\n",
            "\n",
            "First Citizen:\n",
            "You are all resolved rather to die than to famish?\n",
            "\n",
            "All:\n",
            "Resolved. resolved.\n",
            "\n",
            "First Citizen:\n",
            "First, you know Caius Marcius is chief enemy to the people.\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# The unique characters in the file\n",
        "vocab = sorted(set(text))\n",
        "print(f'{len(vocab)} unique characters')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JYsnzlXc4n1B",
        "outputId": "fc7ea4ca-d943-49f1-bb1b-9ef464825e5b"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "65 unique characters\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "example_texts = ['abcdefg', 'xyz']\n",
        "chars = tf.strings.unicode_split(example_texts, input_encoding='UTF-8')\n",
        "chars"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aBYvPtkW4uDU",
        "outputId": "f64a9a27-2990-43eb-bb32-7857cca3173f"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.RaggedTensor [[b'a', b'b', b'c', b'd', b'e', b'f', b'g'], [b'x', b'y', b'z']]>"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ids_from_chars = tf.keras.layers.StringLookup(\n",
        "    vocabulary=list(vocab), mask_token=None)"
      ],
      "metadata": {
        "id": "dEJPUzok46yf"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ids = ids_from_chars(chars)\n",
        "ids"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EgkNIUlt5A4U",
        "outputId": "6c8ab662-e106-45ca-b672-a1603603ec4b"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.RaggedTensor [[40, 41, 42, 43, 44, 45, 46], [63, 64, 65]]>"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chars_from_ids = tf.keras.layers.StringLookup(\n",
        "    vocabulary=ids_from_chars.get_vocabulary(), invert=True, mask_token=None)"
      ],
      "metadata": {
        "id": "oHr28Yi15Oca"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chars = chars_from_ids(ids)\n",
        "chars"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jZNJYVj55Wqg",
        "outputId": "ea07e7c7-736c-485a-abb4-bd8479cfa9f9"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.RaggedTensor [[b'a', b'b', b'c', b'd', b'e', b'f', b'g'], [b'x', b'y', b'z']]>"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tf.strings.reduce_join(chars, axis=-1).numpy()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wdN-phHR5akM",
        "outputId": "5b9ca15d-62b9-489a-8a83-c6695f884b0c"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([b'abcdefg', b'xyz'], dtype=object)"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def text_from_ids(ids):\n",
        "  return tf.strings.reduce_join(chars_from_ids(ids), axis=-1)"
      ],
      "metadata": {
        "id": "KEme6zvU5eZX"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "all_ids = ids_from_chars(tf.strings.unicode_split(text, 'UTF-8'))\n",
        "all_ids"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q583eBir5nZ5",
        "outputId": "35c355cc-c4e5-4957-a197-6dbf9f31368a"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: shape=(1115394,), dtype=int64, numpy=array([19, 48, 57, ..., 46,  9,  1])>"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ids_dataset = tf.data.Dataset.from_tensor_slices(all_ids)"
      ],
      "metadata": {
        "id": "dQy6DHe05sXg"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for ids in ids_dataset.take(10):\n",
        "    print(chars_from_ids(ids).numpy().decode('utf-8'))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UuHIaKEk5w-A",
        "outputId": "7c273767-8857-446e-8a47-1f75407eb4b5"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "F\n",
            "i\n",
            "r\n",
            "s\n",
            "t\n",
            " \n",
            "C\n",
            "i\n",
            "t\n",
            "i\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "seq_length = 100"
      ],
      "metadata": {
        "id": "g80IwT6O50JN"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sequences = ids_dataset.batch(seq_length+1, drop_remainder=True)\n",
        "\n",
        "for seq in sequences.take(1):\n",
        "  print(chars_from_ids(seq))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-lPlO6_Z55Qv",
        "outputId": "a3387660-2127-425a-c2fe-a50fde3bba37"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tf.Tensor(\n",
            "[b'F' b'i' b'r' b's' b't' b' ' b'C' b'i' b't' b'i' b'z' b'e' b'n' b':'\n",
            " b'\\n' b'B' b'e' b'f' b'o' b'r' b'e' b' ' b'w' b'e' b' ' b'p' b'r' b'o'\n",
            " b'c' b'e' b'e' b'd' b' ' b'a' b'n' b'y' b' ' b'f' b'u' b'r' b't' b'h'\n",
            " b'e' b'r' b',' b' ' b'h' b'e' b'a' b'r' b' ' b'm' b'e' b' ' b's' b'p'\n",
            " b'e' b'a' b'k' b'.' b'\\n' b'\\n' b'A' b'l' b'l' b':' b'\\n' b'S' b'p' b'e'\n",
            " b'a' b'k' b',' b' ' b's' b'p' b'e' b'a' b'k' b'.' b'\\n' b'\\n' b'F' b'i'\n",
            " b'r' b's' b't' b' ' b'C' b'i' b't' b'i' b'z' b'e' b'n' b':' b'\\n' b'Y'\n",
            " b'o' b'u' b' '], shape=(101,), dtype=string)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for seq in sequences.take(5):\n",
        "  print(text_from_ids(seq).numpy())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ILrzwWQl59MI",
        "outputId": "e1dac5af-f76f-4ff8-dae2-5a1aa94bc57e"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "b'First Citizen:\\nBefore we proceed any further, hear me speak.\\n\\nAll:\\nSpeak, speak.\\n\\nFirst Citizen:\\nYou '\n",
            "b'are all resolved rather to die than to famish?\\n\\nAll:\\nResolved. resolved.\\n\\nFirst Citizen:\\nFirst, you k'\n",
            "b\"now Caius Marcius is chief enemy to the people.\\n\\nAll:\\nWe know't, we know't.\\n\\nFirst Citizen:\\nLet us ki\"\n",
            "b\"ll him, and we'll have corn at our own price.\\nIs't a verdict?\\n\\nAll:\\nNo more talking on't; let it be d\"\n",
            "b'one: away, away!\\n\\nSecond Citizen:\\nOne word, good citizens.\\n\\nFirst Citizen:\\nWe are accounted poor citi'\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def split_input_target(sequence):\n",
        "    input_text = sequence[:-1]\n",
        "    target_text = sequence[1:]\n",
        "    return input_text, target_text"
      ],
      "metadata": {
        "id": "_5S3JtE16Dui"
      },
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "split_input_target(list(\"Tensorflow\"))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8iT9TiPn6JfY",
        "outputId": "ef215aee-e83c-49b8-93c5-20561c3396a7"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(['T', 'e', 'n', 's', 'o', 'r', 'f', 'l', 'o'],\n",
              " ['e', 'n', 's', 'o', 'r', 'f', 'l', 'o', 'w'])"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = sequences.map(split_input_target)"
      ],
      "metadata": {
        "id": "_S5LSjSo6O9r"
      },
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for input_example, target_example in dataset.take(1):\n",
        "    print(\"Input :\", text_from_ids(input_example).numpy())\n",
        "    print(\"Target:\", text_from_ids(target_example).numpy())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "T1QMOqMA6SN5",
        "outputId": "e9b47e76-347c-4b3a-d52b-3b9d2dc6a6b3"
      },
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input : b'First Citizen:\\nBefore we proceed any further, hear me speak.\\n\\nAll:\\nSpeak, speak.\\n\\nFirst Citizen:\\nYou'\n",
            "Target: b'irst Citizen:\\nBefore we proceed any further, hear me speak.\\n\\nAll:\\nSpeak, speak.\\n\\nFirst Citizen:\\nYou '\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Batch size\n",
        "BATCH_SIZE = 64\n",
        "\n",
        "# Buffer size to shuffle the dataset\n",
        "# (TF data is designed to work with possibly infinite sequences,\n",
        "# so it doesn't attempt to shuffle the entire sequence in memory. Instead,\n",
        "# it maintains a buffer in which it shuffles elements).\n",
        "BUFFER_SIZE = 10000\n",
        "\n",
        "dataset = (\n",
        "    dataset\n",
        "    .shuffle(BUFFER_SIZE)\n",
        "    .batch(BATCH_SIZE, drop_remainder=True)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE))\n",
        "\n",
        "dataset"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BkRz_MWz6VDo",
        "outputId": "7d871ae0-1fca-4794-87a8-67439f58bcd2"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<_PrefetchDataset element_spec=(TensorSpec(shape=(64, 100), dtype=tf.int64, name=None), TensorSpec(shape=(64, 100), dtype=tf.int64, name=None))>"
            ]
          },
          "metadata": {},
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Length of the vocabulary in StringLookup Layer\n",
        "vocab_size = len(ids_from_chars.get_vocabulary())\n",
        "\n",
        "# The embedding dimension\n",
        "embedding_dim = 256\n",
        "\n",
        "# Number of RNN units\n",
        "rnn_units = 1024"
      ],
      "metadata": {
        "id": "1Aai0byK6cR0"
      },
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class MyModel(tf.keras.Model):\n",
        "  def __init__(self, vocab_size, embedding_dim, rnn_units):\n",
        "    super().__init__(self)\n",
        "    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n",
        "    self.gru = tf.keras.layers.GRU(rnn_units,\n",
        "                                   return_sequences=True,\n",
        "                                   return_state=True)\n",
        "    self.dense = tf.keras.layers.Dense(vocab_size)\n",
        "\n",
        "  def call(self, inputs, states=None, return_state=False, training=False):\n",
        "    x = inputs\n",
        "    x = self.embedding(x, training=training)\n",
        "    if states is None:\n",
        "      states = self.gru.get_initial_state(x)\n",
        "    x, states = self.gru(x, initial_state=states, training=training)\n",
        "    x = self.dense(x, training=training)\n",
        "\n",
        "    if return_state:\n",
        "      return x, states\n",
        "    else:\n",
        "      return x"
      ],
      "metadata": {
        "id": "YrJzaOPd6hEw"
      },
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = MyModel(\n",
        "    vocab_size=vocab_size,\n",
        "    embedding_dim=embedding_dim,\n",
        "    rnn_units=rnn_units)"
      ],
      "metadata": {
        "id": "Y6pvrPdY6kGF"
      },
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for input_example_batch, target_example_batch in dataset.take(1):\n",
        "    example_batch_predictions = model(input_example_batch)\n",
        "    print(example_batch_predictions.shape, \"# (batch_size, sequence_length, vocab_size)\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DPZcOMxp6njD",
        "outputId": "58e6b11e-164b-40df-c44e-654daa6bf4b5"
      },
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(64, 100, 66) # (batch_size, sequence_length, vocab_size)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model.summary()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bfqRiyMd6uJ9",
        "outputId": "08cd8f92-2e85-4acf-be5f-b0a104cb0954"
      },
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"my_model\"\n",
            "_________________________________________________________________\n",
            " Layer (type)                Output Shape              Param #   \n",
            "=================================================================\n",
            " embedding (Embedding)       multiple                  16896     \n",
            "                                                                 \n",
            " gru (GRU)                   multiple                  3938304   \n",
            "                                                                 \n",
            " dense (Dense)               multiple                  67650     \n",
            "                                                                 \n",
            "=================================================================\n",
            "Total params: 4022850 (15.35 MB)\n",
            "Trainable params: 4022850 (15.35 MB)\n",
            "Non-trainable params: 0 (0.00 Byte)\n",
            "_________________________________________________________________\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "sampled_indices = tf.random.categorical(example_batch_predictions[0], num_samples=1)\n",
        "sampled_indices = tf.squeeze(sampled_indices, axis=-1).numpy()"
      ],
      "metadata": {
        "id": "cE6fSTUD6xeQ"
      },
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sampled_indices"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "toZoC7ry62Mo",
        "outputId": "875e7707-a4dc-4745-b2aa-8ac3ed855855"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([59, 29, 46, 54, 14, 64, 33, 48, 15, 29,  7, 33, 56, 22, 44, 11, 43,\n",
              "        1, 65, 18, 32, 21, 22, 23, 34, 23, 52, 56, 41, 48, 48,  6, 56, 50,\n",
              "        4, 34, 25,  0, 34, 22, 11,  1, 20, 32,  5, 62, 48, 30, 60, 60, 51,\n",
              "       24, 19,  9, 55, 19, 14,  2, 50, 51, 65, 43, 31, 36, 58, 25, 44, 50,\n",
              "       48, 15, 63, 12, 21, 24, 54, 30, 25,  1,  5, 65, 10, 61, 24, 17, 27,\n",
              "       14, 28,  9, 18, 40, 43,  9, 60, 11, 30, 27,  6, 11, 65, 13])"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Input:\\n\", text_from_ids(input_example_batch[0]).numpy())\n",
        "print()\n",
        "print(\"Next Char Predictions:\\n\", text_from_ids(sampled_indices).numpy())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "L5eJvyhb67kS",
        "outputId": "e12d2f21-2940-4289-f678-7b2a86fdc66b"
      },
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input:\n",
            " b\"on thee; and thou art too full\\nOf the wars' surfeits, to go rove with one\\nThat's yet unbruised: brin\"\n",
            "\n",
            "Next Char Predictions:\n",
            " b\"tPgoAyTiBP,TqIe:d\\nzESHIJUJmqbii'qk$UL[UNK]UI:\\nGS&wiQuulKF.pFA klzdRWsLekiBx;HKoQL\\n&z3vKDNAO.Ead.u:QN':z?\"\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True)"
      ],
      "metadata": {
        "id": "YJcEv-d26_8C"
      },
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "example_batch_mean_loss = loss(target_example_batch, example_batch_predictions)\n",
        "print(\"Prediction shape: \", example_batch_predictions.shape, \" # (batch_size, sequence_length, vocab_size)\")\n",
        "print(\"Mean loss:        \", example_batch_mean_loss)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YPePSKrq7IRb",
        "outputId": "a6ceee9f-a95c-48ca-fbcb-9724bc8d5967"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prediction shape:  (64, 100, 66)  # (batch_size, sequence_length, vocab_size)\n",
            "Mean loss:         tf.Tensor(4.1907787, shape=(), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tf.exp(example_batch_mean_loss).numpy()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "g8vDL3mX7Lk-",
        "outputId": "f0f3aa94-7d49-4045-937f-797a16623d88"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "66.07423"
            ]
          },
          "metadata": {},
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model.compile(optimizer='adam', loss=loss)"
      ],
      "metadata": {
        "id": "Zb3FqBwK7RTr"
      },
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Directory where the checkpoints will be saved\n",
        "checkpoint_dir = './training_checkpoints'\n",
        "# Name of the checkpoint files\n",
        "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt_{epoch}\")\n",
        "\n",
        "checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n",
        "    filepath=checkpoint_prefix,\n",
        "    save_weights_only=True)"
      ],
      "metadata": {
        "id": "aR37xhHo8Fnh"
      },
      "execution_count": 42,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "EPOCHS = 20"
      ],
      "metadata": {
        "id": "MMKSRzRw8Jli"
      },
      "execution_count": 43,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "history = model.fit(dataset, epochs=EPOCHS, callbacks=[checkpoint_callback])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_w1z9BpS8NIe",
        "outputId": "2ca6da2b-dbd0-4452-9f69-c53bea2ba665"
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/20\n",
            "172/172 [==============================] - 16s 55ms/step - loss: 2.7250\n",
            "Epoch 2/20\n",
            "172/172 [==============================] - 12s 55ms/step - loss: 1.9915\n",
            "Epoch 3/20\n",
            "172/172 [==============================] - 11s 55ms/step - loss: 1.7156\n",
            "Epoch 4/20\n",
            "172/172 [==============================] - 11s 56ms/step - loss: 1.5533\n",
            "Epoch 5/20\n",
            "172/172 [==============================] - 12s 56ms/step - loss: 1.4532\n",
            "Epoch 6/20\n",
            "172/172 [==============================] - 12s 56ms/step - loss: 1.3843\n",
            "Epoch 7/20\n",
            "172/172 [==============================] - 12s 57ms/step - loss: 1.3331\n",
            "Epoch 8/20\n",
            "172/172 [==============================] - 11s 57ms/step - loss: 1.2876\n",
            "Epoch 9/20\n",
            "172/172 [==============================] - 11s 57ms/step - loss: 1.2468\n",
            "Epoch 10/20\n",
            "172/172 [==============================] - 12s 58ms/step - loss: 1.2078\n",
            "Epoch 11/20\n",
            "172/172 [==============================] - 12s 58ms/step - loss: 1.1686\n",
            "Epoch 12/20\n",
            "172/172 [==============================] - 12s 58ms/step - loss: 1.1283\n",
            "Epoch 13/20\n",
            "172/172 [==============================] - 12s 58ms/step - loss: 1.0854\n",
            "Epoch 14/20\n",
            "172/172 [==============================] - 12s 62ms/step - loss: 1.0400\n",
            "Epoch 15/20\n",
            "172/172 [==============================] - 12s 59ms/step - loss: 0.9921\n",
            "Epoch 16/20\n",
            "172/172 [==============================] - 14s 62ms/step - loss: 0.9426\n",
            "Epoch 17/20\n",
            "172/172 [==============================] - 13s 58ms/step - loss: 0.8895\n",
            "Epoch 18/20\n",
            "172/172 [==============================] - 12s 59ms/step - loss: 0.8373\n",
            "Epoch 19/20\n",
            "172/172 [==============================] - 13s 59ms/step - loss: 0.7862\n",
            "Epoch 20/20\n",
            "172/172 [==============================] - 13s 59ms/step - loss: 0.7375\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class OneStep(tf.keras.Model):\n",
        "  def __init__(self, model, chars_from_ids, ids_from_chars, temperature=1.0):\n",
        "    super().__init__()\n",
        "    self.temperature = temperature\n",
        "    self.model = model\n",
        "    self.chars_from_ids = chars_from_ids\n",
        "    self.ids_from_chars = ids_from_chars\n",
        "\n",
        "    # Create a mask to prevent \"[UNK]\" from being generated.\n",
        "    skip_ids = self.ids_from_chars(['[UNK]'])[:, None]\n",
        "    sparse_mask = tf.SparseTensor(\n",
        "        # Put a -inf at each bad index.\n",
        "        values=[-float('inf')]*len(skip_ids),\n",
        "        indices=skip_ids,\n",
        "        # Match the shape to the vocabulary\n",
        "        dense_shape=[len(ids_from_chars.get_vocabulary())])\n",
        "    self.prediction_mask = tf.sparse.to_dense(sparse_mask)\n",
        "\n",
        "  @tf.function\n",
        "  def generate_one_step(self, inputs, states=None):\n",
        "    # Convert strings to token IDs.\n",
        "    input_chars = tf.strings.unicode_split(inputs, 'UTF-8')\n",
        "    input_ids = self.ids_from_chars(input_chars).to_tensor()\n",
        "\n",
        "    # Run the model.\n",
        "    # predicted_logits.shape is [batch, char, next_char_logits]\n",
        "    predicted_logits, states = self.model(inputs=input_ids, states=states,\n",
        "                                          return_state=True)\n",
        "    # Only use the last prediction.\n",
        "    predicted_logits = predicted_logits[:, -1, :]\n",
        "    predicted_logits = predicted_logits/self.temperature\n",
        "    # Apply the prediction mask: prevent \"[UNK]\" from being generated.\n",
        "    predicted_logits = predicted_logits + self.prediction_mask\n",
        "\n",
        "    # Sample the output logits to generate token IDs.\n",
        "    predicted_ids = tf.random.categorical(predicted_logits, num_samples=1)\n",
        "    predicted_ids = tf.squeeze(predicted_ids, axis=-1)\n",
        "\n",
        "    # Convert from token ids to characters\n",
        "    predicted_chars = self.chars_from_ids(predicted_ids)\n",
        "\n",
        "    # Return the characters and model state.\n",
        "    return predicted_chars, states"
      ],
      "metadata": {
        "id": "CCckUdXT8lk9"
      },
      "execution_count": 45,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "one_step_model = OneStep(model, chars_from_ids, ids_from_chars)"
      ],
      "metadata": {
        "id": "fFB3488o9RlD"
      },
      "execution_count": 46,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "start = time.time()\n",
        "states = None\n",
        "next_char = tf.constant(['ROMEO:'])\n",
        "result = [next_char]\n",
        "\n",
        "for n in range(1000):\n",
        "  next_char, states = one_step_model.generate_one_step(next_char, states=states)\n",
        "  result.append(next_char)\n",
        "\n",
        "result = tf.strings.join(result)\n",
        "end = time.time()\n",
        "print(result[0].numpy().decode('utf-8'), '\\n\\n' + '_'*80)\n",
        "print('\\nRun time:', end - start)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GjSlai1W9Y5I",
        "outputId": "df3edda4-b9d7-4fe1-f62c-5c817b7f11b6"
      },
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ROMEO:\n",
            "No end of the country, chargest thou became myself?\n",
            "Aland, poor man, you are now 'I' the should be of both?\n",
            "\n",
            "KING HENRY VI:\n",
            "What, train! Come, good morrow; gentle my lords,\n",
            "Both have knew him, not pity of thee;\n",
            "And somieth the tackle with a duzed ill.\n",
            "\n",
            "Messenger:\n",
            "Almost can contagons to bring him\n",
            "As I am out at once, in two in one.\n",
            "\n",
            "BENVOLIO:\n",
            "So clied their mother wrings in spitted body:\n",
            "His time o' the people, seeking slaughter,\n",
            "There's blemedict, at once rough in warmony.\n",
            "\n",
            "Clown:\n",
            "Never my master, Baccas! O thou wilt not be spent for,\n",
            "And ask thou blentments, being revenge my ground.\n",
            "\n",
            "DUKE OF YORK:\n",
            "Is Clarence: ah. What, do I the heaven will revenge\n",
            "This way to cheque to the septent! Titus Lartius, it how but hell!\n",
            "I prisoner than you to-day? but I wish her met-our,\n",
            "I do not know. Your beauteous dewitchs\n",
            "Be endured all hearts from them, 'tis too hate\n",
            "both took a fool.\n",
            "\n",
            "First Gentleman:\n",
            "Then say Hone touch'd with man? Who begin\n",
            "Would throw it from the water of my fine,\n",
            "For resign thre \n",
            "\n",
            "________________________________________________________________________________\n",
            "\n",
            "Run time: 3.789961576461792\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "start = time.time()\n",
        "states = None\n",
        "next_char = tf.constant(['ROMEO:', 'ROMEO:', 'ROMEO:', 'ROMEO:', 'ROMEO:'])\n",
        "result = [next_char]\n",
        "\n",
        "for n in range(1000):\n",
        "  next_char, states = one_step_model.generate_one_step(next_char, states=states)\n",
        "  result.append(next_char)\n",
        "\n",
        "result = tf.strings.join(result)\n",
        "end = time.time()\n",
        "print(result, '\\n\\n' + '_'*80)\n",
        "print('\\nRun time:', end - start)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VyAqEjd49d7F",
        "outputId": "69cfc0cf-4efc-4ac3-f96f-2eb508a33308"
      },
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tf.Tensor(\n",
            "[b\"ROMEO:\\nNo, no, mine weeps from the Lord's righ.\\n\\nKATHARINA:\\nIs't like this news from him?\\n\\nDUKE OF YORK:\\nThou woes, thy crown here enjoy'd:\\nI'll not be seen to do good doubt,\\nOr in propar'd galmant canon but became.\\n\\nSecond Watchman:\\n'Tis well; and she is Christopority.\\nMaking no other name! what tranco bred wilt unquiet of thy gown\\nThe hest of disordenate before deliver'd\\nTimes! I hear those cries aid propers:\\nAnd were they talk'd with country grows to me.\\n\\nHORTENSIO:\\nNo longer Earl of Warwick, Caius Marcius,\\nRepual from thy sickly hate him first,\\nAnd while it was but a ringer wing.\\n\\nMERCUTIO:\\n'Twas I revenge: be brothers, when I was wareward love.\\n\\nKATHARINA:\\nNo doubt, the news, hold-true, there is sickle, since\\nDidst peep while the dead bodes:\\nThis sentence then to prizid the grace!\\n\\nCLARENCE:\\nAlas fool, and break out along!\\nCousin on my sweet Montague, Signior Tapest.\\nThis citizens and good Camillo,\\nPut not your country, and no let him be Murder'd him;\\nAnd do not one of the king dead, th\"\n",
            " b\"ROMEO:\\nNo more rememberent gantorable and fall?\\n\\nKeNNII:\\nAway within, there's none else blood which nays of joy;\\nAnd so she stamp'ed, and will maintain.\\n\\nPETRUCHIO:\\nA hundred without put better: then comes the fire,\\nThe Volsces I rather, this news so hold me,\\nAnd I unperious tears.\\nNow Saint Leontes: I here renounce\\nI'll turn the death is most likely in a sint\\nThe great comforts, let them please to blamn the bed\\nIn those foes thou but man, deposed;\\nFor Yet, I pray.\\n\\nPETRUCHIO:\\nCome, well, well anon. Come, sir?\\n\\nCAils:\\nYour broid's noble and\\nextrumity. The bear half lived the burthen gamb;\\nAnd you were sworn to came we to denial,\\nAddect the people and the Cotis,\\nTo credit, Our presence? Thou art deceived:\\nLet him dead; and what I apply them against me\\nSome heaviness in dembss on eyes;\\nSign'd fortune making me retreat,\\nNor shall anon we stand withwain at home;\\nAnd do not so, for whom we set himself\\nInto his fathers, folkness; for, as thou art,\\nCame aller by passing eegs our brothers\\nMore in hi\"\n",
            " b\"ROMEO:\\nAnd why seem warfind between twrivener betreated\\nThe crown with former likewarches:\\nWho bids me ask your sentence? musters are now;\\nThings alonght awhile.\\nA plot her reworth both?\\nHave worsed! Because they are not pilgrimage.\\n\\nBRUTUS:\\nHe calls, and so fair a name,\\nAnd do not call the world guess then\\nAboun a child, king in sadness, were as they\\nAmpusing so soft-day with the jointuse\\ngentlewomb, that we have nothing else:\\nHave yond my speeches with your bust?\\n\\nGLOUCESTER:\\nWhat is that way Had Henry fine the world;\\nLeast to my Gloucester's death, if not, nor there be spent.\\nWhich on thy tast, go son, and know not when I do\\nBut marry, shall I.\\nShame perish, though he me part; look thee, revenge,\\nAnd to be wash'd with thy eyesire he mistack'd.\\n\\nISABELLA:\\nGentle my dishonour; but, fear well\\nOn this most guilty light to me and\\nTybalt defend me sound your tongue.\\n\\nLEONTES:\\nWhat's the news who? by this pant a flenty\\nAs blessed man in pain deward,\\nThe very emperity of a town is All.\\n\\nGONZALO:\\n\"\n",
            " b\"ROMEO:\\nNow be deep and earth Servive, to be true.\\n\\nCAPULET:\\nGo, give what thou the king?\\n\\nDUKE OF YORK:\\nI'll not be good unto him.\\n\\nCALUSBY:\\nBy Ladland. Answer to the Capitol.\\n\\nAll:\\nNo, on the wrong.\\nWhiles Warwick many women on the breedy;\\nIn some uncherers should take it for peace.\\n\\nPERDITA:\\nO love!\\n\\nNurse:\\nAnd so shows, love.\\n\\nARIEL:\\nNo nights shall not be continution\\nWell seen and drink the general throng.\\n\\nSecond Gentleman:\\nThe violent clouds that fair Marcius do him to\\nproud Hereford's. Romeo!\\nDine-mercy and Catesby!\\nWhat, shall we down infected away?\\nThy father bodies to this place.\\n\\nCORIOLANUS:\\nA cold world, then come; the swird I knew no love\\nThan such and shames receive the westeful country:\\nStop not so deep a man of safety,\\nThrough there not scorn he may be rop to question.\\n\\nGONZALO:\\nI'll be thy goodness shall I be best away.\\nI pray your eyes of men, they know thy kindness,\\nAnd ten to one that lives in the hampest blood.\\nIf thou weo'sand, that do but approaches:\\nHis rashness of yo\"\n",
            " b\"ROMEO:\\nBy show your injuries,\\nThat my offen me and my slight-corns and so,\\nBut bloody light in publect knaves' bred!\\n\\nKING HENRY VI:\\nThis country, and go bod, these battle was\\nafter you: if Warwick can be gone to ben?\\n\\nMENENIUS:\\nHow! know'st thou not\\nThat ever henceforth shall I do know\\nBut first to help him not?\\n\\nLEONTES:\\nGo, and bethink you.\\n\\nCitizen:\\nShe shall, go to be trueves; you are they came.\\n\\nKING RICHARD II:\\nRefusion do mis apparel; he's with the botty\\npreys of those beheld this land\\nAs to be crowned England's good,\\nThis was behavious for the crien,\\nTo stop all povern man in blood,\\nBeing ta'en our stining by the name to me.\\nLet's get the cold and fresh from tears to wear;\\nAnd so this morning to her heart?\\nThis prince your stap, rebells up revenge!\\n\\nJULIET:\\nWe are the good well.\\nHolla! when giving hie the consuls and whose immolts\\nThat yispect she parted withal?\\n\\nQUEEN ELIZABETH:\\nBut thus account my dream of Rome,\\nAnd do not say 'Forth Green Romeo is busk;\\nAnd, as our brother royal \"], shape=(5,), dtype=string) \n",
            "\n",
            "________________________________________________________________________________\n",
            "\n",
            "Run time: 3.8618671894073486\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tf.saved_model.save(one_step_model, 'one_step')\n",
        "one_step_reloaded = tf.saved_model.load('one_step')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1v86eMkr9jrX",
        "outputId": "d659f48a-3769-42a7-c87d-718c9b8bbd27"
      },
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:tensorflow:Skipping full serialization of Keras layer <__main__.OneStep object at 0x7fd674bec2e0>, because it is not built.\n",
            "WARNING:tensorflow:Model's `__init__()` arguments contain non-serializable objects. Please implement a `get_config()` method in the subclassed Model for proper saving and loading. Defaulting to empty config.\n",
            "WARNING:tensorflow:Model's `__init__()` arguments contain non-serializable objects. Please implement a `get_config()` method in the subclassed Model for proper saving and loading. Defaulting to empty config.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "states = None\n",
        "next_char = tf.constant(['ROMEO:'])\n",
        "result = [next_char]\n",
        "\n",
        "for n in range(100):\n",
        "  next_char, states = one_step_reloaded.generate_one_step(next_char, states=states)\n",
        "  result.append(next_char)\n",
        "\n",
        "print(tf.strings.join(result)[0].numpy().decode(\"utf-8\"))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fWTNQtm190J_",
        "outputId": "25bac03e-d7f6-4f69-89dc-a68f3bc0223a"
      },
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ROMEO:\n",
            "'Twill fly: go for the hest; and, seeing thou but O,\n",
            "the drum, arise a knife, in ditiners\n",
            "And under\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "xWHSaCIV96k9"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}