main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f5f0bd2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "from dataclasses import dataclass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "048b4b4b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "419ae745",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelArgs:\n",
    "    batch_size: int = 32\n",
    "    d_model: int = 512\n",
    "    hidden_dim: int = 1024 # hidden dim for feed forward layer\n",
    "    num_blocks: int = 32\n",
    "    num_q_heads: int = 32\n",
    "    num_kv_heads: int = 16\n",
    "    vocab_size: int = -1 # initialized later\n",
    "    eps: float = 1e-6 # eps for RMSNorm\n",
    "    max_batch_size: int = 32\n",
    "    max_seq_len: int = 512\n",
    "    device: torch.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    freq_base: int = 10000\n",
    "    epochs: int = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0524e0d8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 16])\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[[-2.0066e+00, -1.5439e+00,  1.2610e+00,  ..., -1.3504e+00,\n",
       "           -1.2820e+00, -4.1116e-01],\n",
       "          [-1.0628e+00,  8.0939e-01,  8.9014e-01,  ...,  1.0547e+00,\n",
       "           -9.1446e-01,  4.3935e-01],\n",
       "          [-9.1582e-01, -2.2099e-01, -6.9655e-02,  ..., -2.7517e-01,\n",
       "            6.7879e-01, -3.5210e-02],\n",
       "          [-1.4875e-01,  2.6449e+00,  3.2572e-01,  ..., -1.0715e+00,\n",
       "            1.8116e+00,  9.8089e-03],\n",
       "          [-6.5293e-01, -1.3751e+00, -9.0568e-01,  ..., -1.2641e+00,\n",
       "            1.0315e+00, -7.1314e-01]],\n",
       "\n",
       "         [[ 2.5721e-01, -6.7754e-01, -2.2514e+00,  ..., -2.0556e+00,\n",
       "            2.7232e+00,  1.0822e+00],\n",
       "          [-4.7367e-01,  2.5354e-01, -4.7315e-01,  ..., -3.2901e+00,\n",
       "            2.1026e+00, -7.3989e-01],\n",
       "          [-1.4458e+00, -1.3209e+00,  4.0396e-02,  ..., -4.9933e-01,\n",
       "           -1.1447e-01,  8.6106e-01],\n",
       "          [-1.3568e+00, -3.3785e+00, -2.4562e-01,  ..., -3.1407e-01,\n",
       "           -1.4376e+00, -1.8868e-01],\n",
       "          [-3.6285e-01,  7.2419e-01,  4.5145e-01,  ..., -4.5855e-01,\n",
       "           -6.0713e-01, -1.2953e+00]],\n",
       "\n",
       "         [[ 1.4648e+00, -1.0993e-01, -1.0186e-01,  ..., -1.1249e+00,\n",
       "            1.5535e-02,  1.9790e+00],\n",
       "          [ 7.5893e-01, -7.3501e-01, -1.5220e+00,  ..., -1.4437e+00,\n",
       "           -5.5618e-01,  9.2010e-01],\n",
       "          [-2.0649e-02, -1.2605e+00,  3.0798e-01,  ...,  4.0228e-01,\n",
       "           -4.0514e-01,  3.7681e-01],\n",
       "          [-8.6477e-01, -1.2716e+00,  9.8016e-01,  ...,  8.0223e-02,\n",
       "            1.6373e+00, -1.0089e+00],\n",
       "          [ 1.9102e+00, -9.7154e-02,  9.8823e-01,  ..., -1.3751e+00,\n",
       "           -6.6131e-02, -4.2216e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-1.1765e+00, -1.1508e+00,  7.0547e-01,  ...,  1.3365e+00,\n",
       "            1.1585e+00, -1.3741e+00],\n",
       "          [-6.3515e-01,  4.6365e-01, -3.6857e-01,  ..., -5.9913e-02,\n",
       "           -1.4525e+00,  2.0682e+00],\n",
       "          [ 1.6898e-01, -3.9157e-01, -4.1761e-01,  ..., -1.5773e+00,\n",
       "            1.3565e+00, -1.4378e-01],\n",
       "          [-1.0781e+00,  1.1315e+00,  3.2901e-01,  ...,  1.0255e+00,\n",
       "           -4.7318e-01,  1.4581e+00],\n",
       "          [-3.1885e-01,  7.7627e-01,  2.0045e+00,  ...,  8.9667e-01,\n",
       "           -7.4813e-01,  4.7953e-03]],\n",
       "\n",
       "         [[-3.9927e-01, -5.9020e-01,  7.3254e-01,  ..., -2.0896e+00,\n",
       "            9.8633e-02, -3.7544e-01],\n",
       "          [-1.3234e+00,  8.3726e-01,  1.1456e+00,  ..., -4.2602e-01,\n",
       "           -5.7701e-01,  7.7021e-01],\n",
       "          [-1.5189e+00, -6.1889e-02, -1.7725e-01,  ..., -3.3823e-01,\n",
       "           -4.2636e-02,  8.2029e-01],\n",
       "          [ 4.8197e-01,  1.3567e+00,  1.1760e-01,  ..., -1.3285e-01,\n",
       "           -4.3787e-02,  8.6447e-01],\n",
       "          [ 3.7932e-01,  1.2004e+00,  2.2870e+00,  ..., -4.6180e-01,\n",
       "            3.5592e-01, -1.9797e+00]],\n",
       "\n",
       "         [[-1.0618e+00, -1.4746e-01,  1.9027e-01,  ...,  7.8172e-01,\n",
       "           -1.4112e+00,  9.0253e-02],\n",
       "          [ 1.2642e+00, -7.7483e-01, -8.9332e-02,  ...,  8.6511e-01,\n",
       "           -5.0917e-01,  1.3605e+00],\n",
       "          [ 3.3946e-01, -1.6169e+00, -6.2547e-01,  ..., -1.5009e+00,\n",
       "           -1.0171e-02, -3.3712e-02],\n",
       "          [ 5.6375e-01, -3.0481e-01, -3.9877e-01,  ...,  1.0751e+00,\n",
       "            7.0643e-01,  8.1723e-01],\n",
       "          [ 4.1197e-01, -1.1278e+00, -8.3383e-01,  ...,  2.6639e+00,\n",
       "           -1.5876e-02, -1.0073e+00]]],\n",
       "\n",
       "\n",
       "        [[[-7.0014e-01,  5.9468e-01, -4.0757e-01,  ..., -3.6780e-02,\n",
       "           -1.4123e+00, -9.0403e-02],\n",
       "          [-1.0892e+00,  2.1636e+00,  2.3641e-01,  ...,  1.0572e+00,\n",
       "            7.8685e-03, -2.4993e-01],\n",
       "          [ 1.0873e+00, -1.4320e-01,  4.0864e-01,  ...,  4.4946e-01,\n",
       "           -1.2053e+00, -4.6101e-01],\n",
       "          [-1.6777e-02,  5.8845e-01,  3.4405e-01,  ...,  8.0250e-01,\n",
       "           -7.7436e-01,  3.3795e-01],\n",
       "          [ 9.6548e-01, -1.5258e-01,  1.2137e+00,  ..., -1.2219e-01,\n",
       "            1.2458e+00, -7.7141e-01]],\n",
       "\n",
       "         [[ 1.5481e+00,  4.0900e-01,  1.4419e+00,  ...,  1.1452e+00,\n",
       "           -2.6057e-01, -7.2190e-01],\n",
       "          [-2.0257e-01,  2.1721e+00, -2.4774e-01,  ..., -7.1622e-01,\n",
       "            7.1951e-01,  3.6282e-01],\n",
       "          [-1.4345e+00,  9.2900e-01,  2.3656e+00,  ...,  1.4559e+00,\n",
       "           -1.8832e+00,  1.3685e+00],\n",
       "          [-6.3191e-01,  1.3668e+00, -5.7284e-02,  ..., -6.1396e-01,\n",
       "            1.2505e+00, -2.9646e-01],\n",
       "          [ 6.0034e-02,  1.6536e+00,  5.3967e-01,  ...,  7.0357e-01,\n",
       "           -7.7810e-01, -7.5708e-02]],\n",
       "\n",
       "         [[-3.2062e-01,  6.9114e-01, -3.7661e-01,  ...,  8.3855e-01,\n",
       "           -1.2209e-01, -1.6735e+00],\n",
       "          [ 1.8479e-01,  7.6741e-01, -5.2849e-02,  ..., -8.0333e-03,\n",
       "            3.2792e-01,  6.5380e-01],\n",
       "          [ 3.4704e-01,  6.4215e-01,  2.0402e-01,  ..., -4.3641e-01,\n",
       "           -5.6870e-01,  6.4560e-01],\n",
       "          [ 5.5668e-01, -3.7656e-01,  5.9017e-01,  ...,  8.5959e-01,\n",
       "            8.4171e-01,  5.5766e-01],\n",
       "          [ 1.2832e+00,  1.0986e+00, -1.6987e-01,  ..., -8.0364e-01,\n",
       "           -1.3860e+00,  1.4878e+00]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-2.3571e-01,  5.6015e-01,  8.7283e-02,  ..., -1.3244e-01,\n",
       "           -3.4165e-01, -1.0373e+00],\n",
       "          [-9.4622e-02, -5.8899e-01, -2.0091e+00,  ..., -2.7590e-01,\n",
       "            1.5445e+00, -8.8277e-01],\n",
       "          [ 3.7273e-03,  1.8650e+00,  4.2665e-01,  ..., -1.5833e-01,\n",
       "            1.6378e-01,  1.4145e+00],\n",
       "          [ 4.7441e-01, -2.5147e+00, -1.8675e+00,  ..., -9.8884e-01,\n",
       "           -7.3860e-01, -8.4181e-01],\n",
       "          [ 1.6623e-01, -8.4306e-01,  1.1377e+00,  ...,  5.7030e-01,\n",
       "            1.5223e+00,  1.2139e+00]],\n",
       "\n",
       "         [[-9.3500e-01, -6.9319e-01,  2.6616e-01,  ..., -3.6527e-01,\n",
       "            3.6296e-01,  1.9539e+00],\n",
       "          [-7.8794e-01, -3.0549e-01, -1.5884e+00,  ...,  8.0416e-01,\n",
       "            6.8835e-01,  9.7471e-01],\n",
       "          [-5.4176e-02, -4.4060e-01,  2.6040e-01,  ..., -3.2612e-03,\n",
       "           -1.6719e+00, -1.2157e-01],\n",
       "          [ 1.0374e+00,  1.2041e+00, -5.4827e-02,  ..., -9.8862e-01,\n",
       "            3.5002e-01, -9.4213e-01],\n",
       "          [-2.6654e+00,  4.3520e-02, -6.4352e-01,  ..., -2.0385e+00,\n",
       "           -1.5940e+00,  1.5404e+00]],\n",
       "\n",
       "         [[-8.2994e-01,  8.5769e-01, -1.4820e+00,  ..., -2.0890e+00,\n",
       "           -4.0971e-02,  1.1923e+00],\n",
       "          [ 1.1248e+00,  2.4338e-01, -9.1567e-01,  ...,  6.3998e-01,\n",
       "           -1.8640e+00, -3.9025e-01],\n",
       "          [-1.1390e+00, -8.7336e-01, -1.1673e+00,  ...,  2.3848e-01,\n",
       "            5.9850e-01, -6.0154e-01],\n",
       "          [-2.3959e+00,  2.2632e-01,  1.3441e+00,  ..., -1.9686e+00,\n",
       "            1.1736e+00, -8.1735e-01],\n",
       "          [-1.0456e+00,  4.2120e-01,  3.5546e-01,  ...,  1.7878e-01,\n",
       "           -5.5930e-01,  1.1639e+00]]],\n",
       "\n",
       "\n",
       "        [[[-2.8730e-01,  6.9461e-02, -1.2277e-01,  ...,  2.5765e-01,\n",
       "           -3.0880e-01,  2.5602e-01],\n",
       "          [-1.1486e+00, -7.8977e-01,  5.9295e-01,  ...,  4.8510e-01,\n",
       "           -2.2470e+00,  6.9537e-01],\n",
       "          [-1.3814e-01,  3.5833e-01, -1.0256e+00,  ..., -8.9628e-01,\n",
       "            6.4371e-01,  1.0175e-01],\n",
       "          [ 1.1309e+00, -2.3240e+00,  7.5688e-03,  ...,  8.4226e-01,\n",
       "           -9.3342e-01, -1.4444e+00],\n",
       "          [ 4.0725e-01,  3.8776e-01,  1.4677e+00,  ..., -3.9183e-01,\n",
       "            9.4333e-01,  5.7398e-01]],\n",
       "\n",
       "         [[ 5.8001e-01,  6.0729e-01, -7.9890e-01,  ...,  2.5973e-01,\n",
       "            1.0499e+00,  7.9049e-01],\n",
       "          [ 1.4329e+00, -1.0120e-01,  2.4199e+00,  ...,  7.6171e-01,\n",
       "           -1.1704e+00,  3.2148e-01],\n",
       "          [ 1.5925e+00,  2.3101e+00,  1.3279e+00,  ...,  1.2196e+00,\n",
       "            1.6259e+00, -7.0700e-01],\n",
       "          [ 9.3495e-01, -5.3398e-01,  1.1523e+00,  ..., -1.4859e+00,\n",
       "           -6.8994e-01,  1.9950e-01],\n",
       "          [-3.3724e-01, -1.9348e-02,  7.0825e-01,  ..., -1.7272e+00,\n",
       "           -3.5640e-01, -1.1563e-01]],\n",
       "\n",
       "         [[ 6.3546e-01,  1.6407e+00, -1.6750e+00,  ...,  1.1465e+00,\n",
       "           -9.1444e-01, -1.2736e+00],\n",
       "          [-6.8484e-01, -1.1922e+00,  1.6464e+00,  ..., -1.2735e-01,\n",
       "            6.6798e-01,  1.8518e-01],\n",
       "          [-5.6865e-02, -4.6970e-01,  4.9949e-01,  ...,  1.3083e+00,\n",
       "           -4.5296e-01,  1.4069e+00],\n",
       "          [-1.0973e-01, -4.0307e-01,  8.2983e-01,  ...,  1.5707e+00,\n",
       "            1.1781e+00, -1.5278e+00],\n",
       "          [ 8.4788e-02, -6.1234e-01,  1.8878e-01,  ..., -5.5528e-02,\n",
       "           -4.0496e-01,  5.2506e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 1.0148e+00, -1.3552e+00, -2.4865e-01,  ..., -8.1695e-01,\n",
       "           -5.4854e-01,  9.2495e-01],\n",
       "          [-4.8762e-01,  1.1414e+00,  1.2840e+00,  ..., -6.5458e-01,\n",
       "           -1.0621e+00, -6.6608e-01],\n",
       "          [ 3.0157e-01,  3.4385e-01,  1.1425e+00,  ..., -7.2241e-01,\n",
       "           -2.6500e+00,  1.2624e+00],\n",
       "          [ 2.1304e+00,  9.1707e-01, -1.8678e-01,  ..., -2.5436e-01,\n",
       "            6.6260e-01,  9.5578e-02],\n",
       "          [ 3.6192e-01, -8.9761e-01, -1.4479e+00,  ..., -1.2304e+00,\n",
       "           -3.4617e-01, -4.3774e-01]],\n",
       "\n",
       "         [[-7.9894e-01,  9.1024e-01,  3.5579e-01,  ..., -4.9517e-01,\n",
       "            9.4537e-01, -9.7790e-01],\n",
       "          [-7.0701e-01, -4.1658e-01,  8.2708e-01,  ..., -2.0458e-01,\n",
       "            5.9919e-01,  6.7199e-01],\n",
       "          [-2.1694e+00, -4.7610e-01,  3.3066e-01,  ..., -7.1010e-01,\n",
       "            7.1383e-01,  1.0649e-02],\n",
       "          [-2.2056e-01, -1.7150e+00,  8.7383e-01,  ...,  6.6278e-01,\n",
       "           -9.3843e-01, -8.2292e-01],\n",
       "          [ 2.2844e+00,  1.2955e+00,  7.8485e-02,  ...,  1.4646e+00,\n",
       "           -3.4431e-01, -6.1749e-01]],\n",
       "\n",
       "         [[ 2.4135e+00, -1.8237e+00,  8.3735e-02,  ...,  4.6759e-01,\n",
       "            9.7003e-01, -2.1595e+00],\n",
       "          [ 9.7725e-01,  3.4652e-01,  4.0147e-01,  ..., -1.2104e+00,\n",
       "            1.0395e-01, -2.7542e-01],\n",
       "          [-4.5721e-01, -6.0193e-01,  3.4434e-01,  ...,  6.9067e-01,\n",
       "            5.8327e-01,  6.8080e-02],\n",
       "          [ 1.0875e+00, -6.2153e-01, -1.9263e-02,  ..., -1.1249e+00,\n",
       "            4.8815e-01,  1.2159e+00],\n",
       "          [-2.9730e+00,  3.9516e-01, -1.0763e+00,  ...,  5.9868e-01,\n",
       "           -3.1276e-01,  1.7689e-01]]]], device='cuda:0')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def compute_freqs(d_model:int, seq_len: int, device: torch.device | str, base: int):\n",
    "    \"\"\"\n",
    "    d_model: embedding dim\n",
    "    seq_len: sequence length\n",
    "    device: cuda / cpu\n",
    "    base: base for exponential of theta values\n",
    "    \"\"\"\n",
    "\n",
    "    assert d_model % 2 == 0, \"d_model has to be even\"\n",
    "    \n",
    "    theta = 1. / (base ** (torch.arange(0, d_model, 2) / d_model)).to(device)\n",
    "    m = torch.arange(seq_len).to(device)\n",
    "    freqs = torch.outer(m, theta).float() # Since each m value corresponds to a single token, multiply every value of m by every value of theta, kind of like a nested for loop.\n",
    "    freqs_complex = torch.polar(torch.ones_like(freqs), freqs) # turn into complex form, z = r*cis(theta), in this case, r is 1\n",
    "    return freqs_complex\n",
    "\n",
    "def apply_rotary_embeddings(x: torch.Tensor, freqs_complex: torch.Tensor, device: torch.device):\n",
    "    \"\"\"\n",
    "    x: input sequence to add positional embedding, (batch, seq_len, emb_dim)\n",
    "    freqs_complex: frequencies for rotary postitional embeddings\n",
    "    device: cuda / cpu\n",
    "    \"\"\"\n",
    "    # print(x.shape, freqs_complex.shape)\n",
    "    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) # (batch, seq_len, _, 2)\n",
    "    # print(x_complex.shape)\n",
    "    freqs_complex = freqs_complex.unsqueeze(0).unsqueeze(2)\n",
    "    # print(freqs_complex.shape)\n",
    "    x_rotated = x_complex * freqs_complex # * is for element-wise multiplication\n",
    "    # print(x_rotated.shape)\n",
    "    x_rotated = torch.view_as_real(x_rotated)\n",
    "    # print(x_rotated.shape)\n",
    "    x_rotated = x_rotated.reshape(*x.shape)\n",
    "    # print(x_rotated.shape)\n",
    "    return x_rotated.type_as(x).to(device)\n",
    "\n",
    "a = compute_freqs(32, 10, torch.device(\"cuda\"), 10000)\n",
    "print(a.shape)\n",
    "b = torch.randn((3, 10, 5, 32)).to(torch.device(\"cuda\"))\n",
    "apply_rotary_embeddings(b, a, torch.device(\"cuda\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "24a0d6e2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 10, 32])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "class RMSNorm(nn.Module):\n",
    "    def __init__(self, args: ModelArgs):\n",
    "        super().__init__()\n",
    "        self.eps = args.eps\n",
    "        self.gamma = nn.Parameter(torch.ones(args.d_model))\n",
    "\n",
    "    def rms(self, x: torch.Tensor):\n",
    "        x = torch.pow(x, 2)\n",
    "        x = torch.mean(x, dim=-1, keepdim=True)\n",
    "        x = torch.sqrt(x + self.eps) # add eps to in case x = 0 (sqrt(0) is undefined in math)\n",
    "        return x\n",
    "\n",
    "    def forward(self, x: torch.Tensor):\n",
    "        return x / self.rms(x) * self.gamma\n",
    "\n",
    "args = ModelArgs(d_model=32)\n",
    "r = RMSNorm(args)\n",
    "a = torch.randn((5, 10, 32))\n",
    "r(a).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3e82dce1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 10, 32])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, args: ModelArgs):\n",
    "        super().__init__()\n",
    "        self.W = nn.Linear(args.d_model, args.hidden_dim)\n",
    "        self.V = nn.Linear(args.d_model, args.hidden_dim)\n",
    "        self.f = nn.Linear(args.hidden_dim, args.d_model)\n",
    "    \n",
    "    def forward(self, x: torch.Tensor):\n",
    "        swiglu = F.silu(self.W(x)) * self.V(x)\n",
    "        return self.f(swiglu)\n",
    "\n",
    "l = FeedForward(args)\n",
    "a = torch.randn((5, 10, 32))\n",
    "l(a).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "83d8b8de",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, args: ModelArgs):\n",
    "        super().__init__()\n",
    "        self.args = args\n",
    "        self.num_rep = args.num_q_heads // args.num_kv_heads\n",
    "        \n",
    "        self.q = nn.Linear(args.d_model, args.num_q_heads * args.d_model, bias=False)\n",
    "        self.k = nn.Linear(args.d_model, args.num_kv_heads * args.d_model, bias=False)\n",
    "        self.v = nn.Linear(args.d_model, args.num_kv_heads * args.d_model, bias=False)\n",
    "\n",
    "        self.out = nn.Linear(args.num_q_heads * args.d_model, args.d_model, bias=False)\n",
    "        self.cache_k = torch.zeros((args.max_batch_size, args.max_seq_len, args.num_kv_heads, args.d_model))\n",
    "        self.cache_v = torch.zeros((args.max_batch_size, args.max_seq_len, args.num_kv_heads, args.d_model))\n",
    "\n",
    "    \n",
    "    def repeat_kv(self, x: torch.Tensor, n_rep: int):\n",
    "        batch_size, seq_len, num_kv_heads, emb_dim = x.shape\n",
    "        if n_rep == 1:\n",
    "            return x\n",
    "        return x[:, :, :, None, :].expand(batch_size, seq_len, num_kv_heads, n_rep, emb_dim).reshape(batch_size, seq_len, num_kv_heads * n_rep, emb_dim)\n",
    "\n",
    "    \n",
    "    def forward(self, x: torch.Tensor, start_pos: int, freqs_complex: torch.Tensor, use_kv_cache=False):\n",
    "        batch_size, seq_len, _ = x.shape\n",
    "\n",
    "        q = self.q(x) # (batch_size, seq_len, head_num, emb_dim)\n",
    "        k = self.k(x)\n",
    "        v = self.v(x)\n",
    "        \n",
    "        q = q.view(batch_size, seq_len, self.args.num_q_heads, self.args.d_model)\n",
    "        k = k.view(batch_size, seq_len, self.args.num_kv_heads, self.args.d_model)\n",
    "        v = v.view(batch_size, seq_len, self.args.num_kv_heads, self.args.d_model)\n",
    "\n",
    "        q = apply_rotary_embeddings(q, freqs_complex, device=self.args.device)\n",
    "        k = apply_rotary_embeddings(k, freqs_complex, device=self.args.device)\n",
    "        \n",
    "        keys, values = k, v\n",
    "        if use_kv_cache:\n",
    "            self.cache_k[:batch_size, start_pos: start_pos + seq_len] = k\n",
    "            self.cache_v[:batch_size, start_pos: start_pos + seq_len] = v\n",
    "\n",
    "            keys = self.cache_k[:batch_size, :start_pos + seq_len] # all cache including added key\n",
    "            values = self.cache_v[:batch_size, :start_pos + seq_len]\n",
    "\n",
    "        keys = self.repeat_kv(keys, self.num_rep)\n",
    "        values = self.repeat_kv(values, self.num_rep)\n",
    "\n",
    "        q = q.permute(0, 2, 1, 3)\n",
    "        keys = keys.permute(0, 2, 3, 1)\n",
    "        values = values.permute(0, 2, 3, 1) # (batch_size, head_num, emb_dim, seq_len)\n",
    "        \n",
    "        scores = torch.matmul(q, keys) / self.args.d_model # (batch_size, head_num, seq_len, seq_len)\n",
    "        scores = F.softmax(scores.float(), dim=-1).type_as(q)\n",
    "        \n",
    "        output = torch.matmul(scores, values) # (batch_size, head_num, seq_len, emb_dim)\n",
    "        output = output.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_len, -1)\n",
    "\n",
    "        return self.out(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4f361e1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "    def __init__(self, args: ModelArgs):\n",
    "        super().__init__()\n",
    "        self.attention = MultiHeadAttention(args)\n",
    "        self.feed_forward = FeedForward(args)\n",
    "        self.norm1 = RMSNorm(args)\n",
    "        self.norm2 = RMSNorm(args)\n",
    "\n",
    "    def forward(self, x: torch.Tensor, start_pos: int, freqs_complex: torch.Tensor):\n",
    "        h = x + self.attention(self.norm1(x), start_pos, freqs_complex)\n",
    "        out = h + self.feed_forward(self.norm2(h))\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7eb5f3e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, args: ModelArgs):\n",
    "        super().__init__()\n",
    "\n",
    "        self.args = args\n",
    "\n",
    "        self.embedding_layer = nn.Embedding(args.vocab_size, args.d_model)\n",
    "\n",
    "        self.encoder_layers = nn.ModuleList()\n",
    "        for _ in range(args.num_blocks):\n",
    "            self.encoder_layers.append(EncoderBlock(args))\n",
    "\n",
    "        self.norm = RMSNorm(args)\n",
    "        self.output = nn.Linear(args.d_model, args.vocab_size, bias=False)\n",
    "        self.freqs_complex = compute_freqs(args.d_model, args.max_seq_len, device=args.device, base=args.freq_base)\n",
    "\n",
    "    def forward(self, tokens: torch.Tensor, start_pos: int):\n",
    "        h = self.embedding_layer(tokens)\n",
    "        freqs_complex = self.freqs_complex[:, start_pos: start_pos + tokens.shape[1]]\n",
    "\n",
    "        h = self.norm(h)\n",
    "\n",
    "        for layer in self.encoder_layers:\n",
    "            h = layer(h, start_pos, freqs_complex)\n",
    "\n",
    "        output = self.output(h).float()\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "72548858",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class TransformerSparseEmbeddings(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "\n",
    "        self.args = args\n",
    "\n",
    "        # Use Sparse Embedding\n",
    "        self.embedding_layer = nn.EmbeddingBag(args.vocab_size, args.d_model, sparse=True)\n",
    "\n",
    "        self.encoder_layers = nn.ModuleList()\n",
    "        for _ in range(args.num_blocks):\n",
    "            self.encoder_layers.append(EncoderBlock(args))\n",
    "\n",
    "        self.norm = RMSNorm(args)\n",
    "        self.output = nn.Linear(args.d_model, args.vocab_size, bias=False)\n",
    "        self.freqs_complex = compute_freqs(args.d_model, args.max_seq_len, device=args.device, base=args.freq_base)\n",
    "\n",
    "    def forward(self, tokens: torch.Tensor, start_pos: int):\n",
    "        # Modify the embedding lookup to work with sparse embeddings\n",
    "        offsets = torch.arange(0, tokens.size(0) * tokens.size(1), tokens.size(1), device=tokens.device)\n",
    "        h = self.embedding_layer(tokens.view(-1), offsets)\n",
    "\n",
    "        h = h.view(tokens.size(0), tokens.size(1), -1)\n",
    "        freqs_complex = self.freqs_complex[:, start_pos: start_pos + tokens.shape[1]]\n",
    "\n",
    "        h = self.norm(h)\n",
    "\n",
    "        for layer in self.encoder_layers:\n",
    "            h = layer(h, start_pos, freqs_complex)\n",
    "\n",
    "        output = self.output(h).float()\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "32587b8b",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "import sentencepiece as spm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f039c26a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "/tmp/ipykernel_331128/2762783289.py:2: ParserWarning: Falling back to the 'python' engine because the 'c' engine does not support regex separators (separators > 1 char and different from '\\s+' are interpreted as regex); you can avoid this warning by specifying engine='python'.\n",
       "  df = pd.read_csv('./DeepLearning/Ja-En-LLaMA/en-ja.bicleaner05.txt', sep=\"\\\\t\", header=None)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_csv('./DeepLearning/Ja-En-LLaMA/en-ja.bicleaner05.txt', sep=\"\\\\t\", header=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "77a1e44e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>And everyone will not care that it is not you.</td>\n",
       "      <td>鼻・口のところはあらかじめ少し切っておくといいですね。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>And everyone will not care that it is not you.</td>\n",
       "      <td>アドレス置いとくので、消されないうちにメールくれたら嬉しいです。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Sponsored link This advertisement is displayed...</td>\n",
       "      <td>スポンサードリンク この広告は一定期間更新がない場合に表示されます。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Also, it will always be hidden when becoming a...</td>\n",
       "      <td>また、 プレミアムユーザー になると常に非表示になります。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>It will return to non-display when content upd...</td>\n",
       "      <td>コンテンツの更新が行われると非表示に戻ります。</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                   3  \\\n",
       "0     And everyone will not care that it is not you.   \n",
       "1     And everyone will not care that it is not you.   \n",
       "2  Sponsored link This advertisement is displayed...   \n",
       "3  Also, it will always be hidden when becoming a...   \n",
       "4  It will return to non-display when content upd...   \n",
       "\n",
       "                                    4  \n",
       "0         鼻・口のところはあらかじめ少し切っておくといいですね。  \n",
       "1    アドレス置いとくので、消されないうちにメールくれたら嬉しいです。  \n",
       "2  スポンサードリンク この広告は一定期間更新がない場合に表示されます。  \n",
       "3       また、 プレミアムユーザー になると常に非表示になります。  \n",
       "4             コンテンツの更新が行われると非表示に戻ります。  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df.head()[[3, 4]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "36c8e5e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<class 'pandas.core.frame.DataFrame'>\n",
       "RangeIndex: 25740835 entries, 0 to 25740834\n",
       "Data columns (total 5 columns):\n",
       " #   Column  Dtype  \n",
       "---  ------  -----  \n",
       " 0   0       object \n",
       " 1   1       object \n",
       " 2   2       float64\n",
       " 3   3       object \n",
       " 4   4       object \n",
       "dtypes: float64(1), object(4)\n",
       "memory usage: 981.9+ MB\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df.info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d554439e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Go to the original video hierarchy of the conversion source, copy and paste the following is fine. ffmpeg -i sample.mp4 -strict -2 video.webm summary I’ve been using the upload and embed method to Youtube to set up videos on the web.'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df[3][9]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6c507eec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# using trained tokenizers from http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl/\n",
    "english_tokenizer = spm.SentencePieceProcessor(\"./DeepLearning/Ja-En-LLaMA/enja_spm_models/spm.en.nopretok.model\")\n",
    "japanese_tokenizer = spm.SentencePieceProcessor(\"./DeepLearning/Ja-En-LLaMA/enja_spm_models/spm.ja.nopretok.model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "fef895ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([False]), array([False]))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df[3].isna().unique(), df[4].isna().unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "059c109b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['▁Go',\n",
       " '▁to',\n",
       " '▁the',\n",
       " '▁original',\n",
       " '▁video',\n",
       " '▁hierarchy',\n",
       " '▁of',\n",
       " '▁the',\n",
       " '▁conversion',\n",
       " '▁source',\n",
       " ',',\n",
       " '▁copy',\n",
       " '▁and',\n",
       " '▁paste',\n",
       " '▁the',\n",
       " '▁following',\n",
       " '▁is',\n",
       " '▁fine',\n",
       " '.',\n",
       " '▁',\n",
       " 'ff',\n",
       " 'mp',\n",
       " 'eg',\n",
       " '▁-',\n",
       " 'i',\n",
       " '▁sample',\n",
       " '.',\n",
       " 'mp',\n",
       " '4',\n",
       " '▁-',\n",
       " 'strict',\n",
       " '▁-',\n",
       " '2',\n",
       " '▁video',\n",
       " '.',\n",
       " 'web',\n",
       " 'm',\n",
       " '▁summary',\n",
       " '▁I',\n",
       " '’',\n",
       " 've',\n",
       " '▁been',\n",
       " '▁using',\n",
       " '▁the',\n",
       " '▁upload',\n",
       " '▁and',\n",
       " '▁embed',\n",
       " '▁method',\n",
       " '▁to',\n",
       " '▁You',\n",
       " 'tube',\n",
       " '▁to',\n",
       " '▁set',\n",
       " '▁up',\n",
       " '▁videos',\n",
       " '▁on',\n",
       " '▁the',\n",
       " '▁web',\n",
       " '.']"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "english_tokenizer.encode(\"Go to the original video hierarchy of the conversion source, copy and paste the following is fine. ffmpeg -i sample.mp4 -strict -2 video.webm summary I’ve been using the upload and embed method to Youtube to set up videos on the web.\", out_type=str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "25304c3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['▁',\n",
       " '年',\n",
       " '金',\n",
       " '▁日本',\n",
       " 'に住んでいる',\n",
       " '20',\n",
       " '歳',\n",
       " '~',\n",
       " '60',\n",
       " '歳の',\n",
       " '全ての',\n",
       " '人は',\n",
       " '、',\n",
       " '公的',\n",
       " '年',\n",
       " '金',\n",
       " '制度',\n",
       " 'に',\n",
       " '加入',\n",
       " 'しなければなりません',\n",
       " '。']"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "japanese_tokenizer.encode(\"年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。\", out_type=str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "03e3bc5a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(32000, 32000)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "english_tokenizer.vocab_size(), japanese_tokenizer.vocab_size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "41663432",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "/mnt/ianch-Secondary/Programming/linux_venv/lib/python3.12/site-packages/torchtext/vocab/__init__.py:4: UserWarning: \n",
       "/!\\ IMPORTANT WARNING ABOUT TORCHTEXT STATUS /!\\ \n",
       "Torchtext is deprecated and the last released version will be 0.18 (this one). You can silence this warning by calling the following at the beginnign of your scripts: `import torchtext; torchtext.disable_torchtext_deprecation_warning()`\n",
       "  warnings.warn(torchtext._TORCHTEXT_DEPRECATION_MSG)\n",
       "/mnt/ianch-Secondary/Programming/linux_venv/lib/python3.12/site-packages/torchtext/utils.py:4: UserWarning: \n",
       "/!\\ IMPORTANT WARNING ABOUT TORCHTEXT STATUS /!\\ \n",
       "Torchtext is deprecated and the last released version will be 0.18 (this one). You can silence this warning by calling the following at the beginnign of your scripts: `import torchtext; torchtext.disable_torchtext_deprecation_warning()`\n",
       "  warnings.warn(torchtext._TORCHTEXT_DEPRECATION_MSG)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from collections import Counter\n",
    "from torchtext.vocab import vocab\n",
    "\n",
    "\n",
    "def build_vocab(sentences, tokenizer):\n",
    "    counter = Counter()\n",
    "    for sentence in sentences:\n",
    "        counter.update(tokenizer.encode(sentence, out_type=str))\n",
    "    return vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>']) # specials: include special tokens in the mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "4827f7f7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "25740835 2574083\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(len(df), len(df) // 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "3de1fbb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.sample(frac=1)\n",
    "english = df[3].tolist()[:len(df) // 500]\n",
    "japanese = df[4].tolist()[:len(df) // 500]\n",
    "assert len(english) == len(japanese)\n",
    "l = len(english)\n",
    "train_en = english[:int(0.7 * l)]\n",
    "val_en = english[int(0.7 * l): int(0.85 * l)]\n",
    "test_en = english[int(0.85 * l): l]\n",
    "train_ja = japanese[:int(0.7 * l)]\n",
    "val_ja = japanese[int(0.7 * l): int(0.85 * l)]\n",
    "test_ja = japanese[int(0.85 * l): l]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "e32e9a52",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "36036 25740835\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(len(train_en), len(df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "2869b707",
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_ja = build_vocab(japanese, japanese_tokenizer)\n",
    "vocab_en = build_vocab(english, english_tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "b83b9b4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_process(ja, en):\n",
    "    data = []\n",
    "    for (raw_ja, raw_en) in zip(ja, en):\n",
    "        ja_tensor = torch.tensor([vocab_ja[token] for token in japanese_tokenizer.encode(raw_ja.strip(\"\\n\"), out_type=str)], dtype=torch.long)\n",
    "        en_tensor = torch.tensor([vocab_en[token] for token in english_tokenizer.encode(raw_en.rstrip(\"\\n\"), out_type=str)], dtype=torch.long)\n",
    "        data.append((ja_tensor, en_tensor))\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "46458077",
   "metadata": {},
   "outputs": [],
   "source": [
    "train = data_process(train_ja, train_en)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1c518ed1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'batch_size': 32, 'd_model': 512, 'hidden_dim': 1024, 'num_blocks': 32, 'num_q_heads': 32, 'num_kv_heads': 16, 'vocab_size': -1, 'eps': 1e-06, 'max_batch_size': 32, 'max_seq_len': 512, 'device': device(type='cuda'), 'freq_base': 10000, 'epochs': 5}\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "a = ModelArgs()\n",
    "print(vars(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c2f71e6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn.utils.rnn import pad_sequence\n",
    "\n",
    "PAD_IDX = vocab_ja['<pad>']\n",
    "BOS_IDX = vocab_ja['<bos>']\n",
    "EOS_IDX = vocab_ja['<eos>']\n",
    "\n",
    "def generate_batch(data_batch):\n",
    "    ja_batch, en_batch = [], []\n",
    "    for (ja_item, en_item) in data_batch:\n",
    "        ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))\n",
    "        en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))\n",
    "    ja_batch = pad_sequence(ja_batch, padding_value=PAD_IDX) # pad sequences into equal length\n",
    "    en_batch = pad_sequence(en_batch, padding_value=PAD_IDX)\n",
    "    return ja_batch, en_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "ac6ad1c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "args = ModelArgs()\n",
    "\n",
    "train_iter = DataLoader(train, batch_size=args.batch_size, shuffle=True, collate_fn=generate_batch) # type: ignore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "564d6dc5",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "24181"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "len(vocab_ja)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "03f0f787",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'batch_size': 32, 'd_model': 512, 'hidden_dim': 1024, 'num_blocks': 32, 'num_q_heads': 32, 'num_kv_heads': 16, 'vocab_size': 24181, 'eps': 1e-06, 'max_batch_size': 32, 'max_seq_len': 512, 'device': device(type='cuda'), 'freq_base': 10000, 'epochs': 5}\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "args.vocab_size = len(vocab_ja)\n",
    "# args = ModelArgs(batch_size=32, d_model=64, hidden_dim=512, num_blocks=8, num_q_heads=32, num_kv_heads=16, vocab_size=100)\n",
    "print(vars(args))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "40933f7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = ModelArgs(batch_size=32, d_model=64, hidden_dim=512, num_blocks=8, num_q_heads=32, num_kv_heads=16, vocab_size=len(vocab_ja))\n",
    "transformer = TransformerSparseEmbeddings(args)\n",
    "transformer = transformer.to(args.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "6f3379ee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([69, 32]) torch.Size([88, 32])\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for (i, j) in train_iter:\n",
    "    print(i.shape, j.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "ce52b590",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1127"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "len(train_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "e5903dd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)\n",
    "optimizer = torch.optim.Adam(transformer.parameters(), lr=0.0001)\n",
    "\n",
    "def train_epoch(model, train_iter, optimizer, device): # https://chatgpt.com/c/2be13c26-50d8-4583-97b6-397c1fe2d028\n",
    "    model.train()\n",
    "    losses = 0\n",
    "    for i, (x, y) in enumerate(train_iter):\n",
    "        x = x.to(device)\n",
    "        y = y.to(device)\n",
    "        y = y[:-1, :]\n",
    "        logits = model(x, 0)\n",
    "\n",
    "        y_out = y[1:,:]\n",
    "        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), y_out.reshape(-1))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        losses += loss.item()\n",
    "    return losses / len(train_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "697dadb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def evaluate(model, val_iter, device):\n",
    "    model.eval()\n",
    "    losses = 0\n",
    "    for idx, (src, tgt) in (enumerate(val_iter)):\n",
    "        src = src.to(device)\n",
    "        tgt = tgt.to(device)\n",
    "        tgt_input = tgt[:-1, :]\n",
    "        logits = model(src, 0)\n",
    "        tgt_out = tgt[1:,:]\n",
    "        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))\n",
    "        losses += loss.item()\n",
    "    return losses / len(val_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "0e5c0aca",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "The size of tensor a (2) must match the size of tensor b (64) at non-singleton dimension 2",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[56], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m epoch \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m1\u001b[39m, args\u001b[38;5;241m.\u001b[39mepochs\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m):\n\u001b[1;32m      3\u001b[0m     start_time \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[0;32m----> 4\u001b[0m     train_loss \u001b[38;5;241m=\u001b[39m \u001b[43mtrain_epoch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtransformer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_iter\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      5\u001b[0m     end_time \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[1;32m      6\u001b[0m     \u001b[38;5;28mprint\u001b[39m((\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mEpoch: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mepoch\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Train loss: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtrain_loss\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m.3f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m      7\u001b[0m           \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mEpoch time = \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m(end_time\u001b[38;5;250m \u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;250m \u001b[39mstart_time)\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m.3f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124ms\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n",
      "Cell \u001b[0;32mIn[54], line 11\u001b[0m, in \u001b[0;36mtrain_epoch\u001b[0;34m(model, train_iter, optimizer, device)\u001b[0m\n\u001b[1;32m      9\u001b[0m y \u001b[38;5;241m=\u001b[39m y\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m     10\u001b[0m y \u001b[38;5;241m=\u001b[39m y[:\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, :]\n\u001b[0;32m---> 11\u001b[0m logits \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m     13\u001b[0m y_out \u001b[38;5;241m=\u001b[39m y[\u001b[38;5;241m1\u001b[39m:,:]\n\u001b[1;32m     14\u001b[0m loss \u001b[38;5;241m=\u001b[39m loss_fn(logits\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, logits\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]), y_out\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m))\n",
      "File \u001b[0;32m/mnt/ianch-Secondary/Programming/linux_venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1530\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1532\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/mnt/ianch-Secondary/Programming/linux_venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1539\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1540\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1541\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1544\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[14], line 29\u001b[0m, in \u001b[0;36mTransformerSparseEmbeddings.forward\u001b[0;34m(self, tokens, start_pos)\u001b[0m\n\u001b[1;32m     26\u001b[0m h \u001b[38;5;241m=\u001b[39m h\u001b[38;5;241m.\u001b[39mview(tokens\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m0\u001b[39m), tokens\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m1\u001b[39m), \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m     27\u001b[0m freqs_complex \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfreqs_complex[:, start_pos: start_pos \u001b[38;5;241m+\u001b[39m tokens\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m1\u001b[39m]]\n\u001b[0;32m---> 29\u001b[0m h \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnorm\u001b[49m\u001b[43m(\u001b[49m\u001b[43mh\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     31\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mencoder_layers:\n\u001b[1;32m     32\u001b[0m     h \u001b[38;5;241m=\u001b[39m layer(h, start_pos, freqs_complex)\n",
      "File \u001b[0;32m/mnt/ianch-Secondary/Programming/linux_venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1530\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1532\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/mnt/ianch-Secondary/Programming/linux_venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1539\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1540\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1541\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1544\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[9], line 14\u001b[0m, in \u001b[0;36mRMSNorm.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x: torch\u001b[38;5;241m.\u001b[39mTensor):\n\u001b[0;32m---> 14\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mx\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrms\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgamma\u001b[49m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: The size of tensor a (2) must match the size of tensor b (64) at non-singleton dimension 2"
     ]
    }
   ],
   "source": [
    "import time\n",
    "for epoch in range(1, args.epochs+1):\n",
    "    start_time = time.time()\n",
    "    train_loss = train_epoch(transformer, train_iter, optimizer, args.device)\n",
    "    end_time = time.time()\n",
    "    print((f\"Epoch: {epoch}, Train loss: {train_loss:.3f}, \"\n",
    "          f\"Epoch time = {(end_time - start_time):.3f}s\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e54eafe",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "4bf3ea7f",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}