diff --git a/docs/conf.py b/docs/conf.py index 49599bfd..713880f0 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -227,6 +227,7 @@ def new_process_docstring(app, what, name, obj, options, lines): nb_execution_allow_errors = False nb_execution_excludepatterns = [ # slow examples + 'nanogpt.ipynb', 'cifar10_resnet.ipynb', 'adversarial_training.ipynb', 'reduce_on_plateau.ipynb', diff --git a/docs/gallery.rst b/docs/gallery.rst index c99dc1e6..8e784aaa 100644 --- a/docs/gallery.rst +++ b/docs/gallery.rst @@ -150,6 +150,23 @@
Adversarial training of CNN on MNIST.
+ +.. raw:: html + +
+ +.. only:: html + + .. image:: /images/examples/tiny_shakespeare.png + :alt: Small Transformer Language Model on Tiny Shakespeare + + :doc:`_collections/examples/nanogpt` + +.. raw:: html + +
Small Transformer Language Model on Tiny Shakespeare.
+
+ .. raw:: html diff --git a/docs/images/examples/tiny_shakespeare.png b/docs/images/examples/tiny_shakespeare.png new file mode 100644 index 00000000..62e9571b Binary files /dev/null and b/docs/images/examples/tiny_shakespeare.png differ diff --git a/examples/nanogpt.ipynb b/examples/nanogpt.ipynb new file mode 100644 index 00000000..b3966fbd --- /dev/null +++ b/examples/nanogpt.ipynb @@ -0,0 +1,750 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "xpfwcMJHTtfw" + }, + "source": [ + "\n", + "# Small Transformer Language Model on Tiny Shakespeare\n", + "\n", + "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.sandbox.google.com/github/google-deepmind/optax/blob/main/examples/nanogpt.ipynb)\n", + "\n", + "This example demonstrates how to train a small-scale transformer-based language model (inspired by NanoGPT) on the Tiny Shakespeare dataset. The core idea is to train a model that can predict the next character in a sequence of text based on the characters that came before it.\n", + "\n", + "**Why the Tiny Shakespeare Dataset?**\n", + "\n", + "* **Manageable Size:** Since we're building a small-scale model, the Tiny Shakespeare dataset provides a suitable training corpus without overwhelming computational resources.\n", + "* **Linguistic Complexity:** Shakespeare's works offer a rich vocabulary and interesting grammatical patterns, making the dataset a good testbed for our model's language learning abilities.\n", + "* **Accessibility:** Easily accessible through TensorFlow Datasets.\n", + "\n", + "**Libraries Used**\n", + "\n", + "* **JAX:** Provides the foundation for numerical computations and automatic differentiation.\n", + "* **Tensorflow Datasets (`tfds`)** Offers easy access to the Tiny Shakespeare dataset.\n", + "* **Flax's Linen Module:** Provides building blocks for defining our neural network architecture.\n", + "* **Optax:** Contains a library of optimization algorithms for training the model's parameters. In this example we'll use the {py:func}`optax.adamw` solver." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "id": "jIabArrRWFw0", + "outputId": "2b03c2df-9a3a-4ce5-9339-a83051fc3580", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "JAX running on GPU\n" + ] + } + ], + "source": [ + "import functools\n", + "\n", + "import flax.linen as nn\n", + "import jax\n", + "import jax.numpy as jnp\n", + "from matplotlib import pyplot as plt\n", + "import optax\n", + "import tensorflow_datasets as tfds\n", + "\n", + "# platform check\n", + "print(\"JAX running on\", jax.devices()[0].platform.upper())" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "UhFD-uojcAI6" + }, + "source": [ + "# Hyperparameters and dataset download\n", + "\n", + "Next, we set some important hyperparameters. This includes hyperparameters for the training process such as the learning rate `LEARNING_RATE` and the batch size `BATCH_SIZE`, as well as model parameters such as the context window size `BLOCK_SIZE` and the number of layers `NUM_LAYERS`.\n", + "\n", + "\n", + "After setting these, we load the Tiny Shakespeare dataset and print the length of the training set, which is around one million characters, and that of the validation set (around 50k characters). Finally, we print a small snippet of the train set." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "id": "34pKN_bIXt8O" + }, + "outputs": [], + "source": [ + "# @markdown Random seed:\n", + "SEED = 42 # @param{type:\"integer\"}\n", + "# @markdown Learning rate passed to the optimizer:\n", + "LEARNING_RATE = 5e-3 # @param{type:\"number\"}\n", + "# @markdown Batch size:\n", + "BATCH_SIZE = 128 # @param{type:\"integer\"}\n", + "# @markdown Numer of training iterations:\n", + "N_ITERATIONS = 50_000 # @param{type:\"integer\"}\n", + "# @markdown Number of training iterations between two consecutive evaluations:\n", + "N_FREQ_EVAL = 2_000 # @param{type:\"integer\"}\n", + "# @markdown Batch size\n", + "BATCH_SIZE = 512 # @param{type:\"integer\"}\n", + "# @markdown Rate for dropout in the transformer model\n", + "DROPOUT_RATE = 0.2 # @param{type:\"number\"}\n", + "# @markdown Context window for the transformer model\n", + "BLOCK_SIZE = 64 # @param{type:\"integer\"}\n", + "# @markdown Number of layer for the transformer model\n", + "NUM_LAYERS = 6 # @param{type:\"integer\"}\n", + "# @markdown Size of the embedding for the transformer model\n", + "EMBED_SIZE = 256 # @param{type:\"integer\"}\n", + "# @markdown Number of heads for the transformer model\n", + "NUM_HEADS = 8 # @param{type:\"integer\"}\n", + "# @markdown Size of the heads for the transformer model\n", + "HEAD_SIZE = 32 # @param{type:\"integer\"}\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "id": "mghpbB9653Gw" + }, + "outputs": [], + "source": [ + "ds = tfds.load(\"tiny_shakespeare\")\n", + "\n", + "# combine train and test examples into a single string\n", + "text_train = \"\"\n", + "for example in ds[\"train\"].concatenate(ds[\"test\"]).as_numpy_iterator():\n", + " text_train += example[\"text\"].decode(\"utf-8\")\n", + "\n", + "# similarly, create a single string for validation\n", + "text_validation = \"\"\n", + "for example in ds[\"validation\"].as_numpy_iterator():\n", + " text_validation += example[\"text\"].decode(\"utf-8\")" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "id": "USiJ0GjWSPu_", + "outputId": "72983a4a-921c-4dcf-b732-b883c2e5a113", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Length of text for training: 1_059_624 characters\n", + "Length of text for validation: 55_770 characters\n" + ] + } + ], + "source": [ + "print(f\"Length of text for training: {len(text_train):_} characters\")\n", + "print(f\"Length of text for validation: {len(text_validation):_} characters\")" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": { + "id": "wOq-djQ9cueI", + "outputId": "8d0c8238-88ff-44fa-8f7d-66c79d032ad2", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "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", + "All:\n", + "We know't, we know't.\n", + "\n", + "First Citizen:\n", + "Let us kill him, and we'll have corn at our own price.\n", + "Is't a verdict?\n", + "\n", + "All:\n", + "No more talking on't; let it be done: away, away!\n", + "\n", + "Second Citizen:\n", + "One word, good citizens.\n", + "\n", + "First Citizen:\n", + "We are accounted poor citizens, the patricians good.\n", + "What authority surfeits on would relieve us: if they\n", + "would yield us but the superfluity, while it were\n", + "wholesome, we might guess they relieved us humanely;\n", + "but they think we are too dear: the leanness that\n", + "afflicts us, the object of our misery, is as an\n", + "inventory to particularise their abundance; our\n", + "sufferance is a gain to them Let us revenge this with\n", + "our pikes, ere we become rakes: for the gods know I\n", + "speak this in hunger for bread, not in thirst for revenge.\n", + "\n", + "\n" + ] + } + ], + "source": [ + "# small sample of the train set\n", + "print(text_train[:1000])" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "FguiERfTcEPa" + }, + "source": [ + "# Data preparation\n", + "\n", + "To prepare the data for the model, we first create a vocabulary consisting of all the unique characters in the dataset. We print that vocabulary and its size.\n", + "\n", + "We then define encoding and decoding functions to convert text into sequences of integers (representing our characters) and vice versa.\n", + "\n", + "Finally, we define a function `get_batch` that returns random mini-batches of data. This function uses JAX's\n", + "[`dynamic_slice`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.dynamic_slice.html) function to efficiently handle sequences of varying lengths within batches. The `@jax.jit` decorator compiles this function for faster execution. The function randomly samples a batch from the data and prepares input sequences (`x`) and target sequences (`y`). The target sequence is simply the input sequence shifted by one position, as the goal of the language model is to predict the next character given the previous ones.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": { + "id": "rESkNoDXFE-4", + "outputId": "89891f78-5cb6-454d-b78f-304528748840", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Vocabulary:, \n", + " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n", + "Length of vocabulary: 65\n" + ] + } + ], + "source": [ + "vocab = sorted(list(set(text_train)))\n", + "print(\"Vocabulary:, \", \"\".join(vocab))\n", + "print(\"Length of vocabulary: \", len(vocab))" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": { + "id": "F-LSTr86bXrV" + }, + "outputs": [], + "source": [ + "# create a mapping from characters to integers\n", + "stoi = {ch: i for i, ch in enumerate(vocab)}\n", + "itos = {i: ch for i, ch in enumerate(vocab)}\n", + "encode = lambda s: [\n", + " stoi[c] for c in s\n", + "] # encoder: take a string, output a list of integers\n", + "decode = lambda l: \"\".join(\n", + " [itos[i] for i in l]\n", + ") # decoder: take a list of integers, output a string\n", + "\n", + "# encode train and validation data\n", + "train_data = jnp.array(encode(text_train))\n", + "eval_data = jnp.array(encode(text_validation))" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "id": "tZLP1asmb8WY" + }, + "outputs": [], + "source": [ + "dynamic_slice_vmap = jax.vmap(jax.lax.dynamic_slice, in_axes=(None, 0, None))\n", + "\n", + "\n", + "@jax.jit\n", + "def get_batch(random_key, data):\n", + " \"\"\"Prepares a random batch of training data.\n", + "\n", + " Args:\n", + " random_key: A random seed for sampling a batch.\n", + " data: The complete training dataset.\n", + "\n", + " Returns:\n", + " x: Input sequences.\n", + " y: Target sequences (shifted inputs).\n", + " \"\"\"\n", + " ix = jax.random.randint(\n", + " random_key, shape=(BATCH_SIZE, 1), minval=0, maxval=len(data) - BLOCK_SIZE\n", + " )\n", + " x = dynamic_slice_vmap(data, ix, (BLOCK_SIZE,))\n", + " y = dynamic_slice_vmap(data, ix + 1, (BLOCK_SIZE,))\n", + " return x, y" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "PQfTaf2UcTSc" + }, + "source": [ + "# NanoGPT Model Definition\n", + "\n", + "The NanoGPT model itself is defined as a Flax Linen module. The core of the model is a Transformer architecture, designed for sequence-to-sequence tasks like language modeling. Key parameters of the model, such as the number of layers, attention heads, and embedding size, are specified here.\n", + "\n", + "Inside the model's `__call__` method, we first embed our input characters into vector representations. Positional embeddings are added to provide the model with a sense of order in the sequence. The core of the Transformer consists of multiple layers. Each layer has two main components:\n", + "\n", + " * **Multi-Head Attention**: This mechanism allows the model to \"attend\" to different parts of the input sequence, improving its understanding of context and relationships within the text. In the code this is implemented through the `nn.MultiHeadDotProductAttention` class.\n", + "\n", + " * **Feedforward Network**: This network processes the output of the attention layer, applying non-linear transformations to further learn complex patterns in the data. This is implemented through the `nn.Sequential` class.\n", + "\n", + "Normalization and dropout (for regularization) are used within the layers to improve training stability. Finally, a dense layer maps the model's output to the vocabulary size, producing probabilities for each character as the next potential character.\n", + "\n", + "The generate function enables the model to create new text sequences. It iteratively generates one character at a time, conditioned on the previously generated text.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": { + "id": "-c7M35UaYMyD" + }, + "outputs": [], + "source": [ + "class NanoGPT(nn.Module):\n", + " \"\"\"NanoGPT model.\"\"\"\n", + " vocab_size: int\n", + " num_layers: int = 6\n", + " num_heads: int = 8\n", + " head_size: int = 32\n", + " dropout_rate: float = 0.2\n", + " embed_size: int = 256\n", + " block_size: int = 64\n", + "\n", + " @nn.compact\n", + " def __call__(self, x, training: bool):\n", + " seq_len = x.shape[1]\n", + "\n", + " x = nn.Embed(self.vocab_size, self.embed_size)(x) + nn.Embed(\n", + " self.block_size, self.embed_size\n", + " )(jnp.arange(seq_len))\n", + " for _ in range(self.num_layers):\n", + " x_norm = nn.LayerNorm()(x)\n", + " x = x + nn.MultiHeadDotProductAttention(\n", + " num_heads=self.num_heads,\n", + " qkv_features=self.head_size,\n", + " out_features=self.head_size * self.num_heads,\n", + " dropout_rate=self.dropout_rate,\n", + " )(\n", + " x_norm,\n", + " x_norm,\n", + " mask=jnp.tril(jnp.ones((x.shape[-2], x.shape[-2]))),\n", + " deterministic=not training,\n", + " )\n", + "\n", + " x = x + nn.Sequential([\n", + " nn.Dense(4 * self.embed_size),\n", + " nn.relu,\n", + " nn.Dropout(self.dropout_rate, deterministic=not training),\n", + " nn.Dense(self.embed_size),\n", + " ])(nn.LayerNorm()(x))\n", + "\n", + " x = nn.LayerNorm()(x)\n", + " return nn.Dense(self.vocab_size)(x)\n", + "\n", + " @functools.partial(jax.jit, static_argnames=(\"self\", \"length\"))\n", + " def generate(self, rng, params, length):\n", + " def _scan_generate(carry, _):\n", + " random_key, context = carry\n", + " logits = self.apply(params, context, training=False)\n", + " rng, rng_subkey = jax.random.split(random_key)\n", + " new_token = jax.random.categorical(\n", + " rng_subkey, logits[:, -1, :], axis=-1, shape=(1, 1)\n", + " )\n", + " context = jnp.concatenate([context[:, 1:], new_token], axis=1)\n", + " return (rng, context), new_token\n", + "\n", + " _, new_tokens = jax.lax.scan(\n", + " _scan_generate,\n", + " (rng, jnp.zeros((1, self.block_size), dtype=jnp.int32)),\n", + " (),\n", + " length=length,\n", + " )\n", + " return new_tokens" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ylVNKEhscy9d" + }, + "source": [ + "# State, Optimizer, and Loss Definition\n", + "\n", + "This section initializes the model's parameters, defines the loss function used for language modeling, and sets up the training and evaluation processes.\n", + "\n", + "In this case the loss function `loss_fun` is the cross-entropy. It uses dropout for regularization, introduced via the `rngs={\"dropout\": dropout_key}` argument. We also define a function for evaluating the model's performance on unseen data (`eval_step`)." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": { + "id": "sjSnK3yDYIus" + }, + "outputs": [], + "source": [ + "model = NanoGPT(\n", + " vocab_size=len(vocab),\n", + " num_layers=NUM_LAYERS,\n", + " num_heads=NUM_HEADS,\n", + " head_size=HEAD_SIZE,\n", + " dropout_rate=DROPOUT_RATE,\n", + " embed_size=EMBED_SIZE,\n", + " block_size=BLOCK_SIZE,\n", + ")\n", + "\n", + "def loss_fun(params, x, y, dropout_key):\n", + " logits = model.apply(params, x, training=True, rngs={\"dropout\": dropout_key})\n", + " return optax.softmax_cross_entropy_with_integer_labels(\n", + " logits=logits, labels=y\n", + " ).mean()\n", + "\n", + "\n", + "@jax.jit\n", + "def eval_step(params, x, y):\n", + " logits = model.apply(params, x, training=False)\n", + " return optax.softmax_cross_entropy_with_integer_labels(\n", + " logits=logits, labels=y\n", + " ).mean()" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": { + "id": "ejU1Yt8XIH80" + }, + "outputs": [], + "source": [ + "key = jax.random.PRNGKey(SEED)\n", + "key, subkey = jax.random.split(key)\n", + "\n", + "var_params = model.init(\n", + " key,\n", + " jnp.ones((BATCH_SIZE, BLOCK_SIZE), dtype=jnp.int32),\n", + " training=False,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "kgSjWONs4eFp" + }, + "source": [ + "We've now instatiated a NanoGPT model with the following number of parameters" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": { + "id": "9Ckqdkd6QVsl", + "outputId": "c80bb1ea-1686-4d0b-88e5-76418960ac44", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Total number of parameters: 3_408_513\n" + ] + } + ], + "source": [ + "n_params = sum(p.size for p in jax.tree_util.tree_leaves(var_params))\n", + "\n", + "print(f\"Total number of parameters: {n_params:_}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "vreyB_oo4Zch" + }, + "source": [ + "# Model training\n", + "\n", + "We start by creating an optimizer and instantiating its state. In this case we'll use {py:func}`optax.adamw` but I encourage you to try other optimizers.\n", + "\n", + "\n", + "We then proceeded to the training loop. For maximum efficiency we extracted the most computationally intensive tasks inside the `step` function and just-in-time compile this function using `@jax.jit`. This allows JAX to perform some optimizations in our code and generally achieve a much higher efficiency than without.\n", + "\n", + "Inside the training loop, we call the aforementioned `step` functions, as well as computing accuracy on a validation set every `N_FREQ_EVAL` iterations." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": { + "id": "1xwLpjDxccMi" + }, + "outputs": [], + "source": [ + "# choosing a different optimizer is as easy as overwriting the line below.\n", + "# For example, to run with sgd instead just use\n", + "# opt = optax.sgd(learning_rate=LEARNING_RATE)\n", + "opt = optax.adamw(learning_rate=LEARNING_RATE)\n", + "\n", + "opt_state = opt.init(var_params)" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": { + "id": "DhnK0G7AQUCA", + "outputId": "d6017a05-7eca-4e52-ff5c-7806fb01759b", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Step: 0\t train loss: 4.586461067199707\t eval loss: 6.0375471115112305\n", + "Step: 2000\t train loss: 1.3861221075057983\t eval loss: 1.4284186363220215\n", + "Step: 4000\t train loss: 1.280903697013855\t eval loss: 1.404958724975586\n", + "Step: 6000\t train loss: 1.198905348777771\t eval loss: 1.4034368991851807\n", + "Step: 8000\t train loss: 1.173327088356018\t eval loss: 1.382951021194458\n", + "Step: 10000\t train loss: 1.1334363222122192\t eval loss: 1.4140418767929077\n", + "Step: 12000\t train loss: 1.116783618927002\t eval loss: 1.4261202812194824\n", + "Step: 14000\t train loss: 1.1018123626708984\t eval loss: 1.4720300436019897\n", + "Step: 16000\t train loss: 1.0866079330444336\t eval loss: 1.4557700157165527\n", + "Step: 18000\t train loss: 1.054302453994751\t eval loss: 1.4516525268554688\n", + "Step: 20000\t train loss: 1.053168773651123\t eval loss: 1.473431944847107\n", + "Step: 22000\t train loss: 1.037564992904663\t eval loss: 1.5162005424499512\n", + "Step: 24000\t train loss: 1.0463223457336426\t eval loss: 1.5376673936843872\n", + "Step: 26000\t train loss: 1.031209945678711\t eval loss: 1.5154197216033936\n", + "Step: 28000\t train loss: 1.0118529796600342\t eval loss: 1.55436372756958\n", + "Step: 30000\t train loss: 0.9960469007492065\t eval loss: 1.5427621603012085\n", + "Step: 32000\t train loss: 1.00498366355896\t eval loss: 1.508345365524292\n", + "Step: 34000\t train loss: 0.9817172884941101\t eval loss: 1.55613112449646\n", + "Step: 36000\t train loss: 0.9934886693954468\t eval loss: 1.5921552181243896\n", + "Step: 38000\t train loss: 0.9811679124832153\t eval loss: 1.5389573574066162\n", + "Step: 40000\t train loss: 0.9739974141120911\t eval loss: 1.5706850290298462\n", + "Step: 42000\t train loss: 0.9757038354873657\t eval loss: 1.5953700542449951\n", + "Step: 44000\t train loss: 0.9462944269180298\t eval loss: 1.5894972085952759\n", + "Step: 46000\t train loss: 0.9620105028152466\t eval loss: 1.5746440887451172\n", + "Step: 48000\t train loss: 0.9629588723182678\t eval loss: 1.5998138189315796\n", + "CPU times: user 14min 48s, sys: 9min 36s, total: 24min 24s\n", + "Wall time: 22min 48s\n" + ] + } + ], + "source": [ + "%%time\n", + "\n", + "all_train_losses = []\n", + "all_eval_losses = []\n", + "\n", + "# we define one iteration of the optimizer and JIT this function\n", + "@jax.jit\n", + "def step(key, params, opt_state):\n", + " key, subkey = jax.random.split(key)\n", + " batch = get_batch(key, train_data)\n", + " loss, grad = jax.value_and_grad(loss_fun)(params, *batch, subkey)\n", + " updates, opt_state = opt.update(grad, opt_state, params)\n", + " params = optax.apply_updates(params, updates)\n", + " return params, key, opt_state, loss\n", + "\n", + "\n", + "for i in range(N_ITERATIONS):\n", + " var_params, key, opt_state, loss = step(key, var_params, opt_state)\n", + " all_train_losses.append(loss)\n", + "\n", + " # once every N_FREQ_EVAL we compute loss on the validation set\n", + " if i % N_FREQ_EVAL == 0:\n", + " key, subkey = jax.random.split(key)\n", + " eval_loss = eval_step(var_params, *get_batch(subkey, eval_data))\n", + " all_eval_losses.append(eval_loss)\n", + " print(f\"Step: {i}\\t train loss: {loss}\\t eval loss: {eval_loss}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": { + "id": "Gc-V4kAKAA9q", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 449 + }, + "outputId": "267708f0-7673-46ab-fe38-24de472fd021" + }, + "outputs": [ + { + "output_type": "display_data", + "data": { + "text/plain": [ + "
" + ], + "image/png": "iVBORw0KGgoAAAANSUhEUgAAAioAAAGwCAYAAACHJU4LAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAABL1ElEQVR4nO3deXhU1cE/8O+dNetk3wkJELYgIItiRCmWTUDcWvVFflWsra8tvmqpSu1rC2gFa1urVUtbrdi3fSl1A31lkciOBWTfDVvYISGQZCbbZJbz+2OSSWbJZDK5c+8kfD/PM8/MvffMnTMnxnw559xzJSGEABEREVEE0qhdASIiIqK2MKgQERFRxGJQISIioojFoEJEREQRi0GFiIiIIhaDChEREUUsBhUiIiKKWDq1K9AZTqcTFy5cQHx8PCRJUrs6REREFAQhBCwWC7Kzs6HRBO4z6dJB5cKFC8jNzVW7GkRERBSCs2fPokePHgHLdOmgEh8fD8D1RU0mk6znttlsWLNmDSZOnAi9Xi/ruakF21kZbGdlsJ2VwXZWTrja2mw2Izc31/13PJAuHVSah3tMJlNYgkpMTAxMJhN/EcKI7awMtrMy2M7KYDsrJ9xtHcy0DU6mJSIioojFoEJEREQRi0GFiIiIIhaDChEREUUsBhUiIiKKWAwqREREFLEYVIiIiChiMagQERFRxGJQISIioojFoEJEREQRi0GFiIiIIhaDije7FbBcAiqOIan2BFB+WO0aERERKS4/Px9/+MMf1K5G174pYVj8+01g3UvQAxgDwOncCvy/D9WuFRERUbvGjh2L66+/Hq+//nqnz7Vjxw4YDAZs2LCh0+fqDNWDyvnz5zFnzhysWrUKdXV1KCgowOLFizFy5Eh1KhSV4LltNatTDyIiUp3TKVBZ16hqHZJiDNBo2r/LcDCEEHA4HNDp2v/zn5aWBpvNJsvndoaqQaWyshKjR4/GbbfdhlWrViEtLQ3Hjh1DUlKSepXyCiqS1aJSRYiISG2VdY0Y8asvVa3DrhfGIyXO2G65mTNnYuPGjdi4cSPeeOMNAMDixYvxyCOPYOXKlXjhhRdw4MABrFmzBrm5uZg9eza2bduG2tpaDBw4EAsXLsT48ePd58vPz8d//dd/oaCgAAAgSRLeeecdrFixAl988QVycnLwu9/9DnfeeWd4vngTVYPKr3/9a+Tm5mLx4sXufb169WqzvNVqhdVqdW+bza7eDpvNJlvqk3QxHo3irK+GMwISZXfV/HOLhNTenbGdlcF2VoaS7Wyz28P+GcHUwWZrf0rpb3/7W5SUlGDQoEGYO3cuAODwYdc8yzlz5uDXv/41evXqhaSkJJw9exaTJk3CvHnzYDQa8Y9//APTpk3DwYMH0bNnT/c5HQ6Hqw5NbT1//nwsWLAACxYswB//+EfMmDEDx48fR3Jycse+Uwd+dqoGlc8++wyTJk3Cfffdh40bNyInJwc//vGP8cMf/tBv+YULF2L+/Pk++9esWYOYmBhZ6qS9WoI7Wm1ba66ieOVKWc5NbSsuLla7CtcEtrMy2M7KUKKda2yA2rMk1n75JeL0wZW1WCwoKyvD7t27AQAHDhwAANxxxx1obGxESUmJu2xOTg7OnDkDABg1ahTS09Px61//GlOnTgUA1NXVoaSkBP3793e3dVFREUwmE44ePYpbbrkFb731Fv74xz9i+PDhHfpOdXV1QZdVtfVPnjyJRYsWYfbs2fj5z3+OHTt24Mknn4TBYMDDDz/sU/7555/H7Nmz3dtmsxm5ubmYOHEiTCaTLHX6zd+tHkElylmPKZNvByReIBUONpsNxcXFmDBhAvT6IH8TqcPYzspgOytDyXa+UtuI/965Iayf0Z5x48cjJdYQVNnXXnsNvXr1wpQpUwAAsbGxAIDHHnsMOTk57nI1NTV46aWXsHLlSly6dAl2ux319fWIj493vzcmJgb9+/cHAEyYMAEAcM8997iPA4DJZEJeXp7HvmA0j4gEQ9Wg4nQ6MXLkSCxYsAAAMGzYMBw8eBB/+tOf/AYVo9EIo9F3nE6v18v2H+vZBs/zayQBjdMKRMkThMg/OX+G1Da2szLYzspQop3TTTrsemF8+wXDqCOTaSVJgkajcbdL86TZxMREj7Z6/vnnUVxcjN/+9rcoKChAdHQ0vvvd78LhcHiU02q1AODeFx0d7XHc+/OC1ZHyqgaVrKwsFBYWeuwbOHAgPv74Y5VqBNRJsb47rWYGFSKia5BGIwU1kTVSGAwG97ySQL766ivMnDkT99xzDwBXD8upU6fCXLvQqDqeMXr0aI/xMgA4evQo8vLyVKoR0CBFwym8kmtDtTqVISIi6oD8/Hxs374dp06dQkVFBZxOp99yffv2xSeffIK9e/di3759ePDBB9ssqzZVg8pPfvITbNu2DQsWLMDx48exZMkS/OUvf8GsWbNUq5OQNKhBlOfOBq6lQkREke+ZZ56BVqtFYWEh0tLS3JNlvb322mtISkrCzTffjGnTpmHSpEkdnhCrFFWHfm644QYsW7YMzz//PF588UX06tULr7/+OmbMmKFmtWBBDEyob9nBRd+IiKgL6NevH7Zu3eqxb+bMmT7l8vPzsW7dOo993p0Ep06dgs1mw8qmK1+FED7nqaqq6lyFg6D6yrR33HEH7rjjjvYLKsgsYpAjXWnZwR4VIiIiVfCaWz8s8FqTxco5KkRERGpgUPHDIryCCifTEhERqYJBxQ8Loj13cOiHiIhIFQwqfpiF11oqnExLRESkCgYVLxLYo0JERBQpGFT84BwVIiKiyMCg4ofvVT/sUSEiIlIDg4ofvj0qDCpERERqYFDxw8weFSIi6oLGjh2Lp59+WrbzPfroo1iwYIFs5wsFg4ofZvaoEBERRQQGFT985qg0WgBn+7fNJiKibsbpBGor1H0EeVfjmTNnYuPGjXjjjTcgSRIkScKpU6dw8OBBTJ48GXFxccjIyMD3vvc9VFRUuN/30UcfYfDgwYiOjkZKSgrGjx+P2tpazJs3D3//+9/x9ddfw2AwQJIkbNiwIUwN3TbV7/UTiSwi2nen1QxEJylfGSIiUk/9VeA3fdStw7MngNjUdou98cYbOHr0KK677jq8+OKLAAC9Xo8bb7wRP/jBD/D73/8e9fX1mDNnDu6//36sW7cOFy9exPTp0/Hqq6/innvugcViwebNmyGEwDPPPINDhw7h5MmTWL58OfR6PZKTk8P9bX0wqPhhRqzvzgYGFSIiilwJCQkwGAyIiYlBZmYmAOBXv/oVhg0b5jHP5L333kNubi6OHj2Kmpoa2O123HvvvcjLywMADB482F02OjoaOp0OmZmZ0Ov1yn6hJhz68aMORtiFV9NwQi0REXUx+/btw/r16xEXF+d+DBgwAABw4sQJDB06FOPGjcPgwYNx33334Z133kFlZaXKtfbEoOKXhBqf1Wm56BsREXUtNTU1mDZtGvbu3evxOHbsGMaMGQOtVovi4mKsWrUKhYWFePPNN9G/f3+UlpaqXXU3Dv14kSTXs0XEIFGqbTnAK3+IiK490cmuOSJq1yFIBoMBDkfLxR/Dhw/Hxx9/jPz8fOh0/v/kS5KE0aNHY/To0fjlL3+JvLw8LFu2DLNnz4bBYIAzyMm84cKg0gauTktERNBogprIGiny8/Oxfft2nDp1CnFxcZg1axbeeecdTJ8+Hc899xySk5Nx/PhxLF26FO+++y527tyJtWvXYuLEiUhPT8f27dtx+fJlDBw4EACQl5eHzz77DCUlJcjMzERCQoLic1U49NMGn0Xf2KNCREQR7plnnoFWq0VhYSHS0tLQ2NiIr776Cg6HAxMnTsTgwYPx9NNPIzExERqNBiaTCZs2bcKUKVPQr18/vPDCC/jd736HyZMnA3At+JadnY2ioiKkpaXhq6++Uvw7sUelDT7L6Fs5R4WIiCJbv379sHXrVp/9n3zyid/yAwcOxOrVq9s8X1paGubPn48pU6bwqp9IY+ZkWiIiItUxqLSBNyYkIiJSH4NKG3hjQiIiIvUxqLSBPSpERETqY1Bpg8/lyZyjQkREpDgGFS9CuJ59r/phjwoREZHSGFTaYPG56odBhYiISGkMKl6al9A3C687KLNHhYiISHEMKm3w6VGx1QEOmzqVISIiukYxqLTB7D1HBeDwDxERkcIYVNrgc9UPwGX0iYiIFMag0oYGGNAotF472aNCRESkJAaVNkm+vSqcUEtERKQoBpUAfFen5dAPERGRkhhUAuBaKkREROpiUAmAq9MSERGpi0ElADO8Fn1jjwoREZGiGFQCsAjvoR/OUSEiIlISg4oXCZL7te9VPwwqRERESmJQCcAnqHDoh4iISFEMKgH4LKPPybRERESKYlAJwMweFSIiIlUxqATABd+IiIjUxaASgM+Cbxz6ISIiUhSDSgC+PSoMKkREREpiUAnAZ46KwwrYGtSpDBER0TWIQSUAnx4VgMM/RERECmJQCcBnHRWAwz9EREQKYlDxIrUsTItG6GEVes8CXJ2WiIhIMQwq7eBaKkREROphUGkHV6clIiJSD4NKO3zWUuGib0RERIphUGkH11IhIiJSD4NKO3yu/OHQDxERkWIYVNrhM0eFPSpERESKYVBph0+PCueoEBERKUbVoDJv3jxIkuTxGDBggJpV8uEzR4VDP0RERIrRqV2BQYMG4csvv3Rv63TqVkkIz21e9UNERKQe1YOKTqdDZmZmUGWtViusVqt722x29W7YbDbYbDZZ6iOE02PbLGI9jzdUwy7TZxHcPze5fn7kH9tZGWxnZbCdlROutu7I+VQPKseOHUN2djaioqJQVFSEhQsXomfPnn7LLly4EPPnz/fZv2bNGsTE+LkvTwiuXNGg9YiYd49K7dVLWLtypSyfRS2Ki4vVrsI1ge2sDLazMtjOypG7revq6oIuKwnhPdihnFWrVqGmpgb9+/fHxYsXMX/+fJw/fx4HDx5EfHy8T3l/PSq5ubmoqKiAyWSSpU4PLd6JrSevureLNIfwT8PL7m0RnQz77KOyfBa5UnVxcTEmTJgAvV7f/hsoJGxnZbCdlcF2Vk642tpsNiM1NRXV1dXt/v1WtUdl8uTJ7tdDhgzBqFGjkJeXhw8++ACPPvqoT3mj0Qij0eizX6/Xy9aAUuu7EsL38mTJaoZep/O8eyF1mpw/Q2ob21kZbGdlsJ2VI3dbd+RcEXV5cmJiIvr164fjx4+rXRU3n8uTnXbAVq9OZYiIiK4xERVUampqcOLECWRlZaldFTefBd8AXqJMRESkEFWDyjPPPIONGzfi1KlT+Pe//4177rkHWq0W06dPV7NaHmq8L08GeIkyERGRQlSdo3Lu3DlMnz4dV65cQVpaGm655RZs27YNaWlpalbLgx06QB8D2FrNUOYy+kRERIpQNagsXbpUzY8PntHkGVSs7FEhIiJSQkTNUYlYUQme2+xRISIiUgSDSjCivK7x5mRaIiIiRTCoBMPoFVQ4mZaIiEgRDCpe/K7j5t2jwqEfIiIiRTCoBMO7R4VDP0RERIpgUAkGJ9MSERGpgkElGD5DP5yjQkREpAQGlWAYvXpUOPRDRESkCAaVYHAyLRERkSoYVILhPUeFK9MSEREpgkElGD7rqLBHhYiISAkMKsHwtzKtEOrUhYiI6BrCoBIM7x4V4QQaa9SpCxER0TWEQcWLBD9L03r3qAAc/iEiIlIAg0owvHtUAF6iTEREpAAGlWBotIAh3nMfF30jIiIKOwaVYHEtFSIiIsUxqASLNyYkIiJSHINKsHxuTMihHyIionBjUAmWv7VUiIiIKKwYVILlszote1SIiIjCjUHFy6BsP5ciA5xMS0REpAIGFS+DcjznovRLj3O94GRaIiIixTGoBMtnMi2DChERUbgxqHjxs4C+i8/QD+eoEBERhRuDSrCMXj0qHPohIiIKOwaVYHEyLRERkeIYVILlPUeFPSpERERhx6DSDgHheuFz1Y8FcDqVrxAREdE1hEElWN5DPxDsVSEiIgozBpVgefeoAAwqREREYcagEixDHHwuXuaEWiIiorBiUAmWRsMbExIRESmMQaUjvNdS4aJvREREYcWg0hFcS4WIiEhRDCpepDbX0AdvTEhERKQwBpWO8LkxIYd+iIiIwolBpSN4Y0IiIiJFMai0Q4hWGxz6ISIiUhSDSkdwMi0REZGiGFQ6gj0qREREimJQ6QifybQMKkREROHEoNIRnExLRESkKAaVjvBemZZDP0RERGHFoNIRnExLRESkKAaVjvCeo2KrBRx2depCRER0DWBQ8SIhwBr63lf9ABz+ISIiCiMGlY7wHvoBOKGWiIgojBhU2tF6YVroYwBJ61mAPSpERERhw6DSEZLECbVEREQKYlDpKO8JtexRISIiChsGlY7ynlDLOSpERERhw6DSUVxGn4iISDEMKh3FGxMSEREpJmKCyiuvvAJJkvD000+rXZXAfHpUOPRDREQULhERVHbs2IE///nPGDJkiNpVaR9vTEhERKQY1YNKTU0NZsyYgXfeeQdJSUlqVwdSgIVpAXDoh4iISEE6tSswa9YsTJ06FePHj8evfvWrgGWtViusVqt722x2hQSbzQabzSZLfex2z3v3CCE8zq3Rx6L1km/O+mo4ZPrsa1Fz28r18yP/2M7KYDsrg+2snHC1dUfOp2pQWbp0KXbv3o0dO3YEVX7hwoWYP3++z/41a9YgJiZGljrtvSIBraJIbW0tVq5c6d7uWXEaw1qVry47g02tjlNoiouL1a7CNYHtrAy2szLYzsqRu63r6uqCLqtaUDl79iyeeuopFBcXIyoqKqj3PP/885g9e7Z722w2Izc3FxMnToTJ5Oc+PCGQDl7C4qP73duxsbGYMuWWluNH7MDZ99zbidEaTJkyRZbPvhbZbDYUFxdjwoQJ0Ov1alen22I7K4PtrAy2s3LC1dbNIyLBUC2o7Nq1C+Xl5Rg+fLh7n8PhwKZNm/DWW2/BarVCq/W8r47RaITRaPQ5l16vl60BdTrvJpE8zx3rOY9GajDzF0UGcv4MqW1sZ2WwnZXBdlaO3G3dkXOpFlTGjRuHAwcOeOx75JFHMGDAAMyZM8cnpEQMI5fQJyIiUopqQSU+Ph7XXXedx77Y2FikpKT47I8o3pcn2xsAeyOgM6hTHyIiom5M9cuTuxzvBd8A9qoQERGFieqXJ7e2YcMGtavQPu91VADXom+xqcrXhYiIqJtjj0pH6aMArdcwD1enJSIiCgsGlVBwdVoiIiJFMKh4aW8FfQB+7vfDoEJERBQODCrtEr67vCfUskeFiIgoLBhUQuE99MM5KkRERGHBoBIKDv0QEREpgkElFFydloiISBEMKqHwnqPCHhUiIqKwYFAJhc/QT5Uq1SAiIuruGFRCwXVUiIiIFMGgEgpOpiUiIlIEg0oo2KNCRESkCAaVUHAyLRERkSIYVLxIXmvoCz8L0/oO/VS3UZCIiIg6g0ElFN5DP04bYG9Qpy5ERETdGINKKLyHfgAO/xAREYUBg0oovHtUAE6oJSIiCgMGlVDoDIAu2nMfb0xIREQku5CCyt/+9jesWLHCvf3cc88hMTERN998M06fPi1b5SKavwm1REREJKuQgsqCBQsQHe3qUdi6dSvefvttvPrqq0hNTcVPfvITWSsYsbiWChERUdjpQnnT2bNnUVBQAABYvnw5vvOd7+Cxxx7D6NGjMXbsWDnrF7m4lgoREVHYhdSjEhcXhytXrgAA1qxZgwkTJgAAoqKiUF9fL1/tIpn30A97VIiIiGQXUo/KhAkT8IMf/ADDhg3D0aNHMWXKFADAoUOHkJ+fL2f9Ipf30A/nqBAREckupB6Vt99+G0VFRbh8+TI+/vhjpKSkAAB27dqF6dOny1pB5XkuTVt6pc5/Md6YkIiIKOxC6lFJTEzEW2+95bN//vz5na6Q2srMQa4wy8m0REREYRdSj8rq1auxZcsW9/bbb7+N66+/Hg8++CAqKytlq5wa3v/3qeAKRiV6brNHhYiISHYhBZVnn30WZrPrD/OBAwfw05/+FFOmTEFpaSlmz54tawWV1mh3BleQ66gQERGFXUhDP6WlpSgsLAQAfPzxx7jjjjuwYMEC7N692z2xttvzGfphUCEiIpJbSD0qBoMBdXWuSaZffvklJk6cCABITk5297R0VZpgW4STaYmIiMIupB6VW265BbNnz8bo0aPx9ddf41//+hcA4OjRo+jRo4esFVSaRpLaLwT4LvjGybRERESyC6lH5a233oJOp8NHH32ERYsWIScnBwCwatUq3H777bJWUGnaYIOKzzoqZkAI+StERER0DQupR6Vnz574/PPPffb//ve/73SF1BZsTvEZ+hEOoLEWMMbJXiciIqJrVUhBBQAcDgeWL1+OI0eOAAAGDRqEO++8E1qtVrbKqUEKtUcFcA3/MKgQERHJJqSgcvz4cUyZMgXnz59H//79AQALFy5Ebm4uVqxYgT59+shaSSVlJ0bjeHlN+wX9BZUGM2DKlr9SRERE16iQ5qg8+eST6NOnD86ePYvdu3dj9+7dOHPmDHr16oUnn3xS7joq6j/H9A6uoFYHGLx6TzihloiISFYh9ahs3LgR27ZtQ3JysntfSkoKXnnlFYwePVq2yqlBq/Ed+mmwORCl9zOkZTQBja16X7joGxERkaxC6lExGo2wWCw++2tqamAwGDpdKTWdvep7E8KD59sIIFydloiIKKxCCip33HEHHnvsMWzfvh1CCAghsG3bNjz++OO488475a6jovJTY332vbu51H9h3piQiIgorEIKKn/4wx/Qp08fFBUVISoqClFRUbj55ptRUFCA119/XeYqKisjPspnX0mZb+8RAN9F37g6LRERkaxCmqOSmJiITz/9FMePH3dfnjxw4EAUFBTIWjk1ZCb4BpWpg7P8F/Ye+mGPChERkayCDirt3RV5/fr17tevvfZa6DVSmUHn28k0MMvPpciAn9VpOUeFiIhITkEHlT179gRVLugF0yJY/4x4j+Eeu9PpvyBvTEhERBRWQQeV1j0m3Z33JcoOZxv38OFkWiIiorAKaTJtd6fTBhlUOJmWiIgorBhU/NBIoQYVzlEhIiKSE4OKH95DP3YO/RAREamCQcWPXacrPba/udRGAOFkWiIiorBiUAnCP7ad8X/Ae+jHagbaukKIiIiIOoxBpTO8h34ggMY2VrElIiKiDmNQCcKUwZn+D3gP/QAc/iEiIpIRg0oQjlxso5fEEA/Aa4E7TqglIiKSDYNKEEorav0f0GgAY7znPvaoEBERyYZBxY+EaL3H9uCchDZKwv+EWiIiIpIFg4of9w7P8dgelN3GTQkB3piQiIgojBhU/PC+g7LVHuCSY5+1VBhUiIiI5KJqUFm0aBGGDBkCk8kEk8mEoqIirFq1Ss0qAQCqam0e28v2nG+7MFenJSIiChtVg0qPHj3wyiuvYNeuXdi5cye+/e1v46677sKhQ4fUrBb+tfNs8IV5Y0IiIqKw0an54dOmTfPYfvnll7Fo0SJs27YNgwYNUqlWHcShHyIiorBRNai05nA48OGHH6K2thZFRUV+y1itVlitVve22ezqvbDZbLDZbH7fE4rHbs3HXzaf8tjX1vk1+jhoW20766vgkLEu3V1zu8r58yNfbGdlsJ2VwXZWTrjauiPnk4QQbdwaWBkHDhxAUVERGhoaEBcXhyVLlmDKlCl+y86bNw/z58/32b9kyRLExMTIVqetZRKWnmyJH3lxArMHO/yWLSj7HIMufODeLjMNwbY+z8hWFyIiou6mrq4ODz74IKqrq2EyBbiyFhEQVBobG3HmzBlUV1fjo48+wrvvvouNGzeisLDQp6y/HpXc3FxUVFS0+0U74pM95zHnk5Z5MoVZ8fj0x/57eTS7FkO7+ln3trPHjXA8vFK2unR3NpsNxcXFmDBhAvR6fftvoJCwnZXBdlYG21k54Wprs9mM1NTUoIKK6kM/BoMBBQUFAIARI0Zgx44deOONN/DnP//Zp6zRaITRaPTZr9frZW3AKIPnuRxOtH3+2GSPTY3VAg1/cTpM7p8h+cd2VgbbWRlsZ+XI3dYdOVfEraPidDo9ek3UoNd6NovNGWgdFe+rfjiZloiISC6q9qg8//zzmDx5Mnr27AmLxYIlS5Zgw4YN+OKLL9SsFnQazxsN2h0BRse4jgoREVHYqBpUysvL8dBDD+HixYtISEjAkCFD8MUXX2DChAlqVsunR8Xu6MDKtI01gNMBaLT+yxMREVHQVA0qf/3rX9X8+DZ5B5XGQD0q3kM/gKtXJTpJ5loRERFdeyJujkok0Gm9hn4CzVHxHvoBOE+FiIhIJgwqfui9g0qgHhVDLCB5DfNwGX0iIiJZMKj4odN4XfUTaI6KJAHGeM99nFBLREQkCwYVP3wm0zrbWRPP534/DCpERERyYFDxw3vox+EUcAYKK94TatmjQkREJAsGFT90Wt9mCbjom5GLvhEREYUDg4of3gu+Ae1MqOXQDxERUVgwqPhh0Pk2S8dWp2WPChERkRwYVPzw16PSGHB1Wu+hH/aoEBERyYFBxQ9/c1QCLvrmM/TDHhUiIiI5MKj4YfAXVHhjQiIiIsUxqPjhvYQ+0M6ib5xMS0REFBYMKn74m6NiY48KERGR4hhU/JAkySesBO5R4WRaIiKicGBQaYPPjQk7sjItJ9MSERHJgkGlDd5X/tgD9ah4D/3Y6wGHLQy1IiIiurYwqLTBd+inAyvTAhz+ISIikgGDShu876DcoTkqAFenJSIikgGDSht856gECCq6KECj99zHeSpERESdxqDShvNVDR7b1fUB5pxIEtdSISIiCgMGlSCtPngpcAGupUJERCQ7BpUgFaTHBS7AHhUiIiLZMai0oUdilOd2UkzgN3hPqGWPChERUacxqLQhOzHaY7vRHmAyLeA79MPJtERERJ3GoNIGg86zadoNKlxGn4iISHYMKm3YcvyKx/ZXJyoCv8FnMi17VIiIiDqLQSVIG0ouBy7AHhUiIiLZMajIxeeqH/aoEBERdRaDily4jgoREZHsGFTkwnVUiIiIZMegIhf2qBAREcmOQUUunExLREQkOwYVuXgHFYcVsDX4L0tERERBYVDpACFE2we9h34ADv8QERF1EoNKB1yusbZ90HsyLcDhHyIiok5iUGnDj77Vy2fflmMBVqfVGQGd540MuTotERFR5zCotCHXz92S/7qlNPCbeGNCIiIiWTGotGHCwHSffScv1wZ+E9dSISIikhWDShsSY/Q+++ptjsBv4loqREREsmJQkRN7VIiIiGTFoNIBeSm+81Y8eK+lwh4VIiKiTmFQ6YDTV+oCF+BkWiIiIlkxqMiJy+gTERHJikFFTpxMS0REJCsGFTn59Khw6IeIiKgzGFTk5HPVD4MKERFRZzCoyIlDP0RERLJiUOkghzPAHZS5jgoREZGsGFQC6Jfg9Nlnc/juc/PXoyICBBsiIiIKiEElgPHZviEj4P1+vCfTOu2ArV7mWhEREV07GFQCSDb6BpWAd1D2DioAJ9QSERF1AoNKABrJd9/Hu8+1/QZjvO8+TqglIiIKGYNKAEnGDr5Bqwf0XvcD4oRaIiKikDGoBOCvR6VdPjcm5NAPERFRqBhU5MYbExIREclG1aCycOFC3HDDDYiPj0d6ejruvvtulJSUqFmlzuNaKkRERLJRNahs3LgRs2bNwrZt21BcXAybzYaJEyeitjbAJcCRjqvTEhERyUan5oevXr3aY/v9999Heno6du3ahTFjxqhUq05ijwoREZFsVA0q3qqrXfM5kpOT/R63Wq2wWq3ubbPZFQJsNhtsNpusdQl0vvLqWiTFGPwe0xjioW217airhFPmunUnze0s98+PPLGdlcF2VgbbWTnhauuOnE8SIjLWeHc6nbjzzjtRVVWFLVu2+C0zb948zJ8/32f/kiVLEBMT4+cdnffUVt8sN6OPAzem+2+2wvNL0bd8pXv7bNLN2J3/eFjqRkRE1BXV1dXhwQcfRHV1NUwmU8CyERNUfvSjH2HVqlXYsmULevTo4beMvx6V3NxcVFRUtPtFO8pms6G4uNhvUAGAYy9N9Ltfs+U1aDcucG87+06C4/7/lbVu3UlzO0+YMAF6vV7t6nRbbGdlsJ2VwXZWTrja2mw2IzU1NaigEhFDP0888QQ+//xzbNq0qc2QAgBGoxFGo+8qbHq9Pmz/sQ7vmYjdZ6r8fqZfMUkem5rGGmj4i9SucP4MqQXbWRlsZ2WwnZUjd1t35FyqXvUjhMATTzyBZcuWYd26dejVq5ea1fHrvhE5HXuD94JvnExLREQUMlWDyqxZs/CPf/wDS5YsQXx8PC5duoRLly6hvj5y7jh87/XZHXuDz1U/XPCNiIgoVKoGlUWLFqG6uhpjx45FVlaW+/Gvf/1LzWp50HR0HX2fdVQYVIiIiEKl6hyVCJnHKy/vHhWrBRACkEK5cRAREdG1jff6CVFdo93/Ae8eFeEEGmvCXyEiIqJuiEElRA+/97X/A96TaQHOUyEiIgoRg0qIdpyq9H/AGO+7j1f+EBERhYRBJQh/+n/Dgy+s0QIGr7DCGxMSERGFhEElCKN6pXTsDbwxIRERkSwYVIKg13WwmbznqbBHhYiIKCQMKkGIM/q/invf2Sr/b/C+8qehjXJEREQUEINKJ/zgf3b6P8ChHyIiIlkwqHTCZYvV/wGf1WkZVIiIiELBoBIO7FEhIiKSBYNKkJ64rcDv/jJzg+9Onzsoc8E3IiKiUDCoBOnRW3r53e93hVoO/RAREcmCQSVISbEGv/u/uWTx3cmhHyIiIlkwqISDkeuoEBERyYFBRQYNNofnDp85KgwqREREoWBQ6YAtc27zu3/up4c8d/gM/XAyLRERUSgYVDqgR1KM3/3/2nnWc4f3ZNpGC+D06nUhIiIKB1sDUHUGOLcLuFqqdm06zf/a8NRhDTYHovRa14Z3jwoAWC1AdKKidSIiConDDpjPA5WngKrTrucGM5DSB8gaCmQOBozx7Z1FfXYrUFMO1Ja7nt2vL7c8N1QBkgbQ6gGtAdDoXM9aveuh0Xu99j5mALRN7zHEudrFaGp69npotJ34Lo1A7WWgpqzl2eM7tXpYW/Xij34KmPBip5tSTQwqMhnwi9U49cpU14Z3jwrgmlDLoEJEkUAIoL4SqCwFKpuCSOtQUn0OcNoDnyO5KbRkDQWyhgCZQ4HYDt5pPhROB2C+iKTaE5BKVgINV5v+cPsJIdYIG3bXx/gJMK1CjSHOVab+qm8Qqa8M7TNryuX9DipgUOmg+0f2wAc7z/k9dr6qHjmJ0a7/2CQNIJwtBzlPhShyWS1AfZXrj4QhBtBFAZKkdq06x+l0BY8rx1uCSOUpVzCpOt35qxGvnnA9Dn3Ssi8hF8gc4hlg4rOCb0shXD0c1eeA6vNA9VnXa/P5ln2WC9A77RgDAEc79xUUZ6tzPWrKlPtMBpVrz9xpg9oMKlPe2Ix9cycCGo0rHbcOJ7zyh6iFEK4g35mu8I5qMANXTzb9gT3pGru/0vS61ut/5pLGFVr0MYA+GjDEtoQYfWzTc0zT/mj3PkljRHblCUjn0oDU3kBsuuv/B+FWUw6UHwbKDrueyw8D5d8Attrwf3Zr1Wddj5IVLfti05qGi5oCTGpfoLaiVQA52xRKzrkeSte5O9P6X/+rq2FQ6aBYY9tNVl1vQ3W9DQnRetdaKq2DCtdSoe5KCKCxFqi70vS42ur1Ff/766+6uvCjE4GYVCAmBYht/dz8OsX1uvmYPjpwXRqqXcHjyglXEHGHkpOu4YGgv5MTaKxxPTpAB+AGADj1R9cOrRFIzAUS84DEnq0eTdtx6R3rubHWAOVHWoWRpnBSV9GhegbNEAck5bsehlig7BBw+Zv2h4Vaq70MHP/S9VCLRu9q69i0pud0IC7N9RyTAkAAjkbAYXN9t+bXDhvgtDVtN+132lqOORpbytutrt8Dq6Xl0ehnQdBOfxddy/eIy2j6Lq0esU3749KAqMSu3zMIBpWQvHjXIPzS+5LkJhNe24iv/3u8a0Jt69Ee9qhQV1V7Bag42vKwXPQNH3Y/97wKRn2l63HlWHDl9bG+4UWIlp6Suiuh1SNcHFbX0MuV4/6P66JcwyWJPYGkPM8go4tyhYLWPSVVp+Wtn6QFEno0hZG8llCS2PQck+z7h87W4KrLxX3Apf2u57JDof83ECKHpIMmPhOS+w+0nxDSfEytP9jOpsDrHV6s3g9zq+N1rgDvHTqaQ0l0kjK9dBGEQSUEdw7NbjOolFusuFrbiGTvRd/Yo0KRzGF3/RGsOApUHPN8rr+qdu1a2GqBqlrXpZfdgb3BFdKCDWqhiEoEknu3hJDWgcTUw3XFSkfoo4Cc4a5HM4fd9d9K6/BycX8nehQk1x/mhB6eD1MOkNADtthMrNzwNaZMnQq9Xh/iZyhAo3H9o9XflaAUNAaVECTGBB73G/5SMUqHxMMjvzdUhbNKRMGxWoDy0lZhpCmQXDnh6tLu9iTXUExyb9dVK8m9XY+UPq5/wTZ339vqXZMeG2ubnutaJkL63VcH2GrhtNairuIMYu2VkByNyn41XRSQNgDIGASkDwTSC12v4zLC35ug1QEZha4Hprv2OZ2uq4ou7vMMMHVXXEPjCT2AhJxWASS3ZV98NqAL8P9Zm61bDGlQcBhUQrT5udtw66vr2zy+/pQV3269g0M/kclhc/2RvnwEuFziejgaXbdBMDb9S6j52WNfQssxfZR69XfPD6loGYqprWg1NFMB1F2FtuYyJpYdhX5PiJc4Bkujdw3HxKS4hg3cr1vNQWner9G56thc39qKlu9R6/UsOrBgoqRx/dFLaRVEmkNJUh6gM4bt6ztsNqxduRJTJt8OfcNVV89P1RlXb1XV6VbbZ0MPhpLG9X0yCoH0plCSMcjVQ6Lk5OT2aDSun0FKH+C6e137RNNckDD+DKj7YVAJUW6y/1Vqm52t03m2Lod+1OWwuSZXXj7iuhqi+fnK8c73JGiNnoGm+Vkf6/pjrNE2PTe/br3d6rjkXa7pta2u1UTUViGktunZYW23ihoA7UxDbec7GoCUAtcVG0m9XPMBWoeQ5vBhjJf/X7pOp2s9jNqm798caGpdIQzC4foj3RxIEnsG/te4EiQNYMpyPXqO8j3udAI1l1yhpfJ0qzDT9Ny8jkl8dlMgaXpkFAKp/dUNx50hSQwp1GEMKp0w+bpMrDp4ye8xC7yCDHtUlOGwu7qby480TURseq44Fr6hDYfVdWVDR64qiVQxKUBqP1cgSe3X8joxT71/rWs0rgmE0UkACtSpg9w0GsCU7Xr0vMn3uNPh6nlo7yonomsAg0onvDl9GAr+e5XfY2bhGVTspV9B9+kTrn+Juh8pLa9jUlxLMl/LHHbXZMnGVnMBGmtb9rV+batruny06bXV4lrMquKo63/w1DZJ4+qB8A4kKX2VWVmU2qfRAhqGFCKAQaVTdFoNfv2dwZjz8QGfY949Krq6MmDP3wOfMDqpVYhJ9Xwdk+r615WkASABElq91ri6VL23PY612jbEubrqoxLC/6/k5qW63StLNi/ydB4wn4eu+jwmWy5Dt98e1BBG2MRlAukDgLSBrksDG8yu4YYGs2vYrsHsWqOj+bWadW2LRt9qHojn3BCHMRG7jp7DsAkPQJ/ej93vRNRlMKh00gM39ETJpRq895XnHSovieSOn6x5TYkKpdaFllx/lKOTXcHF+9nfvuhkz/Fxq6UpdJxzhw+fbVtdoBpA0dkEcRlAWn9XIGkOJukDmoYVOsBubRVevAKNtWm/vcHVhe+0t3r2ft20Ldo47rC5QkVMqqv9fYJIalDzQ5w2Gy6Wr8SwtAGA7hrvuSOiLoVBRQa/nFboE1T+7RyEQ848DNLIvECTrERLOLp6Ivi36WNcf9itNZF3069msWmuSzXTB3o+x4QQIP3RGZsWYUqT53xEROQXg4pMvnnpdgz4xWr3thUG3NP4Im7SHEaOVIEUmJEimZEqVbtfu7YtAIR6FQ9F89oR4SZpXct2t77PSvPdRd33XIlteR2b2hRIBnKuBRFRN8GgIpMovRbzphVi3v8ddu9rhB6bnEMDvi8tRovtPxkOTf2VlitHaiv8vK5wDQOg6WZuzTd1a73tccx7u1VZVSabNq80meNeXRKmHNjjMrFt/3HcdOu3oYtJaLnRmyHWdUksF3UiIrqmMajIaOboXlhx4CJ2nAp+Ua3LdQ70fnkHAGDJD0fh5utSw1W9FrYG17LodVdbnptvFFdX6f9YQzUC9vxEJzetMpnbFERyXMtzNweT+Cy/a1sImw1XSldCZA8DInkpbCIiUgWDisw+fPxmlJsbcOOCtR1+74PvbAcAxBt1WPLDmzAwKx4aSYJGI3Ovgj4K0Det4RAspwOor2oVYipd5zH1cJ3HEHgBPCIiolAwqIRBuimq3SX2A7FY7Zj21hb3dq/UWPz90RvRI0nFMKDRNq37wrkfRESknGvrXtEKyk2OwdqffkuWc5VW1OKWX69H/s9WIP9nK7Dj1FXYHE402Dpw/xMiIqIuiD0qYdQnLQ6nXpmK77+/A+u+KZftvPf9aavH9oePF+GGfJkuuyUiIoogDCoKeG/mDQCA/J+tCMv5vYNLs/l3DsLEQRnINEVB4tUzRETUBTGoKOjkgin4aPc5/GL5QVjtzrB/3tzPDmHuZ4fc27f2TcV/3NATeq2ECYUZDC9ERBTxGFQUpNFIuH9kLu4fmYvP9l3Ak//co+jnbz5Wgc3HKnz23319NobnJSEhWo/rcxORlxKraL2IiIjawqCikjuHZuPOodmoa7TjJ//aiy8OlalWl+V7L2D53gttHu+bHof//eEonLxci/OV9RjbPw0pcbypHRERhR+DispiDDr8+XsjAQBCCGw+VgFzgw1PLFG2tyWQY+U1uPHltteFual3Mu4d1gNj+6ch3RQFh1PAancgWq/l8BIREXUKg0oEkSQJY/q5bnJ3x5Bs1FrtuFrbGPJ6LErZdvIqtp28GnT5CYUZePme65AeH4XLFisqGsJYOSIi6tIYVCJYrFGHWKMOp16ZCgBY/005Hnl/h8q16rziw2UoPtx6qEuHl/ascW8NzknAY2N6Y9pQz5VzG+1OmBtsSIzWQ6flEkBERNcCBpUu5LYB6e7Q0qyu0Y7tJ6/iTxtPYHtp8L0akezA+Wr81z/34L8CTDbOSohC34x4PD6mN6IMWrz0+WHsOVOFh4ryMHfaIGg1EhpsDnx5pAzJsQbc3EeBeygREZHsGFS6uBiDDrcNSMdtA9IhhMAXh8qw71wVFm04oXbVwupidQMuVjdg09HLHvv/Z+tp/M/W022+L9MUhUvmBjxxWwEyTEY4BfDgqJ7Q++mhKTM3YOPRy+ifEY+huYlyfwUiIgoCg0o3IkkSbr8uE7dfl4k5tw9w779YXY8PdpxDmaUBjXYnPtp1TsVaquuS2TUh5q31x937Wq81055HRuejsrYRy/deQLxRhzmTB6C0ohYzRvVE77Q42etLRHStY1C5BmQlROOp8X3d27+9b6jHcYdT4Fi5BR/uPIe/bilVunpdyuKvTrlfW6x2vLD8IAC0225j+6ehMMuE81X1uGdYDlJijThfVYeSSzXITDCi0SGw50wl+qTFYfqNPZEUo8eV2kZ8vOscMkxRuOv6bF5BRUTXJAYVglYjYUCmCb+4oxC/uKPQvb/Wasc3lyw4dKEaX5dehSlaj37pcZj3f4dVrG3XtKHkMjaUuIapPg2wZg0A/OaLEp99T/9rr8e2TiNBkoCbeqfgsTG9UXy4DOZ6G3omx2BMvzSMbHXvpwabA+ev1qHWBjidAldqrNBqJCTGGNqsgxACFTWNiDVqEWMI/n8TtVY7HELAFKUP+j1ERIEwqFCbYo06jMhLwoi8JDxUlO/eP+OmPGwsuQxTtB439kpGdZ0NAgINNid++elBrDms3uJ11wq7UwDwv9rwH9Yd9/cWADr8fGdxu+e+Y0gWPt9/0b1977AcfHdkD2gkCf/xl20AgDH90jCqVzImFmYgOdaA/eeqUVpRi9+tKUFtowNPjeuLKYOzAAD9M+ND+IZERC4MKtRheq0G4wsz3NsJMS3/ev7LQyMDvnfF/ot48fNDKDNbw1Y/6pzWIQUAPtlzHp/sOe+xb9PRy9h09LLf3h8AeGPtMbyx9pjP/qmDs7DiwEWf/ctnjUaUXoO31h1HXaMDcUYdjpfX4P6RPTDjpjxsP3kV31wy46beKcgwRSHaoEWt1Y6UWAMa7E58vOsctBoJMQYtlmw/g/zUWPxiaqHHf5tE1DUxqJCipg7JwtQhWe5tm82GlStXYsqUKai2OrH1xBX0zYhDRnwUXvr8MD7Zcx5GnUaRmzhS+PkLKQBw99tf+d0/7/8OhzTUuPN0ZZuTxpc+dhP++fUZ9xDc3GmF2HW60h3QHirK87hy7MSCKThwvhp/3VKKMX1TkRpnxMXqBmSYjIjWa3G5xoqtJ65g/MAM3NonCYBrzR9J43Sv93PkohkzF3+NAZkm/HHGcMQadXA6BQRcQ6/BarA50GBzICFazzlLdM2QhBBC7UqEymw2IyEhAdXV1TCZTLKeu/UfUL2e/yoLl1DaudZqx0e7zqFPWhxu6p2MbSev4vTVWqTEGrHz1FW8ywnB1I1E67Wotzn8HvtWvzQ8MjofdY0OWO0OpMVFYWR+EqL0WgghcK6yHmev1mFIbiL0WglGndbj/c1zkaINWsQZff/d6lry4BLKzFbcdX12wHlNSuL/n5UTrrbuyN9vVXtUNm3ahN/85jfYtWsXLl68iGXLluHuu+9Ws0rUBcQadXj45nz39i19U3ELXAu63X5dJl5oNSG4LVdqrFi64yzOVdZhx6lKHC+vgVYjweHssrmduqm2QgoAbDx6GRu91hIKl45cxv/DW3uhV2ocSi6ZcWvfNAzukYC/bz0Nc4MN9Y0O6HUa3Dk0G3kpMShauA4A8MZ/XI+7rs9Bg82B/9t3AUa9FlMHZ0GrkeB0CpRZGlBZa4MkAfEGCZfqgM/2XcTQnskoSG9ZGuCyxYoovQbxIUzoPldZhzJzA4b0SPS7thKpQ9WgUltbi6FDh+L73/8+7r33XjWrQteYlDgjZt1WEHR5m8MJIQCDToMLVfV4ecURmBtseGpcX2SYorD64CW8vPJIGGtM1HW8s7mlV/NvbSzAuGT7GY/tp5buxVNL93rsezLA6tSADth3INQqBmXqkCys2O9/uFIjAb3T4nC8vAYA8Is7CnHPsBxU1FhRWduIWKMOj/5th3s+3qezRmNwTgI+23cBu05X4lv90jC+MAMNNgf2na3C6St1OFZuwcAsE4b1TEKj3YmkWD1OXq7Fa8VHEWfU4Rd3FKJXaiyEEHhz3XH8z9ZTyEuJxewJ/bCrabjzzNU6AMCSH45Cj8QYHLpQjQFZJvRKjYXTKXCyohbJsQYkx/r2jtU12lFmtiIrIQpReq3PcbVEzNCPJEnt9qhYrVZYrS2TMM1mM3Jzc1FRURGWoZ/i4mJMmDCBXYthdK21c1WdDUfLLUiKNkCjkVBdb8OgrHgY9VpYGuyotzlQXWdDg92B6no7Fqz6BsfKa9WuNhFdg3KTovHcxD6wnNiLKZPGITbKKNu5zWYzUlNTgxr66VJBZd68eZg/f77P/iVLliAmJiaMtSOKTEIAZ2uBvVc0uFgHnKmV0M8k0NskMCxFoM4OXG6QUGsH/ve4FgaNwA8GOHHKAnxTpcFJCydkElH7MqIFfn5928OQHVVXV4cHH3yw+wUV9qh0P2xnZbTVzkIIbDl+BXWNDnx7QFqb9zz68kg5RuQlISXWgKo6GwBg99kqlJTVYFhuAk5fqUNuUjRuyE/CqkNleGX1UcW+GxEp49hLE2U7V0d6VLrU5clGoxFGo2/Xk16vD9sfuXCem1qwnZXhr52/XZjVRmmXHil6zLylZdG27KZFbwt7JPkt//hYEx4f29fvscMXzPjXDtc6J82LCJaZG2DUaRBr1MGo00CSXHe+bnQ4EaPXos7mQOnlWhwrr8HGo5cx8+Y8XJeTgHOV9dh64goabA5Y7c4213QhInnI+f/ojpyrSwUVIuraCrNNmH/XdR77shOjfcpF6bXuyXwmrQZDcxMxNDcR3x3Rw12mT1oc+rS6EWRHJke3xeZwosHmWnBu28mrsDmcuKUgFScu16DB5kS/zDg4ncCXR8qQnxKLdJMRxYfLsP9cFT7adQ5OAdyQ6sSc7xSh3i7wzuZSnzt8E1HHMKgQETXRazXu4a+iPinu/X0zPG8DMG1otvv1/7spD0AeXv3uUPeaE0N7JECv1+PWvmkh18XhFH4XgyszN6DGakfv1Fh8tu8CjpZZMHVwNgqzXd3n56vqcazMgqQYAxxCoGdyDC5brHjjy2PYeboSlgabewHFdx8aieQ4A+wOgV9+ehDfXLIAAOKNOlis9pDrTiQnVYNKTU0Njh9vuS9JaWkp9u7di+TkZPTs2VPFmhERqautFWszTFFovoHFXdfn+BzPSYxGjlcvVWqcEX/63oiAn7f66TE++2qtdvy++CjKLVY8NqY3rstJgN3hCjk6r/lMV2qs2H++GrlJMbA7nYg16HDgfDV6pcYiJdaA4iNlsNmdKMxOwJAeCdh28grWHC7D9T0S8Z0RPXC0zIIYgxbZidEw19uw9ptyPPfRfvf5xw9Mx7f6p6NvehyqahqwdusupOQWoN7mRE5SNBKjDTDoNCitqEWf9DgIIbD/XDUGZMZje+lVrNh/MeCaNBTYr+5qf32qcFE1qOzcuRO33Xabe3v27NkAgIcffhjvv/++SrUiIiLAtbii9wKK3gGlWUqcEbf1T/fYl5vccjXmjFF5HsfG9k/H2FblB2a1TKhMiTPi/pG5uH9krt/PstlssJYKTJnQN+Bch+Ygd9/IXPz2vqEex4QQEALQdOAWBt7vt9qdsNqcWL73PFLiDJg6OMvj1gYXqupRUmZBo90JvVbCzX1SEaXXosZqR2VtI7ITo3G1thEayXVBSUWNFdmJ0aiqa8SGksuotdpRkB6Hc5X1OHzBjMJsEy5WN6DkkhnrS5QbUvzRQAceGNmj/YJhompQGTt2LCLkoiMiIrqGSJKEztwuSZIk91yq1itlt5adGO13DlacUee+ZUFafMsFIs2LsMUZdU1DiuprHs5UE9cIJiIioojFoEJEREQRi0GFiIiIIhaDChEREUUsBhUiIiKKWAwqREREFLEYVIiIiChiMagQERFRxGJQISIioojFoEJEREQRi0GFiIiIIhaDChEREUUsBhUiIiKKWKrePbmzmu+8bDabZT+3zWZDXV0dzGZzwNuIU+ewnZXBdlYG21kZbGflhKutm/9uN/8dD6RLBxWLxQIAyM3NVbkmRERE1FEWiwUJCQkBy0gimDgToZxOJy5cuID4+HhIkiTruc1mM3Jzc3H27FmYTCZZz00t2M7KYDsrg+2sDLazcsLV1kIIWCwWZGdnQ6MJPAulS/eoaDQa9OjRI6yfYTKZ+IugALazMtjOymA7K4PtrJxwtHV7PSnNOJmWiIiIIhaDChEREUUsBpU2GI1GzJ07F0ajUe2qdGtsZ2WwnZXBdlYG21k5kdDWXXoyLREREXVv7FEhIiKiiMWgQkRERBGLQYWIiIgiFoMKERERRSwGFT/efvtt5OfnIyoqCqNGjcLXX3+tdpUiyqZNmzBt2jRkZ2dDkiQsX77c47gQAr/85S+RlZWF6OhojB8/HseOHfMoc/XqVcyYMQMmkwmJiYl49NFHUVNT41Fm//79uPXWWxEVFYXc3Fy8+uqrPnX58MMPMWDAAERFRWHw4MFYuXKl7N9XDQsXLsQNN9yA+Ph4pKen4+6770ZJSYlHmYaGBsyaNQspKSmIi4vDd77zHZSVlXmUOXPmDKZOnYqYmBikp6fj2Wefhd1u9yizYcMGDB8+HEajEQUFBXj//fd96tNdfycWLVqEIUOGuBezKioqwqpVq9zH2cbh8corr0CSJDz99NPufWxrecybNw+SJHk8BgwY4D7eJdtZkIelS5cKg8Eg3nvvPXHo0CHxwx/+UCQmJoqysjK1qxYxVq5cKf77v/9bfPLJJwKAWLZsmcfxV155RSQkJIjly5eLffv2iTvvvFP06tVL1NfXu8vcfvvtYujQoWLbtm1i8+bNoqCgQEyfPt19vLq6WmRkZIgZM2aIgwcPin/+858iOjpa/PnPf3aX+eqrr4RWqxWvvvqqOHz4sHjhhReEXq8XBw4cCHsbhNukSZPE4sWLxcGDB8XevXvFlClTRM+ePUVNTY27zOOPPy5yc3PF2rVrxc6dO8VNN90kbr75Zvdxu90urrvuOjF+/HixZ88esXLlSpGamiqef/55d5mTJ0+KmJgYMXv2bHH48GHx5ptvCq1WK1avXu0u051/Jz777DOxYsUKcfToUVFSUiJ+/vOfC71eLw4ePCiEYBuHw9dffy3y8/PFkCFDxFNPPeXez7aWx9y5c8WgQYPExYsX3Y/Lly+7j3fFdmZQ8XLjjTeKWbNmubcdDofIzs4WCxcuVLFWkcs7qDidTpGZmSl+85vfuPdVVVUJo9Eo/vnPfwohhDh8+LAAIHbs2OEus2rVKiFJkjh//rwQQog//vGPIikpSVitVneZOXPmiP79+7u377//fjF16lSP+owaNUr853/+p6zfMRKUl5cLAGLjxo1CCFeb6vV68eGHH7rLHDlyRAAQW7duFUK4AqVGoxGXLl1yl1m0aJEwmUzudn3uuefEoEGDPD7rgQceEJMmTXJvX2u/E0lJSeLdd99lG4eBxWIRffv2FcXFxeJb3/qWO6iwreUzd+5cMXToUL/Humo7c+inlcbGRuzatQvjx49379NoNBg/fjy2bt2qYs26jtLSUly6dMmjDRMSEjBq1Ch3G27duhWJiYkYOXKku8z48eOh0Wiwfft2d5kxY8bAYDC4y0yaNAklJSWorKx0l2n9Oc1luuPPqrq6GgCQnJwMANi1axdsNpvH9x8wYAB69uzp0c6DBw9GRkaGu8ykSZNgNptx6NAhd5lAbXgt/U44HA4sXboUtbW1KCoqYhuHwaxZszB16lSf9mBby+vYsWPIzs5G7969MWPGDJw5cwZA121nBpVWKioq4HA4PH5AAJCRkYFLly6pVKuupbmdArXhpUuXkJ6e7nFcp9MhOTnZo4y/c7T+jLbKdLefldPpxNNPP43Ro0fjuuuuA+D67gaDAYmJiR5lvds51DY0m82or6+/Jn4nDhw4gLi4OBiNRjz++ONYtmwZCgsL2cYyW7p0KXbv3o2FCxf6HGNby2fUqFF4//33sXr1aixatAilpaW49dZbYbFYumw7d+m7JxNdC2bNmoWDBw9iy5YtalelW+rfvz/27t2L6upqfPTRR3j44YexceNGtavVrZw9exZPPfUUiouLERUVpXZ1urXJkye7Xw8ZMgSjRo1CXl4ePvjgA0RHR6tYs9CxR6WV1NRUaLVanxnQZWVlyMzMVKlWXUtzOwVqw8zMTJSXl3sct9vtuHr1qkcZf+do/RltlelOP6snnngCn3/+OdavX48ePXq492dmZqKxsRFVVVUe5b3bOdQ2NJlMiI6OviZ+JwwGAwoKCjBixAgsXLgQQ4cOxRtvvME2ltGuXbtQXl6O4cOHQ6fTQafTYePGjfjDH/4AnU6HjIwMtnWYJCYmol+/fjh+/HiX/W+aQaUVg8GAESNGYO3ate59TqcTa9euRVFRkYo16zp69eqFzMxMjzY0m83Yvn27uw2LiopQVVWFXbt2ucusW7cOTqcTo0aNcpfZtGkTbDabu0xxcTH69++PpKQkd5nWn9Ncpjv8rIQQeOKJJ7Bs2TKsW7cOvXr18jg+YsQI6PV6j+9fUlKCM2fOeLTzgQMHPEJhcXExTCYTCgsL3WUCteG1+DvhdDphtVrZxjIaN24cDhw4gL1797ofI0eOxIwZM9yv2dbhUVNTgxMnTiArK6vr/jfd4em33dzSpUuF0WgU77//vjh8+LB47LHHRGJioscM6GudxWIRe/bsEXv27BEAxGuvvSb27NkjTp8+LYRwXZ6cmJgoPv30U7F//35x1113+b08ediwYWL79u1iy5Ytom/fvh6XJ1dVVYmMjAzxve99Txw8eFAsXbpUxMTE+FyerNPpxG9/+1tx5MgRMXfu3G5zefKPfvQjkZCQIDZs2OBxmWFdXZ27zOOPPy569uwp1q1bJ3bu3CmKiopEUVGR+3jzZYYTJ04Ue/fuFatXrxZpaWl+LzN89tlnxZEjR8Tbb7/t9zLD7vo78bOf/Uxs3LhRlJaWiv3794uf/exnQpIksWbNGiEE2zicWl/1IwTbWi4//elPxYYNG0Rpaan46quvxPjx40VqaqooLy8XQnTNdmZQ8ePNN98UPXv2FAaDQdx4441i27Ztalcpoqxfv14A8Hk8/PDDQgjXJcq/+MUvREZGhjAajWLcuHGipKTE4xxXrlwR06dPF3FxccJkMolHHnlEWCwWjzL79u0Tt9xyizAajSInJ0e88sorPnX54IMPRL9+/YTBYBCDBg0SK1asCNv3VpK/9gUgFi9e7C5TX18vfvzjH4ukpCQRExMj7rnnHnHx4kWP85w6dUpMnjxZREdHi9TUVPHTn/5U2Gw2jzLr168X119/vTAYDKJ3794en9Gsu/5OfP/73xd5eXnCYDCItLQ0MW7cOHdIEYJtHE7eQYVtLY8HHnhAZGVlCYPBIHJycsQDDzwgjh8/7j7eFdtZEkKIjvfDEBEREYUf56gQERFRxGJQISIioojFoEJEREQRi0GFiIiIIhaDChEREUUsBhUiIiKKWAwqREREFLEYVIiIiChiMagQERFRxGJQISJFzZw5E3fffbfa1SCiLoJBhYiIiCIWgwoRhcVHH32EwYMHIzo6GikpKRg/fjyeffZZ/O1vf8Onn34KSZIgSRI2bNgAADh79izuv/9+JCYmIjk5GXfddRdOnTrlPl9zT8z8+fORlpYGk8mExx9/HI2NjQE/s7a2VuFvTkRy0qldASLqfi5evIjp06fj1VdfxT333AOLxYLNmzfjoYcewpkzZ2A2m7F48WIAQHJyMmw2GyZNmoSioiJs3rwZOp0Ov/rVr3D77bdj//79MBgMAIC1a9ciKioKGzZswKlTp/DII48gJSUFL7/8cpufyfuuEnVtDCpEJLuLFy/Cbrfj3nvvRV5eHgBg8ODBAIDo6GhYrVZkZma6y//jH/+A0+nEu+++C0mSAACLFy9GYmIiNmzYgIkTJwIADAYD3nvvPcTExGDQoEF48cUX8eyzz+Kll14K+JlE1HVx6IeIZDd06FCMGzcOgwcPxn333Yd33nkHlZWVbZbft28fjh8/jvj4eMTFxSEuLg7JycloaGjAiRMnPM4bExPj3i4qKkJNTQ3Onj3b4c8koq6BQYWIZKfValFcXIxVq1ahsLAQb775Jvr374/S0lK/5WtqajBixAjs3bvX43H06FE8+OCDYflMIuoaGFSIKCwkScLo0aMxf/587NmzBwaDAcuWLYPBYIDD4fAoO3z4cBw7dgzp6ekoKCjweCQkJLjL7du3D/X19e7tbdu2IS4uDrm5uQE/k4i6LgYVIpLd9u3bsWDBAuzcuRNnzpzBJ598gsuXL2PgwIHIz8/H/v37UVJSgoqKCthsNsyYMQOpqam46667sHnzZpSWlmLDhg148sknce7cOfd5Gxsb8eijj+Lw4cNYuXIl5s6diyeeeAIajSbgZxJR18XJtEQkO5PJhE2bNuH111+H2WxGXl4efve732Hy5MkYOXIkNmzYgJEjR6Kmpgbr16/H2LFjsWnTJsyZMwf33nsvLBYLcnJyMG7cOJhMJvd5x40bh759+2LMmDGwWq2YPn065s2b1+5nElHXJQleu0dEXcDMmTNRVVWF5cuXq10VIlIQh36IiIgoYjGoEBERUcTi0A8RERFFLPaoEBERUcRiUCEiIqKIxaBCREREEYtBhYiIiCIWgwoRERFFLAYVIiIiilgMKkRERBSxGFSIiIgoYv1/eyjoHDjI7rUAAAAASUVORK5CYII=\n" + }, + "metadata": {} + } + ], + "source": [ + "plt.plot(all_train_losses, label=\"train\", lw=3)\n", + "plt.plot(\n", + " jnp.arange(0, len(all_eval_losses) * N_FREQ_EVAL, N_FREQ_EVAL),\n", + " all_eval_losses,\n", + " label=\"test\",\n", + " lw=3,\n", + ")\n", + "plt.xlabel(\"steps\")\n", + "plt.ylabel(\"loss\")\n", + "plt.grid()\n", + "plt.legend(frameon=False)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "E6-aaLDL7RbI" + }, + "source": [ + "# Text generation\n", + "\n", + "Finally, after training, we use the generate function to let the NanoGPT model demonstrate its ability to create text that resembles Shakespeare, albeit in a miniature form." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": { + "id": "6AejKtZnFmhK", + "colab": { + "base_uri": "https://localhost:8080/" + }, + "outputId": "42fb39c9-b8fe-4d98-825b-589770fc42f3" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "HORTENSIO:\n", + "I go, anot away.\n", + "\n", + "PETRUCHIO:\n", + "How! what's your will?\n", + "\n", + "GREGORY:\n", + "Ha, ha.\n", + "\n", + "SAMPSON:\n", + "True; the prince's doom. I am going to them,\n", + "As 'twixt justice to bid thine eyes:\n", + "Execute thy eyes set it down in thy fortune's blood.\n", + "For my part, I'll play too welcome to what;\n", + "But Peter's Capitol,--while we walk her soul\n", + "As deep as years shall be touch'd with such a wandering deed.\n", + "What say'st thou? wilt thou sway thy brother?\n", + "\n", + "YORK:\n", + "Mount I to destroy?\n", + "\n", + "GLOUCESTER:\n", + "Say you can, such a shower,\n", + "As 'twere to grant on in sincerity;\n", + "Made the cunning calm but words to speak the fray?\n", + "Gardener, for she shall not speak:\n", + "The first fault is forget to be touched all,\n", + "That sharp too the steeled stars set it fair.\n", + "\n", + "DUKE OF AUMERLE:\n", + "I know your tongue that seems up crown'd upon;\n", + "And we look'd when the sea, whereon my life must be,\n", + "The queen of tever and most deeply queen, and here stand\n", + "To peck our ceremonious rats:\n", + "And then like to cheek the sad discourse have done.\n", + "See, what a bright then repose his maje\n" + ] + } + ], + "source": [ + "# # Let's now generate some text\n", + "key, subkey = jax.random.split(key)\n", + "text = model.generate(key, var_params, 1000)[:, 0, 0].tolist()\n", + "print(decode(text))" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": { + "id": "irq1sjG0d_2w" + }, + "outputs": [], + "source": [] + } + ], + "metadata": { + "colab": { + "machine_shape": "hm", + "provenance": [], + "gpuType": "A100" + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + }, + "accelerator": "GPU" + }, + "nbformat": 4, + "nbformat_minor": 0 +}