makemore_mlp.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPMVwJTtSiss8jnTQSWDfnt",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/KhawajaAbaid/flax-deep-learning/blob/main/makemore_mlp.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from flax import linen as nn\n",
        "from flax.training import train_state\n",
        "import jax\n",
        "from jax import numpy as jnp, random, tree\n",
        "import optax"
      ],
      "metadata": {
        "id": "MjJ-M-ZFNNDJ"
      },
      "execution_count": 56,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GXTallLxEtZm",
        "outputId": "6e02cb1c-7f72-4b3f-cc17-b29da032b1dd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['emma', 'olivia', 'ava', 'isabella', 'sophia', 'charlotte', 'mia', 'amelia', 'harper', 'evelyn']\n"
          ]
        }
      ],
      "source": [
        "# Load dataset\n",
        "import os\n",
        "\n",
        "filepath = \"./names.txt\"\n",
        "if (not os.path.exists(filepath)):\n",
        "    !wget https://raw.githubusercontent.com/karpathy/makemore/master/names.txt\n",
        "\n",
        "dataset = open(filepath, 'r').read().splitlines()\n",
        "print(dataset[:10])"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Create vocab and mapping\n",
        "vocab = sorted(list(set(''.join(dataset))))\n",
        "\n",
        "stoi = {c: i for i, c in enumerate(vocab, start=1)}\n",
        "itos = {i: c for i, c in enumerate(vocab, start=1)}\n",
        "\n",
        "vocab.append('.')\n",
        "stoi['.'] = 0\n",
        "itos[0] = '.'\n",
        "vocab_size = len(vocab)"
      ],
      "metadata": {
        "id": "2SUg5RhrFFgw"
      },
      "execution_count": 43,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "context_length = 3\n",
        "\n",
        "def build_dataset(data):\n",
        "  X, Y = [], []\n",
        "  for word in data:\n",
        "    context = [0] * context_length\n",
        "    for c in word + '.':\n",
        "      X.append(context)\n",
        "      ix = stoi[c]\n",
        "      Y.append(ix)\n",
        "      context = context[1:] + [ix]\n",
        "  return jnp.asarray(X), jnp.asarray(Y)"
      ],
      "metadata": {
        "id": "fHh-H7vfFc0p"
      },
      "execution_count": 53,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Create dataset splits\n",
        "def shuffle_randomly(dataset):\n",
        "  import random\n",
        "  random.shuffle(dataset)\n",
        "  return dataset\n",
        "\n",
        "dataset = shuffle_randomly(dataset)\n",
        "\n",
        "n_train = int(0.9 * len(dataset))\n",
        "print(n_train)\n",
        "x_train, y_train = build_dataset(dataset[:n_train])\n",
        "x_test, y_test = build_dataset(dataset[n_train:])\n",
        "\n",
        "print(f\"Train Samples: {len(x_train)} | Test Samples {len(x_test)}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XocgELWaGnyc",
        "outputId": "44942bb5-8c0a-4ce3-a7a2-ee06093589a1"
      },
      "execution_count": 62,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28829\n",
            "Train Samples: 205240 | Test Samples 22906\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Create / Initialize model\n",
        "embedding_dim = 10\n",
        "hidden_dim = 200\n",
        "input_dim = context_length * embedding_dim\n",
        "\n",
        "key, subkey = random.split(random.PRNGKey(1337))\n",
        "embeddings = random.normal(key, (vocab_size, embedding_dim))\n",
        "\n",
        "\n",
        "class MLP(nn.Module):\n",
        "  @nn.compact\n",
        "  def __call__(self, x):\n",
        "    x = nn.Dense(features=hidden_dim)(x)\n",
        "    x = nn.tanh(x)\n",
        "    return nn.Dense(features=vocab_size)(x)\n",
        "\n",
        "\n",
        "key, subkey = random.split(subkey)\n",
        "model = MLP()\n",
        "params = model.init(key, jnp.zeros((1, input_dim)))\n",
        "\n",
        "tx = optax.sgd(learning_rate=0.1)\n",
        "\n",
        "state = train_state.TrainState.create(\n",
        "    apply_fn=model.apply,\n",
        "    params=params,\n",
        "    tx=tx)"
      ],
      "metadata": {
        "id": "NMywDdtTGsu-"
      },
      "execution_count": 64,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train_step(state, x, y):\n",
        "  def forward_and_loss(params, x, y):\n",
        "    x = embeddings[x]\n",
        "    x = x.reshape((x.shape[0], -1))\n",
        "    logits = state.apply_fn(params, x)\n",
        "    loss = optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()\n",
        "    return loss\n",
        "  grad_fn = jax.value_and_grad(forward_and_loss)\n",
        "  loss, grads = grad_fn(state.params, x, y)\n",
        "  state = state.apply_gradients(grads=grads)\n",
        "  return state, loss"
      ],
      "metadata": {
        "id": "yQDZHcueJapm"
      },
      "execution_count": 67,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Training loop\n",
        "batch_size = 128\n",
        "for _ in range(10000):\n",
        "  # create random batches\n",
        "  key, subkey = random.split(key)\n",
        "  idx = random.randint(key, (batch_size,), 0, len(x_train))\n",
        "  x_batch = x_train[idx]\n",
        "  y_batch = y_train[idx]\n",
        "\n",
        "  # train step\n",
        "  state, loss = train_step(state, x_batch, y_batch)\n",
        "print(loss)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h-s4qxc2I8nV",
        "outputId": "f5c1acc4-bd17-41cd-c162-665904ae8c8d"
      },
      "execution_count": 111,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.1480448\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate new names\n",
        "def generate(state, n_names=10, seed=2005):\n",
        "  key, subkey = random.split(random.key(seed))\n",
        "  for i in range(n_names):\n",
        "    context = [0] * context_length\n",
        "    name = ''\n",
        "    while True:\n",
        "      e = embeddings[jnp.asarray(context)].reshape((1, -1))\n",
        "      logits = state.apply_fn(state.params, e)\n",
        "      probs = jax.nn.softmax(logits).reshape(-1)\n",
        "      key, subkey = random.split(subkey)\n",
        "      ix = random.choice(key, jnp.arange(vocab_size), p=probs)\n",
        "      if ix == 0:\n",
        "        break\n",
        "      context = context[1:] + [ix]\n",
        "      name += itos[ix.tolist()]\n",
        "    print(name)\n"
      ],
      "metadata": {
        "id": "lhonXg8TPmHf"
      },
      "execution_count": 112,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "generate(state)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QNT2iwwYPNBE",
        "outputId": "9a3adc4d-aa8a-4695-96f9-b0bcef8acca8"
      },
      "execution_count": 113,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "sibh\n",
            "nyia\n",
            "khluns\n",
            "lailtilysiami\n",
            "erily\n",
            "kir\n",
            "ankaham\n",
            "kimlyien\n",
            "mibraquidaniy\n",
            "marian\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# The end."
      ],
      "metadata": {
        "id": "o6a-a6iAYMfq"
      },
      "execution_count": 114,
      "outputs": []
    }
  ]
}