custom_transformer_from_scratch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "k2tzAwI3XEfP",
        "outputId": "ab3d5b51-3bee-439c-84c9-1230c257e9b2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting BPEmb\n",
            "  Downloading bpemb-0.3.4-py3-none-any.whl (19 kB)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.8/dist-packages (from BPEmb) (1.22.4)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from BPEmb) (2.25.1)\n",
            "Requirement already satisfied: gensim in /usr/local/lib/python3.8/dist-packages (from BPEmb) (3.6.0)\n",
            "Collecting sentencepiece\n",
            "  Downloading sentencepiece-0.1.97-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m16.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from BPEmb) (4.64.1)\n",
            "Requirement already satisfied: smart-open>=1.2.1 in /usr/local/lib/python3.8/dist-packages (from gensim->BPEmb) (6.3.0)\n",
            "Requirement already satisfied: scipy>=0.18.1 in /usr/local/lib/python3.8/dist-packages (from gensim->BPEmb) (1.7.3)\n",
            "Requirement already satisfied: six>=1.5.0 in /usr/local/lib/python3.8/dist-packages (from gensim->BPEmb) (1.15.0)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->BPEmb) (4.0.0)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->BPEmb) (2.10)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->BPEmb) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->BPEmb) (2022.12.7)\n",
            "Installing collected packages: sentencepiece, BPEmb\n",
            "Successfully installed BPEmb-0.3.4 sentencepiece-0.1.97\n"
          ]
        }
      ],
      "source": [
        "!pip install BPEmb\n",
        "import math\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from bpemb import BPEmb"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "#Attention(Q, K,V) = softmax(QK¹/√dk)V\n"
      ],
      "metadata": {
        "id": "VH4m3VdJXNrj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def scaled_dot_product_attention (query, key, value, mask=None):\n",
        "  key_dim = tf.cast(tf.shape(key) [-1], tf.float32)\n",
        "  scaled_scores = tf.matmul(query, key, transpose_b=True) / np.sqrt(key_dim)\n",
        "  if mask is not None:\n",
        "    scaled_scores = tf.where (mask==0, -np.inf, scaled_scores)\n",
        "  softmax = tf.keras.layers.Softmax()\n",
        "  weights = softmax(scaled_scores)\n",
        "  return tf.matmul(weights, value), weights"
      ],
      "metadata": {
        "id": "YSVCx8mDXa4x"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "####TESTING \"scaled_dot_product_attention\" using random queries,keys,values matrices"
      ],
      "metadata": {
        "id": "Z4rlt25RYCRS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "seq_len = 3\n",
        "embed_dim = 4\n",
        "queries = np.random.rand(seq_len, embed_dim)\n",
        "keys = np.random.rand(seq_len, embed_dim)\n",
        "values = np.random.rand(seq_len, embed_dim)\n",
        "print(\"Queries: \\n\", queries)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TiD_B9NHXlH1",
        "outputId": "cb5b4fd7-f18f-4e1c-f35c-32a34c1636f6"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Queries: \n",
            " [[0.40459511 0.5565327  0.60170997 0.2437919 ]\n",
            " [0.92323341 0.01134934 0.81965197 0.2727533 ]\n",
            " [0.44399669 0.52053036 0.32596123 0.16678462]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "output, attn_weights = scaled_dot_product_attention(queries, keys, values)\n",
        "print(\"Output\\n\", output, \"\\n\")\n",
        "print(\"Weights\\n\", attn_weights)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DGoP-OYiX58l",
        "outputId": "b3f18eaf-620d-4357-bb29-0ced893e87e3"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Output\n",
            " tf.Tensor(\n",
            "[[0.6590466  0.71541196 0.22662729 0.6481204 ]\n",
            " [0.6560365  0.71541536 0.22508144 0.65121156]\n",
            " [0.6669714  0.7139002  0.22793788 0.6444146 ]], shape=(3, 4), dtype=float32) \n",
            "\n",
            "Weights\n",
            " tf.Tensor(\n",
            "[[0.32888445 0.37981954 0.291296  ]\n",
            " [0.32391748 0.3857926  0.29028997]\n",
            " [0.3362787  0.36108956 0.3026317 ]], shape=(3, 3), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "####GENERATE QUERY,VALUE, KEYS MATRICES FOR MULTI HEAD ATTENTION USING RANDOM WEIGHTS"
      ],
      "metadata": {
        "id": "vWT-YKJ7Yn5G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "batch_size = 1\n",
        "seq_len = 3\n",
        "embed_dim = 12\n",
        "num_heads = 3\n",
        "head_dim = embed_dim // num_heads\n",
        "print (f\"Dimension of each head: {head_dim}\")\n",
        "x = np.random.rand(batch_size, seq_len, embed_dim).round(1)\n",
        "print(\"Input shape: \", x.shape, \"\\n\")\n",
        "print(\"Input:\\n\", x)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_TzIarhAX5_K",
        "outputId": "3c026b0c-f259-45ac-c42d-c057aeea0da9"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dimension of each head: 4\n",
            "Input shape:  (1, 3, 12) \n",
            "\n",
            "Input:\n",
            " [[[0.7 0.1 0.3 0.5 0.8 0.9 0.3 0.1 0.9 0.8 0.1 0.4]\n",
            "  [0.2 0.8 0.5 0.7 0.1 0.7 0.6 0.1 0.9 0.2 0.9 0.4]\n",
            "  [0.  0.  0.4 0.7 0.6 0.8 0.7 0.1 0.2 0.1 1.  1. ]]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# The query weights for each head.\n",
        "wq0 = np. random.rand(embed_dim, head_dim).round(1)\n",
        "wq1 = np.random.rand(embed_dim, head_dim).round(1)\n",
        "wq2 = np.random.rand(embed_dim, head_dim).round(1)\n",
        "# The key weights for each head.\n",
        "wk0 = np.random.rand(embed_dim, head_dim).round(1)\n",
        "wk1 = np.random.rand(embed_dim, head_dim).round(1)\n",
        "wk2 = np.random.rand(embed_dim, head_dim).round(1)\n",
        "# The value weights for each head.\n",
        "wv0 = np.random.rand(embed_dim, head_dim).round(1)\n",
        "wv1 = np.random.rand(embed_dim, head_dim).round(1)\n",
        "wv2 = np.random.rand(embed_dim, head_dim).round(1)"
      ],
      "metadata": {
        "id": "5rvR-p8fX6B3"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"The three sets of query weights (one for each head): \")\n",
        "print(\"wq0: \\n\", wq0)\n",
        "print(\"wq1: \\n\", wq1)\n",
        "print(\"wq2: \\n\", wq1) "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OnAxupvZX6EK",
        "outputId": "f56e4ebd-6eb2-4b07-da3d-eaf93c6ceaa3"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The three sets of query weights (one for each head): \n",
            "wq0: \n",
            " [[0.3 0.9 0.6 0.6]\n",
            " [0.9 0.4 0.9 0.9]\n",
            " [0.3 0.3 0.5 0.5]\n",
            " [0.3 0.8 0.1 0.7]\n",
            " [0.8 0.2 0.3 0.7]\n",
            " [0.9 0.6 0.5 0.6]\n",
            " [0.2 0.3 0.2 0.5]\n",
            " [0.3 0.7 0.3 0.2]\n",
            " [0.5 0.6 0.4 0.5]\n",
            " [0.1 0.1 0.4 0.3]\n",
            " [0.8 0.3 0.3 0.8]\n",
            " [0.9 0.8 0.4 0.4]]\n",
            "wq1: \n",
            " [[0.  0.4 0.6 0.5]\n",
            " [0.5 1.  0.6 0.2]\n",
            " [0.7 0.4 0.6 1. ]\n",
            " [0.8 0.2 0.5 0.7]\n",
            " [0.9 0.8 0.5 0.7]\n",
            " [0.4 0.6 0.6 0.7]\n",
            " [0.9 0.9 0.2 0.5]\n",
            " [0.1 0.2 0.  0.4]\n",
            " [0.3 0.4 0.3 0.9]\n",
            " [0.2 0.3 0.4 0.9]\n",
            " [0.  0.8 0.9 0.5]\n",
            " [0.3 0.  0.8 0. ]]\n",
            "wq2: \n",
            " [[0.  0.4 0.6 0.5]\n",
            " [0.5 1.  0.6 0.2]\n",
            " [0.7 0.4 0.6 1. ]\n",
            " [0.8 0.2 0.5 0.7]\n",
            " [0.9 0.8 0.5 0.7]\n",
            " [0.4 0.6 0.6 0.7]\n",
            " [0.9 0.9 0.2 0.5]\n",
            " [0.1 0.2 0.  0.4]\n",
            " [0.3 0.4 0.3 0.9]\n",
            " [0.2 0.3 0.4 0.9]\n",
            " [0.  0.8 0.9 0.5]\n",
            " [0.3 0.  0.8 0. ]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Geneated queries, keys, and values for the first head.\n",
        "q0 = np.dot(x, wq0)\n",
        "k0 = np.dot(x, wk0)\n",
        "v0 = np.dot(x, wv0)\n",
        "# Geneated queries, keys, and values for the second head.\n",
        "q1 = np.dot(x, wq1)\n",
        "k1 = np.dot(x, wk1)\n",
        "v1 = np.dot(x, wv1)\n",
        "# Geneated queries, keys, and values for the third head.\n",
        "q2 = np.dot(x, wq2)\n",
        "k2= np.dot(x, wk2)\n",
        "v2 = np.dot(x, wv2)"
      ],
      "metadata": {
        "id": "3hUnQyqbX6Ge"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Q, K, and V for first head: \\n\")\n",
        "print(f\" q0 {q0.shape}: \\n\", q0, \"\\n\")\n",
        "print(f\"k0 {k0.shape}: \\n\", k0, \"\\n\")\n",
        "print(f\"v0 {v0.shape}: \\n\", v0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CfkdZYzcZl1U",
        "outputId": "1b27b75f-e98d-49c7-9843-f860db99e934"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Q, K, and V for first head: \n",
            "\n",
            " q0 (1, 3, 4): \n",
            " [[[3.05 2.99 2.36 3.21]\n",
            "  [3.55 3.05 2.56 3.78]\n",
            "  [3.51 2.79 1.84 3.29]]] \n",
            "\n",
            "k0 (1, 3, 4): \n",
            " [[[3.26 3.48 2.15 4.2 ]\n",
            "  [3.47 2.7  2.39 4.1 ]\n",
            "  [3.39 2.39 1.88 4.03]]] \n",
            "\n",
            "v0 (1, 3, 4): \n",
            " [[[2.56 3.71 2.46 2.16]\n",
            "  [2.45 3.17 3.01 2.53]\n",
            "  [2.52 3.13 3.48 2.6 ]]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now that we have our Q, K, V vectors, we can just pass them to our self-attention operation. Here we're calculating the output and attention weights for the first head."
      ],
      "metadata": {
        "id": "ALbQJ0ShZ2eA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "out0, attn_weights0 = scaled_dot_product_attention(q0, k0, v0)\n",
        "out1, attn_weights1 = scaled_dot_product_attention(q1, k1, v1)\n",
        "out2, attn_weights2 = scaled_dot_product_attention(q2, k2, v2)\n",
        "\n",
        "print(\"Output from first attention head: \", out0, \"\\n\")\n",
        "print(\"Output from second attention head: \", out1, \"\\n\")\n",
        "print(\"Output from third attention head: \", out2,)\n",
        "print(\"Attention weights from first head: \", attn_weights0, \"\\n\")\n",
        "print(\"Attention weights from second head: \", attn_weights1, \"\\n\")\n",
        "print(\"Attention weights from third head: \", attn_weights2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b1VT1ZSNZ24U",
        "outputId": "d555f966-44c2-41fd-987f-822c8f3dcfe3"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Output from first attention head:  tf.Tensor(\n",
            "[[[2.5237308 3.5005746 2.7084117 2.3069787]\n",
            "  [2.523203  3.5003133 2.7060704 2.306899 ]\n",
            "  [2.522348  3.4881163 2.727444  2.3161442]]], shape=(1, 3, 4), dtype=float32) \n",
            "\n",
            "Output from second attention head:  tf.Tensor(\n",
            "[[[1.387991  2.8214273 3.0571883 3.2637637]\n",
            "  [1.367864  2.868446  3.0218394 3.2517724]\n",
            "  [1.3553139 2.888186  3.0052195 3.241209 ]]], shape=(1, 3, 4), dtype=float32) \n",
            "\n",
            "Output from third attention head:  tf.Tensor(\n",
            "[[[2.4486806 2.6604676 2.7093885 2.7420452]\n",
            "  [2.4256036 2.601138  2.6989818 2.7314444]\n",
            "  [2.4373336 2.6289148 2.7032614 2.7362404]]], shape=(1, 3, 4), dtype=float32)\n",
            "Attention weights from first head:  tf.Tensor(\n",
            "[[[0.6182333  0.29997995 0.08178674]\n",
            "  [0.61730003 0.30698657 0.0757134 ]\n",
            "  [0.5962577  0.30717665 0.09656565]]], shape=(1, 3, 3), dtype=float32) \n",
            "\n",
            "Attention weights from second head:  tf.Tensor(\n",
            "[[[0.33176637 0.6026759  0.06555772]\n",
            "  [0.28280646 0.6534152  0.06377833]\n",
            "  [0.25085124 0.66698104 0.08216766]]], shape=(1, 3, 3), dtype=float32) \n",
            "\n",
            "Attention weights from third head:  tf.Tensor(\n",
            "[[[0.07759674 0.830996   0.09140725]\n",
            "  [0.1159062  0.766383   0.11771081]\n",
            "  [0.09540661 0.7958004  0.10879309]]], shape=(1, 3, 3), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "combined_out_a = np.concatenate((out0, out1, out2), axis=-1)\n",
        "print(f\"Combined output from all heads {combined_out_a.shape}:\")\n",
        "print(combined_out_a)\n",
        "\n",
        "\n",
        "# The final step would be to run combined_out_a through a linear/dense layer for further processing."
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "A1ccIEbeZ7Lt",
        "outputId": "f3ec84e1-47da-406b-a341-346970b4cd26"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Combined output from all heads (1, 3, 12):\n",
            "[[[2.5237308 3.5005746 2.7084117 2.3069787 1.387991  2.8214273 3.0571883\n",
            "   3.2637637 2.4486806 2.6604676 2.7093885 2.7420452]\n",
            "  [2.523203  3.5003133 2.7060704 2.306899  1.367864  2.868446  3.0218394\n",
            "   3.2517724 2.4256036 2.601138  2.6989818 2.7314444]\n",
            "  [2.522348  3.4881163 2.727444  2.3161442 1.3553139 2.888186  3.0052195\n",
            "   3.241209  2.4373336 2.6289148 2.7032614 2.7362404]]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Query weights for first head: \\n\", wq0, \"\\n\")\n",
        "print(\"Query weights for second head: \\n\", wq1, \"\\n\")\n",
        "print(\"Query weights for third head: \\n\", wq2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B_jKzisFai2F",
        "outputId": "7e36148b-b3ac-401e-fcb3-c5f54bbca8e5"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Query weights for first head: \n",
            " [[0.3 0.9 0.6 0.6]\n",
            " [0.9 0.4 0.9 0.9]\n",
            " [0.3 0.3 0.5 0.5]\n",
            " [0.3 0.8 0.1 0.7]\n",
            " [0.8 0.2 0.3 0.7]\n",
            " [0.9 0.6 0.5 0.6]\n",
            " [0.2 0.3 0.2 0.5]\n",
            " [0.3 0.7 0.3 0.2]\n",
            " [0.5 0.6 0.4 0.5]\n",
            " [0.1 0.1 0.4 0.3]\n",
            " [0.8 0.3 0.3 0.8]\n",
            " [0.9 0.8 0.4 0.4]] \n",
            "\n",
            "Query weights for second head: \n",
            " [[0.  0.4 0.6 0.5]\n",
            " [0.5 1.  0.6 0.2]\n",
            " [0.7 0.4 0.6 1. ]\n",
            " [0.8 0.2 0.5 0.7]\n",
            " [0.9 0.8 0.5 0.7]\n",
            " [0.4 0.6 0.6 0.7]\n",
            " [0.9 0.9 0.2 0.5]\n",
            " [0.1 0.2 0.  0.4]\n",
            " [0.3 0.4 0.3 0.9]\n",
            " [0.2 0.3 0.4 0.9]\n",
            " [0.  0.8 0.9 0.5]\n",
            " [0.3 0.  0.8 0. ]] \n",
            "\n",
            "Query weights for third head: \n",
            " [[0.9 0.7 0.2 0.7]\n",
            " [0.8 0.5 0.  0.2]\n",
            " [1.  0.2 0.6 0.8]\n",
            " [0.8 0.1 0.7 0.4]\n",
            " [0.3 0.8 1.  0.3]\n",
            " [0.3 0.4 0.6 0.9]\n",
            " [0.1 0.7 0.  0.7]\n",
            " [0.9 0.6 0.6 0.2]\n",
            " [0.4 0.4 0.1 0.5]\n",
            " [0.3 0.2 0.9 0.2]\n",
            " [0.9 0.9 0.1 0.2]\n",
            " [0.6 1.  0.3 0.2]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's now get the same thing done using a single query weight matrix, single key weight matrix, and single value weight matrix.\n",
        "\n",
        "Suppose instead of declaring three separate query weight matrices, we had declared one. i.e. a single  d x d  matrix. We're concatenating our per-head query weights here instead of declaring a new set of weights so that we get the same results."
      ],
      "metadata": {
        "id": "X-cHcDUCba9F"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "wq = np.concatenate((wq0, wq1, wq2), axis=1)\n",
        "wk = np.concatenate((wk0, wk1, wk2), axis=1)\n",
        "wv = np.concatenate((wv0, wv1, wv2), axis=1)\n",
        "\n",
        "print(f\"Single key weight matrix {wk.shape}:\\n\", wk, \"\\n\")\n",
        "print(f\"Single value weight matrix {wv.shape}:\\n\", wv)\n",
        "print(f\"Single query weight matrix {wq.shape}: \\n\", wq)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ln3xM6jUaqQr",
        "outputId": "45232d0f-32dc-4abd-af38-3e9544f175aa"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Single key weight matrix (12, 12):\n",
            " [[0.  0.4 0.7 0.5 0.8 0.8 0.3 0.2 0.4 0.7 0.3 0.9]\n",
            " [0.5 0.5 1.  0.9 0.6 0.4 0.8 0.5 0.8 0.2 1.  0.1]\n",
            " [0.9 0.2 0.  0.9 0.4 0.7 0.5 0.5 0.3 0.2 0.6 0.7]\n",
            " [0.1 0.4 0.  0.2 0.6 0.8 0.3 0.6 0.1 0.1 0.9 0.5]\n",
            " [0.8 0.9 0.5 0.9 0.4 0.6 0.1 0.6 0.4 0.4 0.  0.9]\n",
            " [0.8 0.4 0.6 0.9 0.1 0.5 0.3 0.6 0.9 0.2 0.4 0.7]\n",
            " [0.6 0.4 0.6 0.9 0.4 0.8 0.9 0.4 0.6 0.7 0.7 0.9]\n",
            " [0.2 0.4 0.3 0.6 0.6 0.2 0.7 0.4 0.9 0.1 0.7 0.7]\n",
            " [0.9 0.8 0.2 0.5 0.5 0.7 0.2 0.7 0.7 0.4 0.8 0.7]\n",
            " [0.2 0.9 0.1 0.8 0.8 0.1 0.8 0.9 0.6 0.5 0.8 0.4]\n",
            " [0.4 0.1 0.3 0.4 0.8 0.1 0.8 0.1 0.8 0.9 0.9 0.8]\n",
            " [0.8 0.5 0.3 1.  0.5 0.3 1.  0.4 0.4 0.4 0.5 0.7]] \n",
            "\n",
            "Single value weight matrix (12, 12):\n",
            " [[0.5 0.3 0.4 0.1 0.2 0.5 0.4 0.7 0.  0.1 0.1 0.1]\n",
            " [0.3 0.2 0.1 0.4 0.5 1.  0.1 0.8 0.8 0.4 0.6 0.3]\n",
            " [0.9 0.6 0.3 0.3 0.2 0.4 0.5 0.8 0.7 0.9 0.8 0.5]\n",
            " [0.3 0.7 0.4 0.3 0.  0.7 1.  0.4 0.9 0.  0.2 0.9]\n",
            " [0.5 1.  0.4 0.5 0.8 0.7 0.8 1.  0.4 0.1 0.5 0. ]\n",
            " [0.  0.6 0.4 0.2 0.4 0.2 1.  0.7 0.3 0.1 1.  1. ]\n",
            " [0.8 0.9 0.8 0.5 0.2 0.8 0.  0.2 0.2 0.4 0.  0.4]\n",
            " [0.7 0.9 0.7 0.5 0.2 0.8 0.1 0.3 0.1 0.7 0.3 0.9]\n",
            " [0.3 0.6 0.4 0.1 0.1 0.2 0.6 0.6 0.  0.9 0.1 0.2]\n",
            " [0.6 0.8 0.1 0.9 0.2 0.  0.2 0.  0.4 0.1 0.8 0.6]\n",
            " [0.2 0.3 0.9 1.  0.  0.7 0.  0.1 0.2 0.8 0.7 0.2]\n",
            " [0.7 0.1 0.9 0.3 0.2 0.  0.8 0.9 0.6 0.2 0.1 0.3]]\n",
            "Single query weight matrix (12, 12): \n",
            " [[0.3 0.9 0.6 0.6 0.  0.4 0.6 0.5 0.9 0.7 0.2 0.7]\n",
            " [0.9 0.4 0.9 0.9 0.5 1.  0.6 0.2 0.8 0.5 0.  0.2]\n",
            " [0.3 0.3 0.5 0.5 0.7 0.4 0.6 1.  1.  0.2 0.6 0.8]\n",
            " [0.3 0.8 0.1 0.7 0.8 0.2 0.5 0.7 0.8 0.1 0.7 0.4]\n",
            " [0.8 0.2 0.3 0.7 0.9 0.8 0.5 0.7 0.3 0.8 1.  0.3]\n",
            " [0.9 0.6 0.5 0.6 0.4 0.6 0.6 0.7 0.3 0.4 0.6 0.9]\n",
            " [0.2 0.3 0.2 0.5 0.9 0.9 0.2 0.5 0.1 0.7 0.  0.7]\n",
            " [0.3 0.7 0.3 0.2 0.1 0.2 0.  0.4 0.9 0.6 0.6 0.2]\n",
            " [0.5 0.6 0.4 0.5 0.3 0.4 0.3 0.9 0.4 0.4 0.1 0.5]\n",
            " [0.1 0.1 0.4 0.3 0.2 0.3 0.4 0.9 0.3 0.2 0.9 0.2]\n",
            " [0.8 0.3 0.3 0.8 0.  0.8 0.9 0.5 0.9 0.9 0.1 0.2]\n",
            " [0.9 0.8 0.4 0.4 0.3 0.  0.8 0.  0.6 1.  0.3 0.2]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "q_s = np.dot(x, wq)\n",
        "k_s = np.dot(x, wk)\n",
        "v_s = np.dot(x, wv)"
      ],
      "metadata": {
        "id": "g7WHDacEa65R"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(f\"Query vectors using a single weight matrix {q_s.shape}:\\n\", q_s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mwZu2Qn3a7Ni",
        "outputId": "ff1990a0-5c5e-4ce5-cfd6-66319716a315"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Query vectors using a single weight matrix (1, 3, 12):\n",
            " [[[3.05 2.99 2.36 3.21 2.57 2.75 2.91 3.98 2.97 2.93 3.01 2.94]\n",
            "  [3.55 3.05 2.56 3.78 2.66 3.42 3.32 3.59 3.74 3.16 1.89 2.83]\n",
            "  [3.51 2.79 1.84 3.29 2.72 2.82 3.31 3.03 3.15 3.5  2.38 2.53]]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(q0, \"\\n\")\n",
        "print(q1, \"\\n\")\n",
        "print(q2)\n",
        "\n",
        "# we can see how combined and seperated query values are identical "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gMyg9LxZa9ud",
        "outputId": "3883405b-ed29-42e6-cc2e-8bb635f6f418"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[[3.05 2.99 2.36 3.21]\n",
            "  [3.55 3.05 2.56 3.78]\n",
            "  [3.51 2.79 1.84 3.29]]] \n",
            "\n",
            "[[[2.57 2.75 2.91 3.98]\n",
            "  [2.66 3.42 3.32 3.59]\n",
            "  [2.72 2.82 3.31 3.03]]] \n",
            "\n",
            "[[[2.97 2.93 3.01 2.94]\n",
            "  [3.74 3.16 1.89 2.83]\n",
            "  [3.15 3.5  2.38 2.53]]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can split our combined queries into  d x d/h  heads using reshape and transpose"
      ],
      "metadata": {
        "id": "bKfCVST9cEqI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Note: we can achieve the same thing by passing -1 instead of seq_len.\n",
        "q_s_reshaped = tf.reshape(q_s, (batch_size, seq_len, num_heads, head_dim))\n",
        "print(f\"Combined queries: {q_s.shape}\\n\", q_s, \"\\n\")\n",
        "print(f\"Reshaped into separate heads: {q_s_reshaped.shape}\\n\", q_s_reshaped)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "woTMv9ETbAJS",
        "outputId": "a7b6274f-e60a-420e-9372-a3c5575a7b3e"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Combined queries: (1, 3, 12)\n",
            " [[[3.05 2.99 2.36 3.21 2.57 2.75 2.91 3.98 2.97 2.93 3.01 2.94]\n",
            "  [3.55 3.05 2.56 3.78 2.66 3.42 3.32 3.59 3.74 3.16 1.89 2.83]\n",
            "  [3.51 2.79 1.84 3.29 2.72 2.82 3.31 3.03 3.15 3.5  2.38 2.53]]] \n",
            "\n",
            "Reshaped into separate heads: (1, 3, 3, 4)\n",
            " tf.Tensor(\n",
            "[[[[3.05 2.99 2.36 3.21]\n",
            "   [2.57 2.75 2.91 3.98]\n",
            "   [2.97 2.93 3.01 2.94]]\n",
            "\n",
            "  [[3.55 3.05 2.56 3.78]\n",
            "   [2.66 3.42 3.32 3.59]\n",
            "   [3.74 3.16 1.89 2.83]]\n",
            "\n",
            "  [[3.51 2.79 1.84 3.29]\n",
            "   [2.72 2.82 3.31 3.03]\n",
            "   [3.15 3.5  2.38 2.53]]]], shape=(1, 3, 3, 4), dtype=float64)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "q_s_transposed = tf.transpose(q_s_reshaped, perm=[0, 2, 1, 3]).numpy()\n",
        "print(f\"Queries transposed into \\\"separate\\\" heads {q_s_transposed.shape}:\\n\", \n",
        "      q_s_transposed)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ixh5DDxJbzNV",
        "outputId": "cee058d1-e37e-4b87-eb4a-c617b70e3504"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Queries transposed into \"separate\" heads (1, 3, 3, 4):\n",
            " [[[[3.05 2.99 2.36 3.21]\n",
            "   [3.55 3.05 2.56 3.78]\n",
            "   [3.51 2.79 1.84 3.29]]\n",
            "\n",
            "  [[2.57 2.75 2.91 3.98]\n",
            "   [2.66 3.42 3.32 3.59]\n",
            "   [2.72 2.82 3.31 3.03]]\n",
            "\n",
            "  [[2.97 2.93 3.01 2.94]\n",
            "   [3.74 3.16 1.89 2.83]\n",
            "   [3.15 3.5  2.38 2.53]]]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"The separate per-head query matrices from before: \")\n",
        "print(q0, \"\\n\")\n",
        "print(q1, \"\\n\")\n",
        "print(q2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fLx6109Ub2QE",
        "outputId": "d5034a05-31e9-42a1-e35f-68937face9c6"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The separate per-head query matrices from before: \n",
            "[[[3.05 2.99 2.36 3.21]\n",
            "  [3.55 3.05 2.56 3.78]\n",
            "  [3.51 2.79 1.84 3.29]]] \n",
            "\n",
            "[[[2.57 2.75 2.91 3.98]\n",
            "  [2.66 3.42 3.32 3.59]\n",
            "  [2.72 2.82 3.31 3.03]]] \n",
            "\n",
            "[[[2.97 2.93 3.01 2.94]\n",
            "  [3.74 3.16 1.89 2.83]\n",
            "  [3.15 3.5  2.38 2.53]]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "k_s_transposed = tf.transpose(tf.reshape(k_s, (batch_size, -1, num_heads, head_dim)), perm=[0, 2, 1, 3]).numpy()\n",
        "v_s_transposed = tf.transpose(tf.reshape(v_s, (batch_size, -1, num_heads, head_dim)), perm=[0, 2, 1, 3]).numpy()\n",
        "\n",
        "print(f\"Keys for all heads in a single matrix {k_s.shape}: \\n\", k_s_transposed, \"\\n\")\n",
        "print(f\"Values for all heads in a single matrix {v_s.shape}: \\n\", v_s_transposed)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yUdgwVqUcPaG",
        "outputId": "066e8c1f-7ee1-4508-ab1e-7de6282b38b1"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Keys for all heads in a single matrix (1, 3, 12): \n",
            " [[[[3.26 3.48 2.15 4.2 ]\n",
            "   [3.47 2.7  2.39 4.1 ]\n",
            "   [3.39 2.39 1.88 4.03]]\n",
            "\n",
            "  [[3.   3.24 2.58 3.34]\n",
            "   [3.2  3.16 3.45 2.93]\n",
            "   [2.72 2.73 3.33 2.51]]\n",
            "\n",
            "  [[3.25 2.35 3.23 4.1 ]\n",
            "   [3.69 2.51 4.45 3.86]\n",
            "   [3.06 2.48 3.39 4.11]]]] \n",
            "\n",
            "Values for all heads in a single matrix (1, 3, 12): \n",
            " [[[[2.56 3.71 2.46 2.16]\n",
            "   [2.45 3.17 3.01 2.53]\n",
            "   [2.52 3.13 3.48 2.6 ]]\n",
            "\n",
            "  [[1.66 2.23 3.51 3.44]\n",
            "   [1.25 3.17 2.8  3.19]\n",
            "   [1.28 2.61 3.13 3.05]]\n",
            "\n",
            "  [[1.98 1.79 2.64 2.61]\n",
            "   [2.5  2.81 2.74 2.77]\n",
            "   [2.38 2.04 2.49 2.6 ]]]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Set up this way, we can now calculate the outputs from all attention heads with a single call to our self-attention operation."
      ],
      "metadata": {
        "id": "znDzkKj6cW_W"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "all_heads_output, all_attn_weights = scaled_dot_product_attention(q_s_transposed, \n",
        "                                                                  k_s_transposed, \n",
        "                                                                  v_s_transposed)\n",
        "print(\"Self attention output:\\n\", all_heads_output)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SsGo_riBcS5h",
        "outputId": "f1b5e9d6-232e-491f-930b-822d6d7ddd8d"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Self attention output:\n",
            " tf.Tensor(\n",
            "[[[[2.5237305 3.500574  2.7084112 2.3069787]\n",
            "   [2.523203  3.500313  2.70607   2.3068988]\n",
            "   [2.522348  3.4881163 2.727444  2.3161442]]\n",
            "\n",
            "  [[1.387991  2.8214273 3.0571883 3.2637637]\n",
            "   [1.3678638 2.8684459 3.0218391 3.2517724]\n",
            "   [1.3553139 2.888186  3.0052195 3.241209 ]]\n",
            "\n",
            "  [[2.4486809 2.6604676 2.7093887 2.7420454]\n",
            "   [2.4256034 2.601138  2.6989818 2.7314444]\n",
            "   [2.4373336 2.6289148 2.7032614 2.7362404]]]], shape=(1, 3, 3, 4), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Per head outputs from using separate sets of weights per head:\")\n",
        "print(out0, \"\\n\")\n",
        "print(out1, \"\\n\")\n",
        "print(out2)\n",
        "\n",
        "#both seperate and combined results are again identical"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e1bU7IZ2cYGp",
        "outputId": "eef14312-020f-4f45-a624-6515457a421b"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Per head outputs from using separate sets of weights per head:\n",
            "tf.Tensor(\n",
            "[[[2.5237308 3.5005746 2.7084117 2.3069787]\n",
            "  [2.523203  3.5003133 2.7060704 2.306899 ]\n",
            "  [2.522348  3.4881163 2.727444  2.3161442]]], shape=(1, 3, 4), dtype=float32) \n",
            "\n",
            "tf.Tensor(\n",
            "[[[1.387991  2.8214273 3.0571883 3.2637637]\n",
            "  [1.367864  2.868446  3.0218394 3.2517724]\n",
            "  [1.3553139 2.888186  3.0052195 3.241209 ]]], shape=(1, 3, 4), dtype=float32) \n",
            "\n",
            "tf.Tensor(\n",
            "[[[2.4486806 2.6604676 2.7093885 2.7420452]\n",
            "  [2.4256036 2.601138  2.6989818 2.7314444]\n",
            "  [2.4373336 2.6289148 2.7032614 2.7362404]]], shape=(1, 3, 4), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "combined_out_b = tf.reshape(tf.transpose(all_heads_output, perm=[0, 2, 1, 3]), \n",
        "                            shape=(batch_size, seq_len, embed_dim))\n",
        "print(\"Final output from using single query, key, value matrices:\\n\", \n",
        "      combined_out_b, \"\\n\")\n",
        "print(\"Final output from using separate query, key, value matrices per head:\\n\", \n",
        "      combined_out_a)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GIgWLvaJcaaW",
        "outputId": "654b80f9-135f-40b8-c311-107930cbf2f4"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Final output from using single query, key, value matrices:\n",
            " tf.Tensor(\n",
            "[[[2.5237305 3.500574  2.7084112 2.3069787 1.387991  2.8214273 3.0571883\n",
            "   3.2637637 2.4486809 2.6604676 2.7093887 2.7420454]\n",
            "  [2.523203  3.500313  2.70607   2.3068988 1.3678638 2.8684459 3.0218391\n",
            "   3.2517724 2.4256034 2.601138  2.6989818 2.7314444]\n",
            "  [2.522348  3.4881163 2.727444  2.3161442 1.3553139 2.888186  3.0052195\n",
            "   3.241209  2.4373336 2.6289148 2.7032614 2.7362404]]], shape=(1, 3, 12), dtype=float32) \n",
            "\n",
            "Final output from using separate query, key, value matrices per head:\n",
            " [[[2.5237308 3.5005746 2.7084117 2.3069787 1.387991  2.8214273 3.0571883\n",
            "   3.2637637 2.4486806 2.6604676 2.7093885 2.7420452]\n",
            "  [2.523203  3.5003133 2.7060704 2.306899  1.367864  2.868446  3.0218394\n",
            "   3.2517724 2.4256036 2.601138  2.6989818 2.7314444]\n",
            "  [2.522348  3.4881163 2.727444  2.3161442 1.3553139 2.888186  3.0052195\n",
            "   3.241209  2.4373336 2.6289148 2.7032614 2.7362404]]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##We can encapsulate everything we just covered in a single class."
      ],
      "metadata": {
        "id": "b88iIn1KcxAY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MultiHeadSelfAttention(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads):\n",
        "    super(MultiHeadSelfAttention, self).__init__()\n",
        "    self.d_model = d_model\n",
        "    self.num_heads = num_heads\n",
        "\n",
        "    self.d_head = self.d_model // self.num_heads\n",
        "\n",
        "    self.wq = tf.keras.layers.Dense(self.d_model)\n",
        "    self.wk = tf.keras.layers.Dense(self.d_model)\n",
        "    self.wv = tf.keras.layers.Dense(self.d_model)\n",
        "\n",
        "    # Linear layer to generate the final output.\n",
        "    self.dense = tf.keras.layers.Dense(self.d_model)\n",
        "  \n",
        "  def split_heads(self, x):\n",
        "    batch_size = x.shape[0]\n",
        "\n",
        "    split_inputs = tf.reshape(x, (batch_size, -1, self.num_heads, self.d_head))\n",
        "    return tf.transpose(split_inputs, perm=[0, 2, 1, 3])\n",
        "  \n",
        "  def merge_heads(self, x):\n",
        "    batch_size = x.shape[0]\n",
        "\n",
        "    merged_inputs = tf.transpose(x, perm=[0, 2, 1, 3])\n",
        "    return tf.reshape(merged_inputs, (batch_size, -1, self.d_model))\n",
        "\n",
        "  def call(self, q, k, v, mask):\n",
        "    qs = self.wq(q)\n",
        "    ks = self.wk(k)\n",
        "    vs = self.wv(v)\n",
        "\n",
        "    qs = self.split_heads(qs)\n",
        "    ks = self.split_heads(ks)\n",
        "    vs = self.split_heads(vs)\n",
        "\n",
        "    output, attn_weights = scaled_dot_product_attention(qs, ks, vs, mask)\n",
        "    output = self.merge_heads(output)\n",
        "\n",
        "    return self.dense(output), attn_weights\n"
      ],
      "metadata": {
        "id": "7V5WzVEeciDj"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "mhsa = MultiHeadSelfAttention(12, 3)\n",
        "\n",
        "output, attn_weights = mhsa(x, x, x, None)\n",
        "print(f\"MHSA output{output.shape}:\")\n",
        "print(output)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Rx74hhBVcps5",
        "outputId": "6792b105-f697-443d-e749-96e62166ecc0"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "MHSA output(1, 3, 12):\n",
            "tf.Tensor(\n",
            "[[[ 0.2294916   0.12367499  0.5002226  -0.05087875  0.44209313\n",
            "    0.04899372  0.13054031  0.12503248  0.14715973  0.45378336\n",
            "   -0.3372569  -0.30783698]\n",
            "  [ 0.24200596  0.12704034  0.5397398  -0.06036338  0.5002323\n",
            "    0.04581532  0.17330357  0.15491767  0.19630358  0.4839844\n",
            "   -0.3624456  -0.31827798]\n",
            "  [ 0.24138258  0.12508693  0.53865886 -0.06811577  0.45645258\n",
            "    0.02451682  0.1569365   0.14821433  0.182532    0.4636302\n",
            "   -0.3310364  -0.31074744]]], shape=(1, 3, 12), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##ENCODER"
      ],
      "metadata": {
        "id": "UzcDGIjwdZ7a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def feed_forward_network(d_model, hidden_dim):\n",
        "  return tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(hidden_dim, activation='relu'),\n",
        "      tf.keras.layers.Dense(d_model)\n",
        "  ])"
      ],
      "metadata": {
        "id": "H075jw2nc5jx"
      },
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class EncoderBlock(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads, hidden_dim, dropout_rate=0.1):\n",
        "    super(EncoderBlock, self).__init__()\n",
        "\n",
        "    self.mhsa = MultiHeadSelfAttention(d_model, num_heads)\n",
        "    self.ffn = feed_forward_network(d_model, hidden_dim)\n",
        "\n",
        "    self.dropout1 = tf.keras.layers.Dropout(dropout_rate)\n",
        "    self.dropout2 = tf.keras.layers.Dropout(dropout_rate)\n",
        "\n",
        "    self.layernorm1 = tf.keras.layers.LayerNormalization()\n",
        "    self.layernorm2 = tf.keras.layers.LayerNormalization()\n",
        "  \n",
        "  def call(self, x, training, mask):\n",
        "    mhsa_output, attn_weights = self.mhsa(x, x, x, mask)\n",
        "    mhsa_output = self.dropout1(mhsa_output, training=training)\n",
        "    mhsa_output = self.layernorm1(x + mhsa_output)\n",
        "\n",
        "    ffn_output = self.ffn(mhsa_output)\n",
        "    ffn_output = self.dropout2(ffn_output, training=training)\n",
        "    output = self.layernorm2(mhsa_output + ffn_output)\n",
        "\n",
        "    return output, attn_weights\n"
      ],
      "metadata": {
        "id": "Uh08PWmVddo5"
      },
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "encoder_block = EncoderBlock(12, 3, 48)\n",
        "\n",
        "block_output,  _ = encoder_block(x, True, None)\n",
        "print(f\"Output from single encoder block {block_output.shape}:\")\n",
        "print(block_output)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FxgrcTNvde_B",
        "outputId": "c1a23e87-7bec-46c8-9eee-f16f3daca7ea"
      },
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Output from single encoder block (1, 3, 12):\n",
            "tf.Tensor(\n",
            "[[[-2.1197195e-01 -1.3177364e-01  1.5815003e-01  1.0531749e+00\n",
            "   -1.3040812e+00  4.0824369e-02 -3.9061621e-01 -2.8542405e-01\n",
            "    1.0904245e+00  2.2544982e+00 -1.1763124e+00 -1.0968924e+00]\n",
            "  [-6.3340265e-01 -6.7885280e-02  1.6601526e+00  1.2038925e+00\n",
            "   -1.7723087e+00 -2.6691768e-01 -7.9801881e-01  7.8692302e-02\n",
            "    9.3915153e-01  9.6578437e-01 -8.1238325e-04 -1.3083278e+00]\n",
            "  [-1.1207460e+00 -5.5309528e-01  2.0188301e+00  1.3985721e+00\n",
            "   -1.2443501e+00 -8.0472708e-01 -3.1493428e-01 -2.1804489e-02\n",
            "    6.3748491e-01  7.4244398e-01  2.9259631e-01 -1.0302701e+00]]], shape=(1, 3, 12), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Word and Positional Embeddings"
      ],
      "metadata": {
        "id": "YfqCh7TAdnHX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the English tokenizer.\n",
        "bpemb_en = BPEmb(lang=\"en\")\n",
        "bpemb_vocab_size, bpemb_embed_size = bpemb_en.vectors.shape\n",
        "print(\"Vocabulary size:\", bpemb_vocab_size)\n",
        "print(\"Embedding size:\", bpemb_embed_size)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YWAN-24sdhOw",
        "outputId": "a261dd9d-d07a-4571-f14e-1efdd9a65a7d"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "downloading https://nlp.h-its.org/bpemb/en/en.wiki.bpe.vs10000.model\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 400869/400869 [00:00<00:00, 1060653.81B/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "downloading https://nlp.h-its.org/bpemb/en/en.wiki.bpe.vs10000.d100.w2v.bin.tar.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 3784656/3784656 [00:00<00:00, 4791288.21B/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Vocabulary size: 10000\n",
            "Embedding size: 100\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Embedding for the word \"car\".\n",
        "bpemb_en.vectors[bpemb_en.words.index('car')]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "amXRZFATdvSo",
        "outputId": "4a307be0-f39c-461d-eac5-4715f385c742"
      },
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([-0.305548, -0.325598, -0.134716, -0.078735, -0.660545,  0.076211,\n",
              "       -0.735487,  0.124533, -0.294402,  0.459688,  0.030137,  0.174041,\n",
              "       -0.224223,  0.486189, -0.504649, -0.459699,  0.315747,  0.477885,\n",
              "        0.091398,  0.427867,  0.016524, -0.076833, -0.899727,  0.493158,\n",
              "       -0.022309, -0.422785, -0.154148,  0.204981,  0.379834,  0.070588,\n",
              "        0.196073, -0.368222,  0.473406,  0.007409,  0.004303, -0.007823,\n",
              "       -0.19103 , -0.202509,  0.109878, -0.224521, -0.35741 , -0.611633,\n",
              "        0.329958, -0.212956, -0.497499, -0.393839, -0.130101, -0.216903,\n",
              "       -0.105595, -0.076007, -0.483942, -0.139704, -0.161647,  0.136985,\n",
              "        0.415363, -0.360143,  0.038601, -0.078804, -0.030421,  0.324129,\n",
              "        0.223378, -0.523636, -0.048317, -0.032248, -0.117367,  0.470519,\n",
              "        0.225816, -0.222065, -0.225007, -0.165904, -0.334389, -0.20157 ,\n",
              "        0.572352, -0.268794,  0.301929, -0.005563,  0.387491,  0.261031,\n",
              "       -0.11613 ,  0.074982, -0.008433,  0.259987, -0.099893, -0.268875,\n",
              "       -0.054047, -0.534776, -0.111101, -0.051742,  0.214114,  0.04293 ,\n",
              "        0.039873, -0.453112,  0.087382, -0.333201, -0.034079, -0.833045,\n",
              "        0.155232, -1.132393, -0.294766,  0.327572], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 30
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We don't need the embeddings since we're going to use our own embedding layer. What we're interested in are the subword tokens and their respective ids. The ids will be used as indexes into our embedding layer.<br>\n",
        "\n",
        "These are the subword tokens for our example sentence from the slides. **BPEmb** places underscores in front of any tokens which are whole words or intended to begin words.<br>\n",
        "\n",
        "Remember that subword tokenizers are trained using count frequencies over a corpus. So these subword tokens are specific to **BPEmb**. Another subword tokenizer may output something different. This is why it's important that when we use a pretrained model, we make sure to use the pretrained model's tokenizer. We'll see this when we use pretrained transformers later in this module."
      ],
      "metadata": {
        "id": "Ak0s9Mkfe93q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sample_sentence = \"Where can I find a pizzeria?\"\n",
        "tokens = bpemb_en.encode(sample_sentence)\n",
        "print(tokens)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ggdJ-TWUdvwt",
        "outputId": "a9a7c431-6886-4441-f3d7-acbca408adab"
      },
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['▁where', '▁can', '▁i', '▁find', '▁a', '▁p', 'iz', 'zer', 'ia', '?']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "token_seq = np.array(bpemb_en.encode_ids(\"Where can I find a pizzeria?\"))\n",
        "print(token_seq)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "j9kpKZCad16n",
        "outputId": "6e558c89-70b2-4f4d-de18-2a75057cfd28"
      },
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[ 571  280  386 1934    4   24  248 4339  177 9967]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now that we have a way to tokenize and vectorize sentences, we can declare and use an embedding layer with the same vocabulary size as **BPEmb** and a desired embedding size."
      ],
      "metadata": {
        "id": "1OPgyDvQeuNI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "token_embed = tf.keras.layers.Embedding(bpemb_vocab_size, embed_dim)\n",
        "token_embeddings = token_embed(token_seq)\n",
        "\n",
        "# The untrained embeddings for our sample sentence.\n",
        "print(\"Embeddings for: \", sample_sentence)\n",
        "print(token_embeddings)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TUKQNRUTd4NO",
        "outputId": "1106a563-955c-4eb8-fa50-ac5161873ff5"
      },
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Embeddings for:  Where can I find a pizzeria?\n",
            "tf.Tensor(\n",
            "[[-0.01302861  0.03313113 -0.04513763 -0.00997768 -0.02491695 -0.01981989\n",
            "  -0.03860885 -0.0486939  -0.03882251 -0.02593344  0.04286231  0.00894827]\n",
            " [-0.01603675  0.01845826  0.0008552   0.02768201 -0.02777867 -0.00649695\n",
            "   0.00868416 -0.03917243 -0.01836915  0.02388283  0.0146642   0.01673187]\n",
            " [-0.04333204 -0.01979154  0.0175184  -0.044195   -0.02448606 -0.01857346\n",
            "  -0.00193138  0.01074655  0.02854523 -0.03246266  0.01003687  0.03762888]\n",
            " [ 0.04620275  0.02691752 -0.02910745 -0.02630258 -0.01229552 -0.00610255\n",
            "  -0.01044874  0.01956037  0.04262752 -0.04155936 -0.01917005 -0.01994878]\n",
            " [-0.03539585  0.00268234 -0.01060759 -0.01913463 -0.02754086  0.01332739\n",
            "  -0.04862207 -0.01495548  0.02292304  0.02027467 -0.00249114 -0.01782548]\n",
            " [ 0.02176336  0.04785791 -0.03956974  0.00849892  0.00147782 -0.01380372\n",
            "  -0.02697941 -0.0199378   0.0121421  -0.00711564 -0.02658532 -0.02139275]\n",
            " [-0.04082402  0.01951199 -0.02894517  0.0055403   0.01717676 -0.02033845\n",
            "  -0.0492411  -0.04865298  0.02684506  0.01716817  0.02170327 -0.02142456]\n",
            " [-0.04276207  0.03329197  0.00858668 -0.0248852  -0.01038573  0.00981808\n",
            "  -0.01670846 -0.01828521  0.02390087 -0.01332563 -0.02162116 -0.02442855]\n",
            " [-0.04565329 -0.00583225 -0.04463512  0.01812253  0.00633087  0.00535967\n",
            "   0.04065325  0.00667412  0.00202281 -0.02632576 -0.00135802  0.00770247]\n",
            " [ 0.01610572  0.03590734 -0.01354749 -0.01185226 -0.00454631  0.01867985\n",
            "  -0.02560574 -0.03229456 -0.03169658 -0.00980771  0.02210565  0.02151996]], shape=(10, 12), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "max_seq_len = 256\n",
        "pos_embed = tf.keras.layers.Embedding(max_seq_len, embed_dim)\n",
        "\n",
        "# Generate ids for each position of the token sequence.\n",
        "pos_idx = tf.range(len(token_seq))\n",
        "print(pos_idx)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "585SfALed8r8",
        "outputId": "89f228a8-55d9-4c67-b287-e7ba1167da57"
      },
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tf.Tensor([0 1 2 3 4 5 6 7 8 9], shape=(10,), dtype=int32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# These are our positon embeddings.\n",
        "position_embeddings = pos_embed(pos_idx)\n",
        "print(\"Position embeddings for the input sequence\\n\", position_embeddings)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4oGsyEslfKbq",
        "outputId": "493246fd-8f92-4138-fd78-ef02b22e2ddb"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Position embeddings for the input sequence\n",
            " tf.Tensor(\n",
            "[[-0.00267386  0.04497888 -0.01898999  0.02883989 -0.00654338 -0.03075428\n",
            "   0.01430224 -0.03397721  0.01428051 -0.03040561  0.00714212  0.02189949]\n",
            " [ 0.00583997  0.00416914  0.01025517 -0.01858835  0.01152635  0.0453545\n",
            "   0.00875487 -0.04249355  0.02623383 -0.00830251  0.04689712 -0.04844732]\n",
            " [-0.0395951  -0.00850092 -0.04495001  0.04774108 -0.01864243  0.03377538\n",
            "  -0.01878288  0.04108479 -0.03097264  0.04437483 -0.02057176  0.04371068]\n",
            " [ 0.00818767  0.0161379   0.00114429 -0.03537606  0.04595914  0.02060746\n",
            "  -0.03434259  0.03872677 -0.01169735  0.0408063  -0.0116725   0.01446274]\n",
            " [ 0.02735185  0.01735767 -0.034096   -0.01159757 -0.02599738  0.03450202\n",
            "  -0.03427508  0.02789124  0.04310719 -0.0003692  -0.0216522   0.02385452]\n",
            " [-0.04487137  0.0297712   0.00572839 -0.04311861  0.03395119 -0.03924471\n",
            "  -0.02079718  0.04100347  0.02254671 -0.02317665  0.01228994  0.04094117]\n",
            " [ 0.03096701 -0.0092387  -0.00644468  0.04171463  0.03103564  0.04777172\n",
            "  -0.04670488 -0.04088136  0.01079168 -0.0099802  -0.04237262 -0.04883695]\n",
            " [-0.00798281  0.00167825 -0.04544545  0.0009286  -0.0068844   0.03416096\n",
            "  -0.02633656  0.0047374  -0.03104993 -0.04159446 -0.01590217 -0.04526911]\n",
            " [ 0.04720448  0.02071989  0.02709576 -0.0104849   0.02080638  0.04495866\n",
            "  -0.04587611 -0.00335868 -0.0272548   0.03617242 -0.0006305  -0.03822846]\n",
            " [-0.01019701  0.01029786  0.01823778 -0.03126882  0.01395769 -0.02387906\n",
            "  -0.02165006  0.03877784 -0.00389426 -0.04155596 -0.02412431 -0.01842468]], shape=(10, 12), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "input = token_embeddings + position_embeddings\n",
        "print(\"Input to the initial encoder block:\\n\", input)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YqKFBZh-fNPh",
        "outputId": "f0733191-ee62-41a9-8158-382013b244f6"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input to the initial encoder block:\n",
            " tf.Tensor(\n",
            "[[-0.01570247  0.07811001 -0.06412762  0.01886221 -0.03146032 -0.05057417\n",
            "  -0.02430661 -0.08267111 -0.024542   -0.05633905  0.05000442  0.03084775]\n",
            " [-0.01019678  0.0226274   0.01111037  0.00909366 -0.01625233  0.03885754\n",
            "   0.01743903 -0.08166598  0.00786468  0.01558032  0.06156132 -0.03171545]\n",
            " [-0.08292715 -0.02829247 -0.0274316   0.00354609 -0.04312849  0.01520192\n",
            "  -0.02071426  0.05183134 -0.0024274   0.01191217 -0.01053488  0.08133955]\n",
            " [ 0.05439042  0.04305542 -0.02796316 -0.06167864  0.03366362  0.01450491\n",
            "  -0.04479133  0.05828714  0.03093017 -0.00075306 -0.03084254 -0.00548605]\n",
            " [-0.008044    0.02004001 -0.04470359 -0.0307322  -0.05353824  0.04782941\n",
            "  -0.08289715  0.01293576  0.06603023  0.01990546 -0.02414334  0.00602904]\n",
            " [-0.02310801  0.07762911 -0.03384135 -0.03461969  0.035429   -0.05304843\n",
            "  -0.04777659  0.02106567  0.03468881 -0.03029229 -0.01429537  0.01954842]\n",
            " [-0.00985701  0.01027329 -0.03538986  0.04725493  0.04821239  0.02743327\n",
            "  -0.09594598 -0.08953434  0.03763675  0.00718797 -0.02066935 -0.0702615 ]\n",
            " [-0.05074488  0.03497022 -0.03685877 -0.0239566  -0.01727013  0.04397904\n",
            "  -0.04304502 -0.01354782 -0.00714906 -0.05492009 -0.03752333 -0.06969766]\n",
            " [ 0.00155119  0.01488763 -0.01753936  0.00763763  0.02713725  0.05031833\n",
            "  -0.00522286  0.00331544 -0.02523198  0.00984665 -0.00198852 -0.03052599]\n",
            " [ 0.00590872  0.0462052   0.00469029 -0.04312109  0.00941138 -0.00519922\n",
            "  -0.0472558   0.00648328 -0.03559084 -0.05136367 -0.00201867  0.00309528]], shape=(10, 12), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class Encoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, num_blocks, d_model, num_heads, hidden_dim, src_vocab_size,\n",
        "               max_seq_len, dropout_rate=0.1):\n",
        "    super(Encoder, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.max_seq_len = max_seq_len\n",
        "\n",
        "    self.token_embed = tf.keras.layers.Embedding(src_vocab_size, self.d_model)\n",
        "    self.pos_embed = tf.keras.layers.Embedding(max_seq_len, self.d_model)\n",
        "\n",
        "    # The original Attention Is All You Need paper applied dropout to the\n",
        "    # input before feeding it to the first encoder block.\n",
        "    self.dropout = tf.keras.layers.Dropout(dropout_rate)\n",
        "\n",
        "    # Create encoder blocks.\n",
        "    self.blocks = [EncoderBlock(self.d_model, num_heads, hidden_dim, dropout_rate) \n",
        "    for _ in range(num_blocks)]\n",
        "  \n",
        "  def call(self, input, training, mask):\n",
        "    token_embeds = self.token_embed(input)\n",
        "\n",
        "    # Generate position indices for a batch of input sequences.\n",
        "    num_pos = input.shape[0] * self.max_seq_len\n",
        "    pos_idx = np.resize(np.arange(self.max_seq_len), num_pos)\n",
        "    pos_idx = np.reshape(pos_idx, input.shape)\n",
        "    pos_embeds = self.pos_embed(pos_idx)\n",
        "\n",
        "    x = self.dropout(token_embeds + pos_embeds, training=training)\n",
        "\n",
        "    # Run input through successive encoder blocks.\n",
        "    for block in self.blocks:\n",
        "      x, weights = block(x, training, mask)\n",
        "\n",
        "    return x, weights"
      ],
      "metadata": {
        "id": "hjcFzgQ3fNRz"
      },
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Batch of 3 sequences, each of length 10 (10 is also the \n",
        "# maximum sequence length in this case).\n",
        "seqs = np.random.randint(0, 10000, size=(3, 10))\n",
        "print(seqs.shape)\n",
        "print(seqs)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "anfWaCkMfNUa",
        "outputId": "1961a9e1-0dc8-4461-bbf3-d899fde881a6"
      },
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(3, 10)\n",
            "[[7886 2242 8537 8926 8527 7208 1972 3828 9959 8853]\n",
            " [1099 9474 9773 6282 3514 4379 2932 9802 8522 7832]\n",
            " [4356 5064 7960 6436 5539 9761 8952 2611 4558 8429]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pos_ids = np.resize(np.arange(seqs.shape[1]), seqs.shape[0] * seqs.shape[1])\n",
        "print(pos_ids)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "O77tLFKDfbSf",
        "outputId": "666d5502-0754-4737-a2de-9920a97c856c"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pos_ids = np.reshape(pos_ids, (3, 10))\n",
        "print(pos_ids.shape)\n",
        "print(pos_ids)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lu_mFYEPfdAP",
        "outputId": "87f92b2a-b5ca-4216-ebdb-a4404564ee7c"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(3, 10)\n",
            "[[0 1 2 3 4 5 6 7 8 9]\n",
            " [0 1 2 3 4 5 6 7 8 9]\n",
            " [0 1 2 3 4 5 6 7 8 9]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pos_embed(pos_ids)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "m3PZBuZrfeWm",
        "outputId": "c746f87c-a9af-41e5-ce45-5a9ed13d8da0"
      },
      "execution_count": 41,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: shape=(3, 10, 12), dtype=float32, numpy=\n",
              "array([[[-0.00267386,  0.04497888, -0.01898999,  0.02883989,\n",
              "         -0.00654338, -0.03075428,  0.01430224, -0.03397721,\n",
              "          0.01428051, -0.03040561,  0.00714212,  0.02189949],\n",
              "        [ 0.00583997,  0.00416914,  0.01025517, -0.01858835,\n",
              "          0.01152635,  0.0453545 ,  0.00875487, -0.04249355,\n",
              "          0.02623383, -0.00830251,  0.04689712, -0.04844732],\n",
              "        [-0.0395951 , -0.00850092, -0.04495001,  0.04774108,\n",
              "         -0.01864243,  0.03377538, -0.01878288,  0.04108479,\n",
              "         -0.03097264,  0.04437483, -0.02057176,  0.04371068],\n",
              "        [ 0.00818767,  0.0161379 ,  0.00114429, -0.03537606,\n",
              "          0.04595914,  0.02060746, -0.03434259,  0.03872677,\n",
              "         -0.01169735,  0.0408063 , -0.0116725 ,  0.01446274],\n",
              "        [ 0.02735185,  0.01735767, -0.034096  , -0.01159757,\n",
              "         -0.02599738,  0.03450202, -0.03427508,  0.02789124,\n",
              "          0.04310719, -0.0003692 , -0.0216522 ,  0.02385452],\n",
              "        [-0.04487137,  0.0297712 ,  0.00572839, -0.04311861,\n",
              "          0.03395119, -0.03924471, -0.02079718,  0.04100347,\n",
              "          0.02254671, -0.02317665,  0.01228994,  0.04094117],\n",
              "        [ 0.03096701, -0.0092387 , -0.00644468,  0.04171463,\n",
              "          0.03103564,  0.04777172, -0.04670488, -0.04088136,\n",
              "          0.01079168, -0.0099802 , -0.04237262, -0.04883695],\n",
              "        [-0.00798281,  0.00167825, -0.04544545,  0.0009286 ,\n",
              "         -0.0068844 ,  0.03416096, -0.02633656,  0.0047374 ,\n",
              "         -0.03104993, -0.04159446, -0.01590217, -0.04526911],\n",
              "        [ 0.04720448,  0.02071989,  0.02709576, -0.0104849 ,\n",
              "          0.02080638,  0.04495866, -0.04587611, -0.00335868,\n",
              "         -0.0272548 ,  0.03617242, -0.0006305 , -0.03822846],\n",
              "        [-0.01019701,  0.01029786,  0.01823778, -0.03126882,\n",
              "          0.01395769, -0.02387906, -0.02165006,  0.03877784,\n",
              "         -0.00389426, -0.04155596, -0.02412431, -0.01842468]],\n",
              "\n",
              "       [[-0.00267386,  0.04497888, -0.01898999,  0.02883989,\n",
              "         -0.00654338, -0.03075428,  0.01430224, -0.03397721,\n",
              "          0.01428051, -0.03040561,  0.00714212,  0.02189949],\n",
              "        [ 0.00583997,  0.00416914,  0.01025517, -0.01858835,\n",
              "          0.01152635,  0.0453545 ,  0.00875487, -0.04249355,\n",
              "          0.02623383, -0.00830251,  0.04689712, -0.04844732],\n",
              "        [-0.0395951 , -0.00850092, -0.04495001,  0.04774108,\n",
              "         -0.01864243,  0.03377538, -0.01878288,  0.04108479,\n",
              "         -0.03097264,  0.04437483, -0.02057176,  0.04371068],\n",
              "        [ 0.00818767,  0.0161379 ,  0.00114429, -0.03537606,\n",
              "          0.04595914,  0.02060746, -0.03434259,  0.03872677,\n",
              "         -0.01169735,  0.0408063 , -0.0116725 ,  0.01446274],\n",
              "        [ 0.02735185,  0.01735767, -0.034096  , -0.01159757,\n",
              "         -0.02599738,  0.03450202, -0.03427508,  0.02789124,\n",
              "          0.04310719, -0.0003692 , -0.0216522 ,  0.02385452],\n",
              "        [-0.04487137,  0.0297712 ,  0.00572839, -0.04311861,\n",
              "          0.03395119, -0.03924471, -0.02079718,  0.04100347,\n",
              "          0.02254671, -0.02317665,  0.01228994,  0.04094117],\n",
              "        [ 0.03096701, -0.0092387 , -0.00644468,  0.04171463,\n",
              "          0.03103564,  0.04777172, -0.04670488, -0.04088136,\n",
              "          0.01079168, -0.0099802 , -0.04237262, -0.04883695],\n",
              "        [-0.00798281,  0.00167825, -0.04544545,  0.0009286 ,\n",
              "         -0.0068844 ,  0.03416096, -0.02633656,  0.0047374 ,\n",
              "         -0.03104993, -0.04159446, -0.01590217, -0.04526911],\n",
              "        [ 0.04720448,  0.02071989,  0.02709576, -0.0104849 ,\n",
              "          0.02080638,  0.04495866, -0.04587611, -0.00335868,\n",
              "         -0.0272548 ,  0.03617242, -0.0006305 , -0.03822846],\n",
              "        [-0.01019701,  0.01029786,  0.01823778, -0.03126882,\n",
              "          0.01395769, -0.02387906, -0.02165006,  0.03877784,\n",
              "         -0.00389426, -0.04155596, -0.02412431, -0.01842468]],\n",
              "\n",
              "       [[-0.00267386,  0.04497888, -0.01898999,  0.02883989,\n",
              "         -0.00654338, -0.03075428,  0.01430224, -0.03397721,\n",
              "          0.01428051, -0.03040561,  0.00714212,  0.02189949],\n",
              "        [ 0.00583997,  0.00416914,  0.01025517, -0.01858835,\n",
              "          0.01152635,  0.0453545 ,  0.00875487, -0.04249355,\n",
              "          0.02623383, -0.00830251,  0.04689712, -0.04844732],\n",
              "        [-0.0395951 , -0.00850092, -0.04495001,  0.04774108,\n",
              "         -0.01864243,  0.03377538, -0.01878288,  0.04108479,\n",
              "         -0.03097264,  0.04437483, -0.02057176,  0.04371068],\n",
              "        [ 0.00818767,  0.0161379 ,  0.00114429, -0.03537606,\n",
              "          0.04595914,  0.02060746, -0.03434259,  0.03872677,\n",
              "         -0.01169735,  0.0408063 , -0.0116725 ,  0.01446274],\n",
              "        [ 0.02735185,  0.01735767, -0.034096  , -0.01159757,\n",
              "         -0.02599738,  0.03450202, -0.03427508,  0.02789124,\n",
              "          0.04310719, -0.0003692 , -0.0216522 ,  0.02385452],\n",
              "        [-0.04487137,  0.0297712 ,  0.00572839, -0.04311861,\n",
              "          0.03395119, -0.03924471, -0.02079718,  0.04100347,\n",
              "          0.02254671, -0.02317665,  0.01228994,  0.04094117],\n",
              "        [ 0.03096701, -0.0092387 , -0.00644468,  0.04171463,\n",
              "          0.03103564,  0.04777172, -0.04670488, -0.04088136,\n",
              "          0.01079168, -0.0099802 , -0.04237262, -0.04883695],\n",
              "        [-0.00798281,  0.00167825, -0.04544545,  0.0009286 ,\n",
              "         -0.0068844 ,  0.03416096, -0.02633656,  0.0047374 ,\n",
              "         -0.03104993, -0.04159446, -0.01590217, -0.04526911],\n",
              "        [ 0.04720448,  0.02071989,  0.02709576, -0.0104849 ,\n",
              "          0.02080638,  0.04495866, -0.04587611, -0.00335868,\n",
              "         -0.0272548 ,  0.03617242, -0.0006305 , -0.03822846],\n",
              "        [-0.01019701,  0.01029786,  0.01823778, -0.03126882,\n",
              "          0.01395769, -0.02387906, -0.02165006,  0.03877784,\n",
              "         -0.00389426, -0.04155596, -0.02412431, -0.01842468]]],\n",
              "      dtype=float32)>"
            ]
          },
          "metadata": {},
          "execution_count": 41
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Lets test our Encoder"
      ],
      "metadata": {
        "id": "sdOUw61cfvUj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "input_batch = [\n",
        "    \"Where can I find a pizzeria?\",\n",
        "    \"Mass hysteria over listeria.\",\n",
        "    \"I ain't no circle back girl.\"\n",
        "]\n",
        "\n",
        "bpemb_en.encode(input_batch)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iInW8KQFfgGu",
        "outputId": "597b2502-e377-4872-ffb3-505a49c4636d"
      },
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[['▁where', '▁can', '▁i', '▁find', '▁a', '▁p', 'iz', 'zer', 'ia', '?'],\n",
              " ['▁mass', '▁hy', 'ster', 'ia', '▁over', '▁l', 'ister', 'ia', '.'],\n",
              " ['▁i', '▁a', 'in', \"'\", 't', '▁no', '▁circle', '▁back', '▁girl', '.']]"
            ]
          },
          "metadata": {},
          "execution_count": 42
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "input_seqs = bpemb_en.encode_ids(input_batch)\n",
        "print(\"Vectorized inputs:\")\n",
        "input_seqs"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Be-ODuzkfijZ",
        "outputId": "01ded8ab-970f-4b27-bfe3-9cbac84d43d5"
      },
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Vectorized inputs:\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[[571, 280, 386, 1934, 4, 24, 248, 4339, 177, 9967],\n",
              " [1535, 1354, 1238, 177, 380, 43, 871, 177, 9935],\n",
              " [386, 4, 6, 9937, 9915, 467, 5410, 810, 3692, 9935]]"
            ]
          },
          "metadata": {},
          "execution_count": 43
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "padded_input_seqs = tf.keras.preprocessing.sequence.pad_sequences(input_seqs, padding=\"post\")\n",
        "print(\"Input to the encoder:\")\n",
        "print(padded_input_seqs.shape)\n",
        "print(padded_input_seqs)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3szzwzhofj9L",
        "outputId": "84405303-6750-4600-9d5e-38cc1c65c885"
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input to the encoder:\n",
            "(3, 10)\n",
            "[[ 571  280  386 1934    4   24  248 4339  177 9967]\n",
            " [1535 1354 1238  177  380   43  871  177 9935    0]\n",
            " [ 386    4    6 9937 9915  467 5410  810 3692 9935]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "enc_mask = tf.cast(tf.math.not_equal(padded_input_seqs, 0), tf.float32)\n",
        "print(\"Input:\")\n",
        "print(padded_input_seqs, '\\n')\n",
        "print(\"Encoder mask:\")\n",
        "print(enc_mask)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dn-nknGDflWH",
        "outputId": "76a1350c-b8c1-45dc-ac05-f0bbf20274d0"
      },
      "execution_count": 45,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input:\n",
            "[[ 571  280  386 1934    4   24  248 4339  177 9967]\n",
            " [1535 1354 1238  177  380   43  871  177 9935    0]\n",
            " [ 386    4    6 9937 9915  467 5410  810 3692 9935]] \n",
            "\n",
            "Encoder mask:\n",
            "tf.Tensor(\n",
            "[[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]\n",
            " [1. 1. 1. 1. 1. 1. 1. 1. 1. 0.]\n",
            " [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]], shape=(3, 10), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "enc_mask = enc_mask[:, tf.newaxis, tf.newaxis, :]\n",
        "enc_mask"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tpJHsZhRfnzn",
        "outputId": "2842251d-2106-4711-9d07-435ae2d3cba7"
      },
      "execution_count": 46,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: shape=(3, 1, 1, 10), dtype=float32, numpy=\n",
              "array([[[[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]]],\n",
              "\n",
              "\n",
              "       [[[1., 1., 1., 1., 1., 1., 1., 1., 1., 0.]]],\n",
              "\n",
              "\n",
              "       [[[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]]]], dtype=float32)>"
            ]
          },
          "metadata": {},
          "execution_count": 46
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "num_encoder_blocks = 6\n",
        "\n",
        "# d_model is the embedding dimension used throughout.\n",
        "d_model = 12\n",
        "\n",
        "num_heads = 3\n",
        "\n",
        "# Feed-forward network hidden dimension width.\n",
        "ffn_hidden_dim = 48\n",
        "\n",
        "src_vocab_size = bpemb_vocab_size\n",
        "max_input_seq_len = padded_input_seqs.shape[1]\n",
        "\n",
        "encoder = Encoder(\n",
        "    num_encoder_blocks,\n",
        "    d_model,\n",
        "    num_heads,\n",
        "    ffn_hidden_dim,\n",
        "    src_vocab_size,\n",
        "    max_input_seq_len)"
      ],
      "metadata": {
        "id": "5SmqguXHfpmR"
      },
      "execution_count": 47,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "encoder_output, attn_weights = encoder(padded_input_seqs, training=True, \n",
        "                                       mask=enc_mask)\n",
        "print(f\"Encoder output {encoder_output.shape}:\")\n",
        "print(encoder_output)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Z7cKWhtQfq8h",
        "outputId": "fad7e4b5-27a9-4cc8-95ee-828e5d9bc69d"
      },
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Encoder output (3, 10, 12):\n",
            "tf.Tensor(\n",
            "[[[ 5.15816629e-01  7.61540473e-01 -5.77319562e-01 -1.17828798e+00\n",
            "   -7.44615853e-01  1.65449226e+00 -9.72729862e-01  2.48307199e-03\n",
            "   -3.40688735e-01  1.36467981e+00  1.00068784e+00 -1.48605800e+00]\n",
            "  [ 3.57278019e-01  5.87776482e-01 -7.65292645e-01 -1.69315350e+00\n",
            "   -9.02665973e-01  1.44819975e+00 -8.42213154e-01  1.35431662e-01\n",
            "    3.02950680e-01  1.25711429e+00  1.22967255e+00 -1.11509860e+00]\n",
            "  [-9.24525186e-02  3.44464451e-01 -8.33958745e-01 -1.31060040e+00\n",
            "   -3.82650867e-02  9.61784005e-01 -7.66187847e-01  4.31145787e-01\n",
            "   -1.68807566e-01  1.72790456e+00  1.40664506e+00 -1.66167176e+00]\n",
            "  [ 5.31727314e-01  7.21970201e-01 -4.14329678e-01 -1.24973392e+00\n",
            "   -1.39198756e+00  9.73598361e-01 -1.07928050e+00  3.52839112e-01\n",
            "    1.39878631e-01  1.25907505e+00  1.43182182e+00 -1.27557898e+00]\n",
            "  [ 3.78314406e-01  5.09653270e-01 -4.96241033e-01 -1.19794273e+00\n",
            "   -9.56186414e-01  2.11426353e+00 -9.38754022e-01  4.33031589e-01\n",
            "   -3.46625894e-01  8.95030141e-01  9.11093533e-01 -1.30563617e+00]\n",
            "  [ 3.74054730e-01  4.62736815e-01 -3.65884870e-01 -8.96167099e-01\n",
            "   -5.76828957e-01  1.98536205e+00 -1.33179808e+00  7.19928384e-01\n",
            "   -4.86694276e-01  9.81458247e-01  6.98537171e-01 -1.56470442e+00]\n",
            "  [ 3.59841764e-01 -2.26230592e-01 -4.72908080e-01 -1.03639507e+00\n",
            "   -4.39433604e-01  1.53143454e+00 -8.93973887e-01  8.48780692e-01\n",
            "   -9.47102681e-02  1.51274705e+00  7.90965736e-01 -1.88011813e+00]\n",
            "  [-1.08151376e-01  5.50602257e-01 -9.54607725e-01 -1.31323433e+00\n",
            "   -1.08419788e+00  1.13881767e+00 -4.59152609e-01  6.35621071e-01\n",
            "    5.30900776e-01  1.38718188e+00  1.19784975e+00 -1.52162957e+00]\n",
            "  [ 3.56226832e-01  1.91129133e-01 -8.76889944e-01 -1.08539999e+00\n",
            "   -9.45012808e-01  1.72985375e+00 -5.92675269e-01  4.54504013e-01\n",
            "   -9.32951346e-02  1.28036845e+00  1.15154374e+00 -1.57035303e+00]\n",
            "  [ 5.03177166e-01 -1.12271480e-01 -6.32555187e-01 -1.28470004e+00\n",
            "   -1.71926355e+00  1.01490426e+00 -2.52934963e-01  6.54459834e-01\n",
            "    5.33813477e-01  1.63247526e+00  8.78486753e-01 -1.21559179e+00]]\n",
            "\n",
            " [[ 1.09361792e+00 -2.05463450e-03  4.53847468e-01 -1.05023408e+00\n",
            "   -1.65246829e-01  2.25297880e+00 -1.27722120e+00  4.80630368e-01\n",
            "   -6.63040817e-01  6.36511520e-02  2.35166818e-01 -1.42209470e+00]\n",
            "  [ 1.42339155e-01 -5.51407158e-01 -6.53665483e-01 -1.28806341e+00\n",
            "   -5.33288062e-01  9.31522846e-01 -2.27918059e-01  6.85249805e-01\n",
            "    5.49648285e-01  1.90958416e+00  8.39977682e-01 -1.80397975e+00]\n",
            "  [ 3.57926607e-01  1.07610889e-01 -2.65526682e-01 -5.76402783e-01\n",
            "   -4.87648040e-01  1.84455860e+00 -8.23633015e-01  2.82331079e-01\n",
            "   -2.47583613e-01  1.59801257e+00  2.96172351e-01 -2.08581805e+00]\n",
            "  [ 1.51779335e-02  5.34209728e-01 -6.00911379e-01 -1.03433096e+00\n",
            "   -1.00187635e+00  1.25276017e+00 -7.13451982e-01  5.64028732e-02\n",
            "    2.52887666e-01  1.58154595e+00  1.34002113e+00 -1.68243444e+00]\n",
            "  [-2.13713557e-01  5.28257370e-01 -5.75346529e-01 -9.76852715e-01\n",
            "   -9.36101615e-01  2.12735891e+00 -6.13257468e-01  4.75200146e-01\n",
            "   -1.44138277e-01  1.44210625e+00  3.54290426e-01 -1.46780312e+00]\n",
            "  [ 6.92993641e-01  2.92699248e-01 -4.43218678e-01 -8.02154124e-01\n",
            "   -3.58005285e-01  2.13156128e+00 -1.09615088e+00  1.72479853e-01\n",
            "   -5.36251664e-01  1.19435847e+00  4.22568440e-01 -1.67088008e+00]\n",
            "  [-3.69189054e-01  1.21255612e+00 -6.42224610e-01 -1.43472910e+00\n",
            "   -9.90074813e-01  8.31020892e-01 -5.96948266e-01  6.06930077e-01\n",
            "    5.91083229e-01  1.24536133e+00  1.09937835e+00 -1.55316412e+00]\n",
            "  [ 5.26515663e-01  7.94657886e-01 -9.86919522e-01 -1.46026051e+00\n",
            "   -1.20361876e+00  9.02531445e-01  2.93904722e-01  3.40915948e-01\n",
            "    3.28255832e-01  1.57547557e+00  5.17955542e-01 -1.62941360e+00]\n",
            "  [ 1.43138230e+00  8.62317264e-01 -3.99584472e-02 -7.92880058e-01\n",
            "   -1.12130165e+00  1.23324835e+00 -8.01862299e-01 -6.34642184e-01\n",
            "   -3.82830352e-01  1.35534990e+00  4.82175112e-01 -1.59099770e+00]\n",
            "  [-6.11292601e-01  3.94999295e-01 -6.47897601e-01 -8.77423584e-01\n",
            "   -1.31916106e+00  1.10528779e+00  1.47389984e-02  6.96344554e-01\n",
            "    9.62121189e-01  1.69550717e+00  3.00416052e-01 -1.71364045e+00]]\n",
            "\n",
            " [[-4.54108566e-02  1.84979308e+00  1.92561328e-01 -3.76882493e-01\n",
            "   -6.69534028e-01 -3.04782629e-01 -2.03406048e+00  8.57195914e-01\n",
            "    3.13062876e-01  1.08045980e-01  1.25090778e+00 -1.14089632e+00]\n",
            "  [-3.55992258e-01  1.53063893e+00  1.30234802e+00 -3.61703499e-03\n",
            "   -1.35828227e-01 -9.80784833e-01 -2.21486425e+00 -6.23804927e-02\n",
            "    7.32109487e-01 -4.19766963e-01  1.02487624e+00 -4.16738421e-01]\n",
            "  [-5.52536547e-01  1.43290246e+00  1.04551196e+00 -3.59130174e-01\n",
            "    5.85717022e-01 -1.67257917e+00 -1.81453669e+00  1.36227119e+00\n",
            "    4.57653850e-02 -3.00675005e-01  1.48920894e-01  7.83685744e-02]\n",
            "  [ 1.09398298e-01  1.46964276e+00  2.17564180e-01 -5.65911174e-01\n",
            "    1.21644586e-01 -1.11298060e+00 -2.09345150e+00  6.28284037e-01\n",
            "    5.72520971e-01 -3.10687542e-01  1.60693359e+00 -6.42957389e-01]\n",
            "  [-4.97186696e-03  7.12431371e-01 -1.00300699e-01 -4.65544432e-01\n",
            "   -4.04693484e-01 -8.32320750e-01 -2.17943192e+00  1.29228115e+00\n",
            "    8.38950932e-01  2.90186852e-01  1.61570501e+00 -7.62292445e-01]\n",
            "  [ 1.42428660e+00  1.92555225e+00 -2.32462317e-01  6.83987141e-02\n",
            "    2.51298528e-02 -3.16130638e-01 -1.54352808e+00  4.88125026e-01\n",
            "   -4.51251984e-01 -1.83578059e-01  5.18310785e-01 -1.72285235e+00]\n",
            "  [ 6.48983598e-01  7.85135210e-01 -4.81863886e-01 -4.54810560e-01\n",
            "   -9.49205756e-02 -8.05930793e-01 -1.75606143e+00  1.21190262e+00\n",
            "    1.27952182e+00 -1.22286588e-01  1.24225831e+00 -1.45192778e+00]\n",
            "  [ 1.48662671e-01  8.75391424e-01 -1.28298080e+00 -1.05999064e+00\n",
            "   -3.16148132e-01 -2.40330189e-01 -7.45933235e-01  1.39286685e+00\n",
            "    1.51205325e+00  1.28485588e-02  1.16493785e+00 -1.46137714e+00]\n",
            "  [-2.67122179e-01  1.50357747e+00 -1.24591775e-01 -6.92552507e-01\n",
            "   -5.78572452e-01 -3.02846730e-01 -1.70091033e+00  1.66002297e+00\n",
            "    1.75439104e-01  9.55013633e-02  1.34503078e+00 -1.11297572e+00]\n",
            "  [ 8.02782536e-01  1.93554020e+00  2.69401431e-01 -2.02535093e-01\n",
            "   -9.34683383e-01 -4.92501557e-01 -1.93541086e+00  9.50577080e-01\n",
            "    1.59799993e-01  2.76657809e-02  5.82923234e-01 -1.16355932e+00]]], shape=(3, 10, 12), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##DECODER\n",
        "1. a **Multi-Head Cross-Attention** layer which uses the encoder's outputs as the keys and values.\n",
        "\n",
        "2. an extra skip/residual connection along with an extra layer normalization step.\n",
        "\n",
        "<div>\n",
        "<img src=\"https://drive.google.com/uc?export=view&id=1WVT4SX49bnta4uscOTF4xrsxFI4PbPER\" width=\"500\"/>\n",
        "</div>"
      ],
      "metadata": {
        "id": "YA-fJMpEfs25"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class DecoderBlock(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads, hidden_dim, dropout_rate=0.1):\n",
        "    super(DecoderBlock, self).__init__()\n",
        "\n",
        "    self.mhsa1 = MultiHeadSelfAttention(d_model, num_heads)\n",
        "    self.mhsa2 = MultiHeadSelfAttention(d_model, num_heads)\n",
        "\n",
        "    self.ffn = feed_forward_network(d_model, hidden_dim)\n",
        "\n",
        "    self.dropout1 = tf.keras.layers.Dropout(dropout_rate)\n",
        "    self.dropout2 = tf.keras.layers.Dropout(dropout_rate)\n",
        "    self.dropout3 = tf.keras.layers.Dropout(dropout_rate)\n",
        "\n",
        "    self.layernorm1 = tf.keras.layers.LayerNormalization()\n",
        "    self.layernorm2 = tf.keras.layers.LayerNormalization()\n",
        "    self.layernorm3 = tf.keras.layers.LayerNormalization()\n",
        "    \n",
        "  # Note the decoder block takes two masks. One for the first MHSA(Multi-Head Skip-Attention), another\n",
        "  # for the second MHSA.\n",
        "  def call(self, encoder_output, target, training, decoder_mask, memory_mask):\n",
        "    mhsa_output1, attn_weights = self.mhsa1(target, target, target, decoder_mask)\n",
        "    mhsa_output1 = self.dropout1(mhsa_output1, training=training)\n",
        "    mhsa_output1 = self.layernorm1(mhsa_output1 + target)\n",
        "\n",
        "    mhsa_output2, attn_weights = self.mhsa2(mhsa_output1, encoder_output, \n",
        "                                            encoder_output, \n",
        "                                            memory_mask)\n",
        "    mhsa_output2 = self.dropout2(mhsa_output2, training=training)\n",
        "    mhsa_output2 = self.layernorm2(mhsa_output2 + mhsa_output1)\n",
        "\n",
        "    ffn_output = self.ffn(mhsa_output2)\n",
        "    ffn_output = self.dropout3(ffn_output, training=training)\n",
        "    output = self.layernorm3(ffn_output + mhsa_output2)\n",
        "\n",
        "    return output, attn_weights\n"
      ],
      "metadata": {
        "id": "NdNRcNpcfsNV"
      },
      "execution_count": 49,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The decoder is almost the same as the encoder except it takes the encoder's output as part of its input, and it takes two masks: the decoder mask and memory mask."
      ],
      "metadata": {
        "id": "zTKO2620iux6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Decoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, num_blocks, d_model, num_heads, hidden_dim, target_vocab_size,\n",
        "               max_seq_len, dropout_rate=0.1):\n",
        "    super(Decoder, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.max_seq_len = max_seq_len\n",
        "\n",
        "    self.token_embed = tf.keras.layers.Embedding(target_vocab_size, self.d_model)\n",
        "    self.pos_embed = tf.keras.layers.Embedding(max_seq_len, self.d_model)\n",
        "\n",
        "    self.dropout = tf.keras.layers.Dropout(dropout_rate)\n",
        "\n",
        "    self.blocks = [DecoderBlock(self.d_model, num_heads, hidden_dim, dropout_rate) for _ in range(num_blocks)]\n",
        "\n",
        "  def call(self, encoder_output, target, training, decoder_mask, memory_mask):\n",
        "    token_embeds = self.token_embed(target)\n",
        "\n",
        "    # Generate position indices.\n",
        "    num_pos = target.shape[0] * self.max_seq_len\n",
        "    pos_idx = np.resize(np.arange(self.max_seq_len), num_pos)\n",
        "    pos_idx = np.reshape(pos_idx, target.shape)\n",
        "\n",
        "    pos_embeds = self.pos_embed(pos_idx)\n",
        "\n",
        "    x = self.dropout(token_embeds + pos_embeds, training=training)\n",
        "\n",
        "    for block in self.blocks:\n",
        "      x, weights = block(encoder_output, x, training, decoder_mask, memory_mask)\n",
        "\n",
        "    return x, weights"
      ],
      "metadata": {
        "id": "teDJm1P6hg0A"
      },
      "execution_count": 50,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Dummy values\n",
        "target_input_seqs = [\n",
        "    [1, 652, 723, 123, 62],\n",
        "    [1, 25,  98, 129, 248, 215, 359, 249],\n",
        "    [1, 2369, 1259, 125, 486],\n",
        "]"
      ],
      "metadata": {
        "id": "9O57Xds3hlcq"
      },
      "execution_count": 54,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As we did with the encoder input sequences, we need to pad out this batch so that all sequences within it are the same length."
      ],
      "metadata": {
        "id": "vNx7RWDjjP7N"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "padded_target_input_seqs = tf.keras.preprocessing.sequence.pad_sequences(target_input_seqs, padding=\"post\")\n",
        "print(\"Padded target inputs to the decoder:\")\n",
        "print(padded_target_input_seqs.shape)\n",
        "print(padded_target_input_seqs)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gIFAoOdOhlfT",
        "outputId": "381b0e1c-8466-4166-940a-abd44e97a237"
      },
      "execution_count": 55,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Padded target inputs to the decoder:\n",
            "(3, 8)\n",
            "[[   1  652  723  123   62    0    0    0]\n",
            " [   1   25   98  129  248  215  359  249]\n",
            " [   1 2369 1259  125  486    0    0    0]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can create the padding mask the same way we did for the encoder."
      ],
      "metadata": {
        "id": "ijHV-8mAjS-f"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "dec_padding_mask = tf.cast(tf.math.not_equal(padded_target_input_seqs, 0), tf.float32)\n",
        "dec_padding_mask = dec_padding_mask[:, tf.newaxis, tf.newaxis, :]\n",
        "print(dec_padding_mask)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c9x8gUg-hlh2",
        "outputId": "0e91abd7-0890-491d-b407-1003c29bc0bb"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tf.Tensor(\n",
            "[[[[1. 1. 1. 1. 1. 0. 0. 0.]]]\n",
            "\n",
            "\n",
            " [[[1. 1. 1. 1. 1. 1. 1. 1.]]]\n",
            "\n",
            "\n",
            " [[[1. 1. 1. 1. 1. 0. 0. 0.]]]], shape=(3, 1, 1, 8), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "As we covered in the slides, the look-ahead mask is a diagonal where the lower half are 1s and the upper half are zeros. This is easy to create using the *band_part* method:<br>\n",
        "https://www.tensorflow.org/api_docs/python/tf/linalg/band_part"
      ],
      "metadata": {
        "id": "lUIVXV_ljZ7G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "target_input_seq_len = padded_target_input_seqs.shape[1]\n",
        "look_ahead_mask = tf.linalg.band_part(tf.ones((target_input_seq_len, \n",
        "                                               target_input_seq_len)), -1, 0)\n",
        "print(look_ahead_mask)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OVorEyHOhlkr",
        "outputId": "e6b8d2d5-1d2d-42cf-8993-e4a1cd7f9197"
      },
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tf.Tensor(\n",
            "[[1. 0. 0. 0. 0. 0. 0. 0.]\n",
            " [1. 1. 0. 0. 0. 0. 0. 0.]\n",
            " [1. 1. 1. 0. 0. 0. 0. 0.]\n",
            " [1. 1. 1. 1. 0. 0. 0. 0.]\n",
            " [1. 1. 1. 1. 1. 0. 0. 0.]\n",
            " [1. 1. 1. 1. 1. 1. 0. 0.]\n",
            " [1. 1. 1. 1. 1. 1. 1. 0.]\n",
            " [1. 1. 1. 1. 1. 1. 1. 1.]], shape=(8, 8), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "To create the decoder mask, we just need to combine the padding and look-ahead masks. Note how the columns of the resulting decoder mask are all zero for padding positions."
      ],
      "metadata": {
        "id": "JhcbMgEijdSn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "dec_mask = tf.minimum(dec_padding_mask, look_ahead_mask)\n",
        "print(\"The decoder mask:\")\n",
        "print(dec_mask)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rPTUoIh_hlnY",
        "outputId": "0076ec84-7b15-4fb1-fab6-298dff688553"
      },
      "execution_count": 58,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The decoder mask:\n",
            "tf.Tensor(\n",
            "[[[[1. 0. 0. 0. 0. 0. 0. 0.]\n",
            "   [1. 1. 0. 0. 0. 0. 0. 0.]\n",
            "   [1. 1. 1. 0. 0. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 0. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 0. 0. 0.]]]\n",
            "\n",
            "\n",
            " [[[1. 0. 0. 0. 0. 0. 0. 0.]\n",
            "   [1. 1. 0. 0. 0. 0. 0. 0.]\n",
            "   [1. 1. 1. 0. 0. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 0. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 1. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 1. 1. 0.]\n",
            "   [1. 1. 1. 1. 1. 1. 1. 1.]]]\n",
            "\n",
            "\n",
            " [[[1. 0. 0. 0. 0. 0. 0. 0.]\n",
            "   [1. 1. 0. 0. 0. 0. 0. 0.]\n",
            "   [1. 1. 1. 0. 0. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 0. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 0. 0. 0.]\n",
            "   [1. 1. 1. 1. 1. 0. 0. 0.]]]], shape=(3, 1, 8, 8), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can now declare a decoder and pass it everything it needs. In our case, the *memory* mask is the same as the *encoder* mask."
      ],
      "metadata": {
        "id": "NAk4sJdnjgiU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "decoder = Decoder(6, 12, 3, 48, 10000, 8)\n",
        "decoder_output, _ = decoder(encoder_output, padded_target_input_seqs, \n",
        "                            True, dec_mask, enc_mask)\n",
        "print(f\"Decoder output {decoder_output.shape}:\")\n",
        "print(decoder_output)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UC04AT4mhlpx",
        "outputId": "133554be-35eb-410f-a348-3be1826dfdde"
      },
      "execution_count": 59,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Decoder output (3, 8, 12):\n",
            "tf.Tensor(\n",
            "[[[-1.3903046e-01 -1.1472348e+00 -9.4525301e-01 -9.9262261e-01\n",
            "    9.6591866e-01  7.7667969e-01  7.3278648e-01 -1.5465268e+00\n",
            "    1.8527468e+00  8.8189763e-01 -2.6752219e-01 -1.7183934e-01]\n",
            "  [-1.3411474e-01 -9.7858495e-01 -8.6579299e-01 -1.0426285e+00\n",
            "    9.2134255e-01  1.2728901e+00  3.7006342e-01 -1.4032898e+00\n",
            "    2.0083435e+00  4.4674009e-01 -6.0845405e-01  1.3485389e-02]\n",
            "  [ 6.5341175e-01 -8.9217401e-01 -6.6149551e-01 -9.4876671e-01\n",
            "    6.9195396e-01  1.4535443e+00 -8.3237670e-02 -1.5115743e+00\n",
            "    1.8314381e+00  6.5290809e-01 -4.1486269e-01 -7.7114522e-01]\n",
            "  [-1.1089486e-01 -1.0670661e+00 -9.7008365e-01 -6.3454688e-01\n",
            "    7.4080920e-01  1.3388698e+00  6.3307554e-01 -1.3387581e+00\n",
            "    1.8285021e+00  8.3866543e-01 -3.8097706e-01 -8.7759578e-01]\n",
            "  [ 4.8765421e-01 -1.3967913e+00 -1.1272376e+00 -6.4497453e-01\n",
            "    6.0428280e-01  1.7554371e+00 -3.9304355e-01 -4.1379380e-01\n",
            "    1.6058282e+00  8.9691055e-01 -5.9938884e-01 -7.7488327e-01]\n",
            "  [ 5.2802533e-01 -9.0742582e-01 -1.1197897e+00 -1.1759536e+00\n",
            "    6.5396935e-01  8.9867163e-01  9.7464585e-01 -1.3486465e+00\n",
            "    1.8556057e+00  5.5319852e-01 -6.5896797e-01 -2.5333294e-01]\n",
            "  [-5.0053068e-02 -1.3131275e+00 -9.9020070e-01 -3.3731627e-01\n",
            "    5.9454900e-01  1.4855628e+00  5.8544481e-01 -1.4843473e+00\n",
            "    1.6835660e+00  8.2878894e-01 -3.6218232e-01 -6.4068466e-01]\n",
            "  [ 9.6693084e-02 -1.0139909e+00 -1.0110612e+00 -3.5038200e-01\n",
            "    3.4956723e-01  1.1333768e+00  9.2544425e-01 -1.5041984e+00\n",
            "    1.6778642e+00  1.1651890e+00 -6.1107987e-01 -8.5742253e-01]]\n",
            "\n",
            " [[-7.5796169e-01 -9.3521649e-01 -7.8262597e-01 -6.6394544e-01\n",
            "   -2.7162069e-01  1.8863254e+00  1.0651569e+00 -6.0208380e-01\n",
            "    1.6013741e+00  7.1398574e-01 -1.2765259e-01 -1.1257356e+00]\n",
            "  [-4.1633919e-01 -5.8914447e-01 -5.9481609e-01 -4.4020981e-01\n",
            "    4.4318303e-01  1.4185096e+00  6.5690756e-01 -1.8911653e+00\n",
            "    1.9521792e+00  5.2859908e-01 -3.8355693e-01 -6.8414700e-01]\n",
            "  [-3.0454400e-01 -6.5575939e-01 -1.0982703e+00 -7.4728616e-02\n",
            "    2.5427601e-01  1.4978842e+00  7.0195788e-01 -1.6493645e+00\n",
            "    1.8945208e+00  6.4357406e-01 -5.9981734e-01 -6.0972857e-01]\n",
            "  [-1.2957951e-01 -9.4930261e-01 -1.1160412e+00 -3.3425453e-01\n",
            "    6.6838551e-01  1.4404687e+00 -2.6637599e-01 -1.5514085e+00\n",
            "    2.0302558e+00  6.9133282e-01 -2.4869370e-01 -2.3478672e-01]\n",
            "  [-1.1266639e-01 -9.4637811e-01 -1.1557088e+00 -1.8994085e-01\n",
            "    2.9376164e-01  1.5206290e+00  1.2926977e+00 -1.2533451e+00\n",
            "    9.6820974e-01  1.1993514e+00 -3.6935633e-01 -1.2472539e+00]\n",
            "  [-2.5229353e-01 -1.1049886e+00 -2.3400262e-01 -5.8590591e-01\n",
            "    8.9764142e-01  1.1541018e+00  7.3276520e-01 -1.7735856e+00\n",
            "    1.8064053e+00  5.9852904e-01 -3.9558366e-01 -8.4308290e-01]\n",
            "  [-2.0778522e-01 -7.8208768e-01 -9.8682600e-01 -6.4997423e-01\n",
            "    4.7618002e-01  1.3913099e+00  9.5980245e-01 -1.6348306e+00\n",
            "    1.7840071e+00  6.8317866e-01 -5.0893855e-01 -5.2403557e-01]\n",
            "  [-1.9871655e-01 -1.4587231e-01 -9.0782696e-01 -5.6127363e-01\n",
            "    3.0599374e-01  1.6869731e+00  6.4546287e-01 -1.7212809e+00\n",
            "    1.8089339e+00  4.7736773e-01 -7.5156784e-01 -6.3819313e-01]]\n",
            "\n",
            " [[ 5.8169377e-01  1.1465380e+00 -1.5342370e+00 -1.5896785e+00\n",
            "    1.3850477e+00 -1.1313349e+00  3.1082696e-01  4.1418335e-01\n",
            "    8.0972433e-02 -8.2122225e-01  1.1451843e+00  1.2026289e-02]\n",
            "  [-4.2612270e-02  6.4934289e-01 -1.5280372e+00 -9.2143482e-01\n",
            "    1.5231307e+00 -2.0014284e+00  5.9747881e-01  9.2906421e-01\n",
            "    4.0758613e-02  4.1500875e-04  9.0310353e-01 -1.4978114e-01]\n",
            "  [-3.3313850e-01  7.8373748e-01 -1.4596215e+00 -1.2543616e+00\n",
            "    1.7145994e+00 -1.6572342e+00  3.7716079e-01  1.1062441e+00\n",
            "   -2.4150939e-01  1.5510558e-01  6.2961316e-01  1.7940472e-01]\n",
            "  [-3.4624693e-01  1.0208209e+00 -1.2596997e+00 -7.6649022e-01\n",
            "    1.4347005e+00 -1.8724710e+00  5.7572925e-01  6.7549616e-01\n",
            "   -7.6131034e-01 -2.5403476e-01  1.2525532e+00  3.0095297e-01]\n",
            "  [-5.4379267e-01  5.8702755e-01 -1.3142017e+00 -1.3920190e+00\n",
            "    1.4674563e+00 -1.5525002e+00  6.0573244e-01  1.2218111e+00\n",
            "   -3.3602220e-01 -4.0796455e-02  1.0414540e+00  2.5585136e-01]\n",
            "  [ 2.3518236e-02  1.0470022e+00 -1.6452729e+00 -1.4938921e+00\n",
            "    1.7300802e+00 -1.1922512e+00  4.4831535e-01  9.8330683e-01\n",
            "   -4.9220359e-01  1.2159913e-01  1.8499440e-01  2.8480336e-01]\n",
            "  [ 1.4433057e-02  8.8069642e-01 -1.5166434e+00 -1.2483981e+00\n",
            "    1.9378603e+00 -1.5452719e+00  8.6961609e-01  4.8320150e-01\n",
            "   -3.3132932e-01 -7.9195770e-03  1.7420870e-01  2.8954601e-01]\n",
            "  [-1.0821687e-01  1.4309313e+00 -1.4018582e+00 -1.6735029e+00\n",
            "    6.3635391e-01 -7.5449747e-01  8.3758354e-01  1.0017796e+00\n",
            "   -9.9950790e-01 -3.9770207e-01  1.1087228e+00  3.1991416e-01]]], shape=(3, 8, 12), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#TRANSFORMER"
      ],
      "metadata": {
        "id": "YANny9BJh6RB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Transformer(tf.keras.Model):\n",
        "  def __init__(self, num_blocks, d_model, num_heads, hidden_dim, source_vocab_size,\n",
        "               target_vocab_size, max_input_len, max_target_len, dropout_rate=0.1):\n",
        "    super(Transformer, self).__init__()\n",
        "\n",
        "    self.encoder = Encoder(num_blocks, d_model, num_heads, hidden_dim, source_vocab_size, \n",
        "                           max_input_len, dropout_rate)\n",
        "    \n",
        "    self.decoder = Decoder(num_blocks, d_model, num_heads, hidden_dim, target_vocab_size,\n",
        "                           max_target_len, dropout_rate)\n",
        "    \n",
        "    # The final dense layer to generate logits from the decoder output.\n",
        "    self.output_layer = tf.keras.layers.Dense(target_vocab_size)\n",
        "\n",
        "  def call(self, input_seqs, target_input_seqs, training, encoder_mask,\n",
        "           decoder_mask, memory_mask):\n",
        "    encoder_output, encoder_attn_weights = self.encoder(input_seqs, \n",
        "                                                        training, encoder_mask)\n",
        "\n",
        "    decoder_output, decoder_attn_weights = self.decoder(encoder_output, \n",
        "                                                        target_input_seqs, training,\n",
        "                                                        decoder_mask, memory_mask)\n",
        "\n",
        "    return self.output_layer(decoder_output), encoder_attn_weights, decoder_attn_weights\n"
      ],
      "metadata": {
        "id": "cqjhnvj5h3_x"
      },
      "execution_count": 60,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transformer = Transformer(\n",
        "    num_blocks = 6,\n",
        "    d_model = 12,\n",
        "    num_heads = 3,\n",
        "    hidden_dim = 48,\n",
        "    source_vocab_size = bpemb_vocab_size,\n",
        "    target_vocab_size = 7000, # made-up target vocab size.\n",
        "    max_input_len = padded_input_seqs.shape[1],\n",
        "    max_target_len = padded_target_input_seqs.shape[1])\n",
        "\n",
        "transformer_output, _, _ = transformer(padded_input_seqs, \n",
        "                                       padded_target_input_seqs, True, \n",
        "                                       enc_mask, dec_mask, memory_mask=enc_mask)\n",
        "print(f\"Transformer output {transformer_output.shape}:\")\n",
        "print(transformer_output) # If training, we would use this output to calculate losses."
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YIvpPvr9h_cW",
        "outputId": "7678f6ac-6ed1-4e15-f371-f65043d6dc04"
      },
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Transformer output (3, 8, 7000):\n",
            "tf.Tensor(\n",
            "[[[-0.01586961 -0.01493736  0.10750537 ...  0.00607951  0.08848661\n",
            "    0.015226  ]\n",
            "  [-0.04400273 -0.03207124  0.13835101 ... -0.0017617   0.07536489\n",
            "    0.02428343]\n",
            "  [-0.03911787 -0.03066409  0.12442512 ...  0.00308029  0.0952151\n",
            "    0.01540167]\n",
            "  ...\n",
            "  [-0.0881561   0.03407094  0.08948196 ... -0.01739982  0.05045627\n",
            "    0.02708782]\n",
            "  [-0.0687721   0.03003042  0.08036731 ...  0.00281868  0.08661909\n",
            "    0.02208049]\n",
            "  [-0.04451494 -0.00923184  0.10660914 ...  0.01369279  0.09513992\n",
            "    0.01580442]]\n",
            "\n",
            " [[ 0.01396652 -0.04162994  0.10592167 ...  0.05138298 -0.00354845\n",
            "   -0.06019162]\n",
            "  [-0.02629372 -0.02878474  0.12608218 ...  0.00743495  0.07350878\n",
            "    0.02871001]\n",
            "  [-0.05186695 -0.03338452  0.08886429 ...  0.00248026  0.08310422\n",
            "   -0.00086674]\n",
            "  ...\n",
            "  [-0.04867133  0.00328152  0.11373516 ...  0.03528256  0.05657118\n",
            "   -0.01047443]\n",
            "  [-0.06869268 -0.0065574   0.13306202 ...  0.01718214  0.06655375\n",
            "   -0.00289538]\n",
            "  [-0.00519407 -0.02498811  0.0770431  ...  0.02835968  0.0722328\n",
            "   -0.0069797 ]]\n",
            "\n",
            " [[-0.00780144 -0.0386172   0.08690532 ... -0.01948372  0.12111165\n",
            "    0.00133108]\n",
            "  [ 0.00095356 -0.04524592  0.11314684 ... -0.0260413   0.09355883\n",
            "   -0.00945984]\n",
            "  [ 0.02721036 -0.03657889  0.072322   ...  0.01787562  0.10605418\n",
            "   -0.01261308]\n",
            "  ...\n",
            "  [-0.04423098  0.01033567  0.09104352 ... -0.00541039  0.13506898\n",
            "    0.01912957]\n",
            "  [ 0.03766819 -0.08432054  0.10706931 ... -0.02329985  0.07013367\n",
            "   -0.02522813]\n",
            "  [ 0.01194641 -0.02333062  0.08841702 ... -0.00164321  0.09414118\n",
            "   -0.01150878]]], shape=(3, 8, 7000), dtype=float32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "5M1KtC3Cil1G"
      },
      "execution_count": 62,
      "outputs": []
    }
  ]
}