encoder_decoder_poem.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3 (ipykernel)"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "metadata": {
    "id": "sqH2SEO-X9T6"
   },
   "cell_type": "markdown",
   "source": [
    "# Neural Networks on MNIST\n",
    "\n",
    "This Jupyter notebook explains various approaches for implementing neural networks that recognize digits on [MNIST](http://yann.lecun.com/exdb/mnist/) dataset."
   ]
  },
  {
   "metadata": {
    "id": "OaLeoEcDYVO_"
   },
   "cell_type": "markdown",
   "source": [
    "## Preparing the MNIST dataset\n",
    "\n",
    "Most deep learning frameworks provide APIs for loading famous datasets like MNIST (e.g., `torchvision.datasets.MNIST` in pytorch). The APIs are handy, but hide the important step for preparing a training data for a deep learning framework; when graduating from an example dataset to the real data, we must convert a training data of our interest into the data structure that is acceptable by a deep learning framework.\n",
    "\n",
    "The cell below downloads the original distribution of the MNIST dataset on the Web, converts the dataset into `numpy` arrays, and saves the arrays as the file `mnist.npz` with keyword names."
   ]
  },
  {
   "metadata": {
    "id": "OMaOcZMBPuQY",
    "ExecuteTime": {
     "end_time": "2023-10-24T14:25:30.535142369Z",
     "start_time": "2023-10-24T14:25:24.741474404Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import gzip\n",
    "import os\n",
    "import sys\n",
    "import struct\n",
    "import numpy as np\n",
    "\n",
    "def read_image(fi):\n",
    "    magic, n, rows, columns = struct.unpack(\">IIII\", fi.read(16))\n",
    "    assert magic == 0x00000803\n",
    "    assert rows == 28\n",
    "    assert columns == 28\n",
    "    rawbuffer = fi.read()\n",
    "    assert len(rawbuffer) == n * rows * columns\n",
    "    rawdata = np.frombuffer(rawbuffer, dtype='>u1', count=n*rows*columns)\n",
    "    return rawdata.reshape(n, rows, columns).astype(np.float32) / 255.0\n",
    "\n",
    "def read_label(fi):\n",
    "    magic, n = struct.unpack(\">II\", fi.read(8))\n",
    "    assert magic == 0x00000801\n",
    "    rawbuffer = fi.read()\n",
    "    assert len(rawbuffer) == n\n",
    "    return np.frombuffer(rawbuffer, dtype='>u1', count=n)\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    os.system('wget -N http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz')\n",
    "    os.system('wget -N http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz')\n",
    "    os.system('wget -N http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz')\n",
    "    os.system('wget -N http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz')\n",
    "\n",
    "    np.savez_compressed(\n",
    "        'mnist',\n",
    "        train_x=read_image(gzip.open('train-images-idx3-ubyte.gz', 'rb')),\n",
    "        train_y=read_label(gzip.open('train-labels-idx1-ubyte.gz', 'rb')),\n",
    "        test_x=read_image(gzip.open('t10k-images-idx3-ubyte.gz', 'rb')),\n",
    "        test_y=read_label(gzip.open('t10k-labels-idx1-ubyte.gz', 'rb'))\n",
    "    )"
   ],
   "execution_count": 1,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "--2023-10-24 16:25:24--  http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
      "Resolving yann.lecun.com (yann.lecun.com)... 2606:4700:3036::ac43:ab4c, 2606:4700:3034::6815:1d24, 172.67.171.76, ...\n",
      "Connecting to yann.lecun.com (yann.lecun.com)|2606:4700:3036::ac43:ab4c|:80... connected.\n",
      "HTTP request sent, awaiting response... 304 Not Modified\n",
      "File ‘train-images-idx3-ubyte.gz’ not modified on server. Omitting download.\n",
      "\n",
      "--2023-10-24 16:25:25--  http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
      "Resolving yann.lecun.com (yann.lecun.com)... 2606:4700:3034::6815:1d24, 2606:4700:3036::ac43:ab4c, 172.67.171.76, ...\n",
      "Connecting to yann.lecun.com (yann.lecun.com)|2606:4700:3034::6815:1d24|:80... connected.\n",
      "HTTP request sent, awaiting response... 304 Not Modified\n",
      "File ‘train-labels-idx1-ubyte.gz’ not modified on server. Omitting download.\n",
      "\n",
      "--2023-10-24 16:25:25--  http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
      "Resolving yann.lecun.com (yann.lecun.com)... 2606:4700:3034::6815:1d24, 2606:4700:3036::ac43:ab4c, 172.67.171.76, ...\n",
      "Connecting to yann.lecun.com (yann.lecun.com)|2606:4700:3034::6815:1d24|:80... connected.\n",
      "HTTP request sent, awaiting response... 304 Not Modified\n",
      "File ‘t10k-images-idx3-ubyte.gz’ not modified on server. Omitting download.\n",
      "\n",
      "--2023-10-24 16:25:25--  http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
      "Resolving yann.lecun.com (yann.lecun.com)... 2606:4700:3036::ac43:ab4c, 2606:4700:3034::6815:1d24, 104.21.29.36, ...\n",
      "Connecting to yann.lecun.com (yann.lecun.com)|2606:4700:3036::ac43:ab4c|:80... connected.\n",
      "HTTP request sent, awaiting response... 304 Not Modified\n",
      "File ‘t10k-labels-idx1-ubyte.gz’ not modified on server. Omitting download.\n",
      "\n"
     ]
    }
   ]
  },
  {
   "metadata": {
    "id": "TStlGwaUaZKC"
   },
   "cell_type": "markdown",
   "source": [
    "The file contains four numpy arrays (one tensor and array for each split of training and test sets) with the keywords:\n",
    "\n",
    "+ `train_x`: $60000 \\text{ (images)} \\times 28 \\text{ (y)} \\times 28 \\text{ (x)}$\n",
    "+ `train_y`: $60000 \\text{ (labels)}$\n",
    "+ `test_x`: $10000 \\text{ (images)} \\times 28 \\text{ (y)} \\times 28 \\text{ (x)}$\n",
    "+ `test_y`: $10000 \\text{ (labels)}$\n"
   ]
  },
  {
   "metadata": {
    "id": "QwF495MIe0-e"
   },
   "cell_type": "markdown",
   "source": [
    "### Install pytorch"
   ]
  },
  {
   "metadata": {
    "id": "h8U839iGXCH3",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "600dff09-c0ce-4a34-88b4-27b05fee529b",
    "ExecuteTime": {
     "end_time": "2023-10-24T14:25:33.339547070Z",
     "start_time": "2023-10-24T14:25:30.540359820Z"
    }
   },
   "cell_type": "code",
   "source": [
    "!pip install torch torchvision transformers spacy ftfy==4.4.3 pandas"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pypi.idi.ntnu.no\r\n",
      "Requirement already satisfied: torch in ./venv/lib/python3.11/site-packages (2.1.0)\r\n",
      "Requirement already satisfied: torchvision in ./venv/lib/python3.11/site-packages (0.16.0)\r\n",
      "Requirement already satisfied: transformers in ./venv/lib/python3.11/site-packages (4.34.0)\r\n",
      "Requirement already satisfied: spacy in ./venv/lib/python3.11/site-packages (3.7.2)\r\n",
      "Requirement already satisfied: ftfy==4.4.3 in ./venv/lib/python3.11/site-packages (4.4.3)\r\n",
      "Requirement already satisfied: pandas in ./venv/lib/python3.11/site-packages (2.1.1)\r\n",
      "Requirement already satisfied: html5lib in ./venv/lib/python3.11/site-packages (from ftfy==4.4.3) (1.1)\r\n",
      "Requirement already satisfied: wcwidth in ./venv/lib/python3.11/site-packages (from ftfy==4.4.3) (0.2.8)\r\n",
      "Requirement already satisfied: filelock in ./venv/lib/python3.11/site-packages (from torch) (3.12.4)\r\n",
      "Requirement already satisfied: typing-extensions in ./venv/lib/python3.11/site-packages (from torch) (4.8.0)\r\n",
      "Requirement already satisfied: sympy in ./venv/lib/python3.11/site-packages (from torch) (1.12)\r\n",
      "Requirement already satisfied: networkx in ./venv/lib/python3.11/site-packages (from torch) (3.1)\r\n",
      "Requirement already satisfied: jinja2 in ./venv/lib/python3.11/site-packages (from torch) (3.1.2)\r\n",
      "Requirement already satisfied: fsspec in ./venv/lib/python3.11/site-packages (from torch) (2023.9.2)\r\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in ./venv/lib/python3.11/site-packages (from torch) (12.1.105)\r\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in ./venv/lib/python3.11/site-packages (from torch) (12.1.105)\r\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in ./venv/lib/python3.11/site-packages (from torch) (12.1.105)\r\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in ./venv/lib/python3.11/site-packages (from torch) (8.9.2.26)\r\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in ./venv/lib/python3.11/site-packages (from torch) (12.1.3.1)\r\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in ./venv/lib/python3.11/site-packages (from torch) (11.0.2.54)\r\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in ./venv/lib/python3.11/site-packages (from torch) (10.3.2.106)\r\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in ./venv/lib/python3.11/site-packages (from torch) (11.4.5.107)\r\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in ./venv/lib/python3.11/site-packages (from torch) (12.1.0.106)\r\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.18.1 in ./venv/lib/python3.11/site-packages (from torch) (2.18.1)\r\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in ./venv/lib/python3.11/site-packages (from torch) (12.1.105)\r\n",
      "Requirement already satisfied: triton==2.1.0 in ./venv/lib/python3.11/site-packages (from torch) (2.1.0)\r\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12 in ./venv/lib/python3.11/site-packages (from nvidia-cusolver-cu12==11.4.5.107->torch) (12.2.140)\r\n",
      "Requirement already satisfied: numpy in ./venv/lib/python3.11/site-packages (from torchvision) (1.26.0)\r\n",
      "Requirement already satisfied: requests in ./venv/lib/python3.11/site-packages (from torchvision) (2.31.0)\r\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in ./venv/lib/python3.11/site-packages (from torchvision) (10.0.1)\r\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in ./venv/lib/python3.11/site-packages (from transformers) (0.17.3)\r\n",
      "Requirement already satisfied: packaging>=20.0 in ./venv/lib/python3.11/site-packages (from transformers) (23.2)\r\n",
      "Requirement already satisfied: pyyaml>=5.1 in ./venv/lib/python3.11/site-packages (from transformers) (6.0.1)\r\n",
      "Requirement already satisfied: regex!=2019.12.17 in ./venv/lib/python3.11/site-packages (from transformers) (2023.10.3)\r\n",
      "Requirement already satisfied: tokenizers<0.15,>=0.14 in ./venv/lib/python3.11/site-packages (from transformers) (0.14.1)\r\n",
      "Requirement already satisfied: safetensors>=0.3.1 in ./venv/lib/python3.11/site-packages (from transformers) (0.4.0)\r\n",
      "Requirement already satisfied: tqdm>=4.27 in ./venv/lib/python3.11/site-packages (from transformers) (4.66.1)\r\n",
      "Requirement already satisfied: spacy-legacy<3.1.0,>=3.0.11 in ./venv/lib/python3.11/site-packages (from spacy) (3.0.12)\r\n",
      "Requirement already satisfied: spacy-loggers<2.0.0,>=1.0.0 in ./venv/lib/python3.11/site-packages (from spacy) (1.0.5)\r\n",
      "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in ./venv/lib/python3.11/site-packages (from spacy) (1.0.10)\r\n",
      "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in ./venv/lib/python3.11/site-packages (from spacy) (2.0.8)\r\n",
      "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in ./venv/lib/python3.11/site-packages (from spacy) (3.0.9)\r\n",
      "Requirement already satisfied: thinc<8.3.0,>=8.1.8 in ./venv/lib/python3.11/site-packages (from spacy) (8.2.1)\r\n",
      "Requirement already satisfied: wasabi<1.2.0,>=0.9.1 in ./venv/lib/python3.11/site-packages (from spacy) (1.1.2)\r\n",
      "Requirement already satisfied: srsly<3.0.0,>=2.4.3 in ./venv/lib/python3.11/site-packages (from spacy) (2.4.8)\r\n",
      "Requirement already satisfied: catalogue<2.1.0,>=2.0.6 in ./venv/lib/python3.11/site-packages (from spacy) (2.0.10)\r\n",
      "Requirement already satisfied: weasel<0.4.0,>=0.1.0 in ./venv/lib/python3.11/site-packages (from spacy) (0.3.3)\r\n",
      "Requirement already satisfied: typer<0.10.0,>=0.3.0 in ./venv/lib/python3.11/site-packages (from spacy) (0.9.0)\r\n",
      "Requirement already satisfied: smart-open<7.0.0,>=5.2.1 in ./venv/lib/python3.11/site-packages (from spacy) (6.4.0)\r\n",
      "Requirement already satisfied: pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4 in ./venv/lib/python3.11/site-packages (from spacy) (2.4.2)\r\n",
      "Requirement already satisfied: setuptools in ./venv/lib/python3.11/site-packages (from spacy) (68.2.2)\r\n",
      "Requirement already satisfied: langcodes<4.0.0,>=3.2.0 in ./venv/lib/python3.11/site-packages (from spacy) (3.3.0)\r\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in ./venv/lib/python3.11/site-packages (from pandas) (2.8.2)\r\n",
      "Requirement already satisfied: pytz>=2020.1 in ./venv/lib/python3.11/site-packages (from pandas) (2023.3.post1)\r\n",
      "Requirement already satisfied: tzdata>=2022.1 in ./venv/lib/python3.11/site-packages (from pandas) (2023.3)\r\n",
      "Requirement already satisfied: annotated-types>=0.4.0 in ./venv/lib/python3.11/site-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy) (0.6.0)\r\n",
      "Requirement already satisfied: pydantic-core==2.10.1 in ./venv/lib/python3.11/site-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy) (2.10.1)\r\n",
      "Requirement already satisfied: six>=1.5 in ./venv/lib/python3.11/site-packages (from python-dateutil>=2.8.2->pandas) (1.16.0)\r\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in ./venv/lib/python3.11/site-packages (from requests->torchvision) (3.3.0)\r\n",
      "Requirement already satisfied: idna<4,>=2.5 in ./venv/lib/python3.11/site-packages (from requests->torchvision) (3.4)\r\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in ./venv/lib/python3.11/site-packages (from requests->torchvision) (2.0.6)\r\n",
      "Requirement already satisfied: certifi>=2017.4.17 in ./venv/lib/python3.11/site-packages (from requests->torchvision) (2023.7.22)\r\n",
      "Requirement already satisfied: blis<0.8.0,>=0.7.8 in ./venv/lib/python3.11/site-packages (from thinc<8.3.0,>=8.1.8->spacy) (0.7.11)\r\n",
      "Requirement already satisfied: confection<1.0.0,>=0.0.1 in ./venv/lib/python3.11/site-packages (from thinc<8.3.0,>=8.1.8->spacy) (0.1.3)\r\n",
      "Requirement already satisfied: click<9.0.0,>=7.1.1 in ./venv/lib/python3.11/site-packages (from typer<0.10.0,>=0.3.0->spacy) (8.1.7)\r\n",
      "Requirement already satisfied: cloudpathlib<0.17.0,>=0.7.0 in ./venv/lib/python3.11/site-packages (from weasel<0.4.0,>=0.1.0->spacy) (0.16.0)\r\n",
      "Requirement already satisfied: webencodings in ./venv/lib/python3.11/site-packages (from html5lib->ftfy==4.4.3) (0.5.1)\r\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in ./venv/lib/python3.11/site-packages (from jinja2->torch) (2.1.3)\r\n",
      "Requirement already satisfied: mpmath>=0.19 in ./venv/lib/python3.11/site-packages (from sympy->torch) (1.3.0)\r\n",
      "\r\n",
      "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m A new release of pip is available: \u001B[0m\u001B[31;49m23.2.1\u001B[0m\u001B[39;49m -> \u001B[0m\u001B[32;49m23.3\u001B[0m\r\n",
      "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m To update, run: \u001B[0m\u001B[32;49mpip install --upgrade pip\u001B[0m\r\n"
     ]
    }
   ]
  },
  {
   "metadata": {
    "id": "kjBd7K4lI3W5"
   },
   "cell_type": "markdown",
   "source": [
    "### Transformers"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "\n",
    "import re\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "from collections import Counter\n",
    "from typing import List, Tuple\n",
    "import pandas\n",
    "from transformers import OpenAIGPTTokenizer\n",
    "import pickle\n",
    "class BeowulfDataset(Dataset):\n",
    "    def __init__(self, file_path: str, sequence_length: int, sequence_length_enc: int, picklefile=None):\n",
    "        # Read the text file\n",
    "        text = pandas.read_csv(\"./gutenberg-poetry-dataset.csv\")\n",
    "        text = text[[\"title\", \"content\"]]\n",
    "        text = text.dropna()\n",
    "        # Remove non-alphanumeric characters and tokenize\n",
    "\n",
    "        # Build vocabulary\n",
    "        self.tokenizer = OpenAIGPTTokenizer.from_pretrained(\"openai-gpt\")\n",
    "        self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})\n",
    "        \n",
    "        # Create input-output pairs using sliding window\n",
    "        self.sequence_length = sequence_length\n",
    "        self.pairs = [(x_a[0], x_a[1][:10000]) for x_a in text.to_numpy()]\n",
    "        if picklefile is None:\n",
    "            for i in range(len(self.pairs)):\n",
    "                if len(self.tokenizer(self.pairs[i][0],return_tensors=\"pt\", truncation=True, max_length=self.sequence_length, padding=\"max_length\")[\"input_ids\"].squeeze()) == 0 or len(self.tokenizer(self.pairs[i][0],return_tensors=\"pt\", truncation=True, max_length=self.sequence_length, padding=\"max_length\")[\"input_ids\"].squeeze()) == 0:\n",
    "                    self.pairs.pop(i)\n",
    "                i -= 1\n",
    "            self.threes = []\n",
    "            o = 0\n",
    "            for pair in self.pairs:\n",
    "                o+=1\n",
    "                print(\"\\r\",o, end=\"\")\n",
    "                x, y = pair\n",
    "                x = self.tokenizer(x, return_tensors=\"pt\", truncation=True, max_length=sequence_length_enc, padding=\"max_length\")[\"input_ids\"].squeeze()\n",
    "                y = self.tokenizer(y, return_tensors=\"pt\", truncation=True, padding=\"max_length\", max_length=self.sequence_length+1)[\"input_ids\"].squeeze()\n",
    "                if y is None or len(y) == 0:\n",
    "                    continue\n",
    "                for i in range(min(len(y), self.sequence_length+1)):\n",
    "                    z = y[i]\n",
    "                    y_z = y.clone()\n",
    "                    y_z[i:] = self.tokenizer.pad_token_id\n",
    "                    y_z = y_z[:self.sequence_length]\n",
    "                    self.threes.append((x,y_z,z))\n",
    "            with open(\"./threes.pickle\", \"wb\") as f:\n",
    "                pickle.dump(self.threes, f)\n",
    "        else:\n",
    "            with open(picklefile, \"rb\") as f:\n",
    "                self.threes = pickle.load(f)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.threes)\n",
    "\n",
    "    def __getitem__(self, idx) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n",
    "        return self.threes[idx]\n",
    "\n",
    "    def num_unique_tokens(self):\n",
    "        return len(self.tokenizer)\n",
    "    \n",
    "\n",
    "    def tokens_to_text(self, tokens:torch.Tensor) -> str:\n",
    "        \"\"\"\n",
    "        Convert a list of token IDs back to a string of text.\n",
    "        \"\"\"\n",
    "        if tokens.shape[0] == 1:\n",
    "            tokens = [tokens.item()]\n",
    "        else: \n",
    "            tokens = tokens.tolist()\n",
    "        \n",
    "        txt = [self.tokenizer.convert_tokens_to_string(self.tokenizer.convert_ids_to_tokens(token)) for token in tokens]\n",
    "        return ' '.join(txt)\n",
    "    \n",
    "    def text_to_tokens(self, text: str) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Convert a list of token IDs back to a string of text.\n",
    "        \"\"\"\n",
    "        out = self.tokenizer(text)\n",
    "        return out\n",
    "\n",
    "\n",
    "file_path = 'gutenberg-poetry-dataset.csv'\n",
    "batch_size = 128\n",
    "shuffle = True\n",
    "sequence_length = 100  # for example, predicting 10th word based on previous 9 words\n",
    "sequence_length_enc = 25\n",
    "\n",
    "dataset = BeowulfDataset(file_path, sequence_length, sequence_length_enc)\n",
    "dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)\n",
    "\n",
    "# Example usage:\n",
    "print(f\"Number of unique tokens in the dataset: {dataset.num_unique_tokens()}\")\n",
    "\n",
    "print(dataset[128])\n",
    "uniques = dataset.num_unique_tokens()\n",
    "print(uniques)\n",
    "\n"
   ],
   "metadata": {
    "id": "dSW1Tas9HKyi",
    "outputId": "99a3dfa3-a102-43e4-a15f-45479ddc4a90",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "ExecuteTime": {
     "end_time": "2023-10-24T14:52:00.577208129Z",
     "start_time": "2023-10-24T14:51:01.634144306Z"
    }
   },
   "execution_count": 18,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 1150Number of unique tokens in the dataset: 40479\n",
      "(tensor([13329, 30488,   535,  4562,  1024, 14585,   702,  3878, 16262,  1580,\n",
      "        26621, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478,\n",
      "        40478, 40478, 40478, 40478, 40478]), tensor([  498,   525, 10519,  2455,   240,  3690,  4665,   632,   523, 40477,\n",
      "         1739,  1767,   666,   481,  1276,   240,   488,   589,   622, 29992,\n",
      "          240, 40477,   556,  4250,   498,  7756,   240, 40478, 40478, 40478,\n",
      "        40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478,\n",
      "        40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478,\n",
      "        40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478,\n",
      "        40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478,\n",
      "        40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478,\n",
      "        40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478,\n",
      "        40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478, 40478]), tensor(5802))\n",
      "40479\n"
     ]
    }
   ]
  },
  {
   "metadata": {
    "id": "DdSGWuLffzEg",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 436
    },
    "outputId": "64364756-6394-4c7e-87f9-9db3ba4e6bfc",
    "ExecuteTime": {
     "end_time": "2023-10-24T14:26:35.154734341Z",
     "start_time": "2023-10-24T14:26:35.145583281Z"
    }
   },
   "cell_type": "code",
   "source": [
    "\n"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "import math\n",
    "from matplotlib import pyplot as plt\n",
    "#torch.manual_seed(0)\n",
    "class self_attention_pure(nn.Module):\n",
    "  def __init__(self, device=\"cuda:0\", p_in=12, mid = 512, dims=0, num_heads=4):\n",
    "    super().__init__()\n",
    "    self.num_words = p_in\n",
    "    self.num_dims = p_in\n",
    "    self.num_heads = num_heads\n",
    "    if dims != 0:\n",
    "      self.num_dims = dims\n",
    "    if self.num_dims % num_heads  != 0:\n",
    "      raise \"num heads must divide num dims\"\n",
    "\n",
    "    self.Q = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
    "    self.K = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
    "    self.V = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
    "    self.softmax = nn.Softmax(dim=3)\n",
    "    self.preoutput_0 = nn.Linear(self.num_dims,mid).to(device)\n",
    "    self.preoutput_1 = nn.Linear(mid, self.num_dims).to(device)\n",
    "    self.gelu = nn.GELU()\n",
    "    self.device=device\n",
    "    self.normalize = nn.LayerNorm([p_in, self.num_dims]).to(device)\n",
    "\n",
    "\n",
    "  def forward(self, x, device=\"cuda:0\", visualization=False):\n",
    "    x = x.to(self.device)\n",
    "    id = x\n",
    "    x_q = self.Q(x).view(x.shape[0], -1, self.num_heads, self.num_dims//self.num_heads)\n",
    "    x_k = self.K(x).view(x.shape[0], -1, self.num_heads, self.num_dims//self.num_heads)\n",
    "    x_v = self.V(x).view(x.shape[0], -1, self.num_heads, self.num_dims//self.num_heads)\n",
    "    x_q = x_q.permute(0, 2,1,3)\n",
    "    x_k = x_k.permute(0, 2,1,3)\n",
    "    x_v = x_v.permute(0, 2,1,3)\n",
    "\n",
    "    x_qk = x_q.matmul(x_k.transpose(2,3))/math.sqrt(self.num_dims)\n",
    "    x = self.softmax(x_qk).matmul(x_v)\n",
    "    if visualization:\n",
    "      vis_softmax = nn.Softmax(dim=1)\n",
    "      for i in range(x_qk.shape[1]):\n",
    "        print(x_qk.shape)\n",
    "        plt.imshow(vis_softmax(x_qk[0][i]).clone().detach().cpu(), vmin=0, vmax=1)\n",
    "        plt.show()\n",
    "      print(self.softmax(x_qk).shape)\n",
    "    x = x.permute(0, 2, 1, 3)\n",
    "    x = x.reshape(x.shape[0], -1, self.num_dims)\n",
    "    x = self.preoutput_0(x)\n",
    "    x = self.preoutput_1(x)\n",
    "    x += id\n",
    "    x = self.normalize(x)\n",
    "    x = self.gelu(x)\n",
    "\n",
    "    return x\n",
    "\n",
    "\n",
    "class self_attention_w_pos_enc(nn.Module):\n",
    "  def __init__(self, device=\"cuda:0\", words = 12, dims = 12):\n",
    "    super().__init__()\n",
    "    self.num_words = words\n",
    "    self.num_dims = dims\n",
    "    self.embedding = nn.Embedding(uniques,self.num_dims).to(device)\n",
    "    self.pos_encoding = torch.zeros((self.num_words,self.num_dims))\n",
    "\n",
    "    for i in range(self.num_words):\n",
    "      for j in range(self.num_dims):\n",
    "        if j % 2 == 0:\n",
    "          self.pos_encoding[i,j] = math.sin(i/math.pow(10000, j/self.num_dims))\n",
    "        else:\n",
    "          self.pos_encoding[i,j] = math.cos(i/math.pow(10000, j-1/self.num_dims))\n",
    "    self.pos_encoding = self.pos_encoding.to(device)\n",
    "\n",
    "    self.attention = self_attention_pure(device=device, p_in = self.num_dims, dims = self.num_dims, mid=512)\n",
    "    self.output = nn.Linear(512, 519).to(device)\n",
    "\n",
    "  def forward(self, x, device=\"cuda:0\", visualization=False):\n",
    "    x = x.to(device)\n",
    "    x = self.embedding(x)\n",
    "    x += self.pos_encoding\n",
    "    x = self.attention(x, visualization=visualization)\n",
    "    x = self.output(x)\n",
    "\n",
    "    return x\n",
    "\n",
    "class self_attention_w_pos_enc_pure(nn.Module):\n",
    "  def __init__(self, device=\"cuda:0\", num_words = 12, num_dims=12, out=512, num_heads=4, emb_size=uniques):\n",
    "    super().__init__()\n",
    "    self.num_words = num_words\n",
    "    self.num_dims = num_dims\n",
    "    self.embedding = nn.Embedding(emb_size,self.num_dims).to(device)\n",
    "    self.pos_encoding = torch.zeros((self.num_words,self.num_dims))\n",
    "    self.device=device\n",
    "\n",
    "    for i in range(self.num_words):\n",
    "      for j in range(self.num_dims):\n",
    "        if j % 2 == 0:\n",
    "          self.pos_encoding[i,j] = math.sin(i/math.pow(10000, j/self.num_dims))\n",
    "        else:\n",
    "          self.pos_encoding[i,j] = math.cos(i/math.pow(10000, (j-1)/self.num_dims))\n",
    "    self.pos_encoding = self.pos_encoding.to(device)\n",
    "\n",
    "    self.attention = self_attention_pure(device=device, p_in = self.num_words, dims = self.num_dims, mid=out, num_heads=num_heads)\n",
    "\n",
    "  def forward(self, x, device=\"cuda:0\", visualization=False):\n",
    "    x = x.to(self.device)\n",
    "    x = self.embedding(x)\n",
    "    x += self.pos_encoding\n",
    "    x = self.attention(x, visualization=visualization)\n",
    "    return x\n",
    "\n",
    "\n",
    "class multi_layer_self_attention(nn.Module):\n",
    "  def __init__(self, device=\"cuda:0\", outs=[144,256,512,519], seq_length=12, dims=12, num_heads=4):\n",
    "    super().__init__()\n",
    "    self.attention_w_enc = self_attention_w_pos_enc_pure(device=device, num_words=seq_length, num_dims=dims, out=outs[0], num_heads=num_heads)\n",
    "    self.attention_1 = self_attention_pure(device=device, p_in=seq_length, dims=dims, mid=outs[0], num_heads=num_heads)\n",
    "    self.attention_2 = self_attention_pure(device=device, p_in=seq_length, dims=dims, mid=outs[1], num_heads=num_heads)\n",
    "    self.attention_3 = self_attention_pure(device=device, p_in=seq_length, dims=dims, mid=outs[2], num_heads=num_heads)\n",
    "    self.linear_1 = nn.Linear(seq_length*dims, outs[0]).to(device)\n",
    "    self.linear_2 = nn.Linear(outs[0], outs[1]).to(device)\n",
    "    self.flatten = nn.Flatten(start_dim=1)\n",
    "    self.output = nn.Linear(outs[1], outs[3]).to(device)\n",
    "    self.device=device\n",
    "\n",
    "  def forward(self, x, visualization=False):\n",
    "    x = x.to(self.device)\n",
    "    x = self.attention_w_enc(x, visualization=visualization)\n",
    "    x = self.attention_1(x, visualization=visualization)\n",
    "    x = self.attention_2(x, visualization=False)\n",
    "    x = self.attention_3(x, visualization=False)\n",
    "\n",
    "    x = self.flatten(x)\n",
    "    x = self.linear_1(x)\n",
    "    x = self.linear_2(x)\n",
    "    x = self.output(x)\n",
    "    return x\n",
    "\n",
    "seq_length = sequence_length\n",
    "self_atten = multi_layer_self_attention(seq_length=seq_length, dims=256, num_heads=8, outs=[768,1024,512,uniques])\n",
    "#p = self_atten(numdata[:12])\n",
    "\n",
    "optimizer = torch.optim.Adam(self_atten.parameters(), lr = 0.00001, weight_decay=0.000005)\n",
    "loss = nn.CrossEntropyLoss()\n",
    "schedule = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=2, cooldown=3, factor=0.333)\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "for k in range(25):\n",
    "  loss_n = 0\n",
    "  acc = 0\n",
    "  for i, batch in enumerate(dataloader):\n",
    "    optimizer.zero_grad()\n",
    "    print(batch[0][\"input_ids\"].shape)\n",
    "    x = batch[0][\"input_ids\"].view((-1, 24))[:, :23]\n",
    "    print(x.shape)\n",
    "    y = torch.zeros((batch[0][\"input_ids\"].shape[0], uniques))\n",
    "    print(y.shape)\n",
    "    y = y.to(\"cuda:0\")\n",
    "    y[:,batch[0][\"input_ids\"].view((-1, 24))[:,23]] = 1\n",
    "    \n",
    "    y_p = self_atten(x, visualization=False)\n",
    "    l = loss(y_p, y)\n",
    "    l.backward()\n",
    "    optimizer.step()\n",
    "    loss_n += l/batch[0][\"input_ids\"].shape[0]\n",
    "    acc += torch.eq(y.argmax(dim=1), y_p.argmax(dim=1)).sum().item()/batch[0][\"input_ids\"].shape[0]\n",
    "    print(f\"\\rit: {i+1}/{len(dataloader)}, loss: {loss_n/(i+1)}, acc: {acc/(i+1)}\", end=\"\")\n",
    "  loss_n /= len(dataloader)\n",
    "  acc /= len(dataloader)\n",
    "  schedule.step(loss_n)\n",
    "  print(\"epoch:\", k ,\"loss:\", loss_n, \"acc:\",acc)\n",
    "\n",
    "\n",
    "#p = self_atten(numdata[:12])\n",
    "\n",
    "#print(p, words[p.argmax()], data[13])\n",
    "\n",
    "accuracy = 0\n",
    "for i, batch in enumerate(dataloader):\n",
    "  x = batch[0]\n",
    "  y = torch.zeros((batch[1].shape[0], uniques))\n",
    "  y[:,batch[1]] = 1\n",
    "  y_p = self_atten(x, visualization=False)\n",
    "  y_p = y_p.to(\"cpu\")\n",
    "  #print(y, words[y_p.argmax()])\n",
    "  accuracy += torch.eq(y.argmax(dim=1), y_p.argmax(dim=1)).sum().item()/batch[0].shape[0]\n",
    "\n",
    "accuracy /= len(dataloader)-seq_length-1\n",
    "print(\"accuracy:\", accuracy)\n",
    "\n",
    "print(dataset.tokens_to_text(dataset[128][0]), dataset.tokens_to_text(torch.Tensor([dataset[128][1]])))\n",
    "print(dataset[128][0].shape)\n",
    "x = torch.Tensor(dataset[128][0]).unsqueeze(0)\n",
    "self_atten(x, visualization=True)\n",
    "x = torch.Tensor(dataset[129][0]).unsqueeze(0)\n",
    "self_atten(x, visualization=True)\n",
    "\"\"\""
   ],
   "metadata": {
    "id": "X8H7JxBoLXeh",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "80e2fc86-f9a0-407d-bd70-2dd1357bedf3",
    "ExecuteTime": {
     "end_time": "2023-10-24T14:35:19.853781167Z",
     "start_time": "2023-10-24T14:35:18.921999596Z"
    }
   },
   "execution_count": 9,
   "outputs": [
    {
     "data": {
      "text/plain": "'\\n\\nfor k in range(25):\\n  loss_n = 0\\n  acc = 0\\n  for i, batch in enumerate(dataloader):\\n    optimizer.zero_grad()\\n    print(batch[0][\"input_ids\"].shape)\\n    x = batch[0][\"input_ids\"].view((-1, 24))[:, :23]\\n    print(x.shape)\\n    y = torch.zeros((batch[0][\"input_ids\"].shape[0], uniques))\\n    print(y.shape)\\n    y = y.to(\"cuda:0\")\\n    y[:,batch[0][\"input_ids\"].view((-1, 24))[:,23]] = 1\\n    \\n    y_p = self_atten(x, visualization=False)\\n    l = loss(y_p, y)\\n    l.backward()\\n    optimizer.step()\\n    loss_n += l/batch[0][\"input_ids\"].shape[0]\\n    acc += torch.eq(y.argmax(dim=1), y_p.argmax(dim=1)).sum().item()/batch[0][\"input_ids\"].shape[0]\\n    print(f\"\\rit: {i+1}/{len(dataloader)}, loss: {loss_n/(i+1)}, acc: {acc/(i+1)}\", end=\"\")\\n  loss_n /= len(dataloader)\\n  acc /= len(dataloader)\\n  schedule.step(loss_n)\\n  print(\"epoch:\", k ,\"loss:\", loss_n, \"acc:\",acc)\\n\\n\\n#p = self_atten(numdata[:12])\\n\\n#print(p, words[p.argmax()], data[13])\\n\\naccuracy = 0\\nfor i, batch in enumerate(dataloader):\\n  x = batch[0]\\n  y = torch.zeros((batch[1].shape[0], uniques))\\n  y[:,batch[1]] = 1\\n  y_p = self_atten(x, visualization=False)\\n  y_p = y_p.to(\"cpu\")\\n  #print(y, words[y_p.argmax()])\\n  accuracy += torch.eq(y.argmax(dim=1), y_p.argmax(dim=1)).sum().item()/batch[0].shape[0]\\n\\naccuracy /= len(dataloader)-seq_length-1\\nprint(\"accuracy:\", accuracy)\\n\\nprint(dataset.tokens_to_text(dataset[128][0]), dataset.tokens_to_text(torch.Tensor([dataset[128][1]])))\\nprint(dataset[128][0].shape)\\nx = torch.Tensor(dataset[128][0]).unsqueeze(0)\\nself_atten(x, visualization=True)\\nx = torch.Tensor(dataset[129][0]).unsqueeze(0)\\nself_atten(x, visualization=True)\\n'"
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-10-24T14:26:36.677662179Z",
     "start_time": "2023-10-24T14:26:36.673109700Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "it: 908/908, loss: 5.731567859649658, acc: 0.027730933370044054epoch: 0 loss: tensor(5.7316, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.027730933370044054\n",
      "it: 908/908, loss: 5.446281909942627, acc: 0.04154047356828194epoch: 1 loss: tensor(5.4463, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.04154047356828194\n",
      "it: 908/908, loss: 5.4168877601623535, acc: 0.03192972191629956epoch: 2 loss: tensor(5.4169, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.03192972191629956\n",
      "it: 908/908, loss: 5.401157379150391, acc: 0.023222398127753303epoch: 3 loss: tensor(5.4012, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.023222398127753303\n",
      "it: 908/908, loss: 5.4023966789245605, acc: 0.04499070759911894epoch: 4 loss: tensor(5.4024, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.04499070759911894\n",
      "it: 908/908, loss: 5.397100448608398, acc: 0.042116946585903085epoch: 5 loss: tensor(5.3971, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.042116946585903085\n",
      "it: 908/908, loss: 5.389350891113281, acc: 0.05094472742290749epoch: 6 loss: tensor(5.3894, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.05094472742290749\n",
      "it: 908/908, loss: 5.388332843780518, acc: 0.04357963931718062epoch: 7 loss: tensor(5.3883, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.04357963931718062\n",
      "it: 908/908, loss: 5.381434440612793, acc: 0.05180513491189427epoch: 8 loss: tensor(5.3814, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.05180513491189427\n",
      "it: 908/908, loss: 5.385585308074951, acc: 0.04655664922907489epoch: 9 loss: tensor(5.3856, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.04655664922907489\n",
      "it: 908/908, loss: 5.385111331939697, acc: 0.03524229074889868epoch: 10 loss: tensor(5.3851, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.03524229074889868\n",
      "it: 908/908, loss: 5.385129928588867, acc: 0.024340927863436123epoch: 11 loss: tensor(5.3851, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.024340927863436123\n",
      "it: 908/908, loss: 5.3625922203063965, acc: 0.021716685022026432epoch: 12 loss: tensor(5.3626, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.021716685022026432\n",
      "it: 908/908, loss: 5.361545085906982, acc: 0.014773196585903084epoch: 13 loss: tensor(5.3615, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.014773196585903084\n",
      "it: 908/908, loss: 5.359376430511475, acc: 0.012140349669603524epoch: 14 loss: tensor(5.3594, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.012140349669603524\n",
      "it: 908/908, loss: 5.359329700469971, acc: 0.009653772026431718epoch: 15 loss: tensor(5.3593, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.009653772026431718\n",
      "it: 908/908, loss: 5.3623785972595215, acc: 0.01682096640969163epoch: 16 loss: tensor(5.3624, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.01682096640969163\n",
      "it: 908/908, loss: 5.360917568206787, acc: 0.006607929515418502epoch: 17 loss: tensor(5.3609, device='cuda:0', grad_fn=<DivBackward0>) acc: 0.006607929515418502\n",
      "it: 816/908, loss: 5.355347633361816, acc: 0.00219247855392156868"
     ]
    }
   ],
   "source": [
    "class self_attention_pure_encoder_qk_decoder_v(nn.Module):\n",
    "  def __init__(self, device=\"cuda:0\", p_in=12, mid = 512, dims=0, num_heads=4):\n",
    "    super().__init__()\n",
    "    self.num_words = p_in\n",
    "    self.num_dims = p_in\n",
    "    self.num_heads = num_heads\n",
    "    if dims != 0:\n",
    "      self.num_dims = dims\n",
    "    if self.num_dims % num_heads  != 0:\n",
    "      raise \"num heads must divide num dims\"\n",
    "\n",
    "    self.Q = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
    "    self.K = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
    "    self.V = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
    "    self.softmax = nn.Softmax(dim=3)\n",
    "    self.preoutput_0 = nn.Linear(self.num_dims,mid).to(device)\n",
    "    self.preoutput_1 = nn.Linear(mid, self.num_dims).to(device)\n",
    "    self.gelu = nn.GELU()\n",
    "    self.device=device\n",
    "    self.normalize = nn.LayerNorm([p_in, self.num_dims]).to(device)\n",
    "\n",
    "\n",
    "  def forward(self, x, keyvalue,  device=\"cuda:0\", visualization=False):\n",
    "    x = x.to(self.device)\n",
    "    keyvalue = keyvalue.to(self.device)\n",
    "    id = x\n",
    "    x_q = self.Q(x).view(x.shape[0],-1, self.num_heads, self.num_dims//self.num_heads)\n",
    "    x_k = self.K(keyvalue).view(x.shape[0], -1, self.num_heads, self.num_dims//self.num_heads)\n",
    "    x_v = self.V(keyvalue).view(x.shape[0], -1, self.num_heads, self.num_dims//self.num_heads)\n",
    "\n",
    "    x_q = x_q.permute(0, 2, 1, 3)\n",
    "    x_k = x_k.permute(0, 2, 1, 3)\n",
    "    x_v = x_v.permute(0, 2, 1, 3)\n",
    "\n",
    "    x_qk = x_q.matmul(x_k.transpose(2,3))/math.sqrt(self.num_dims)\n",
    "    x = self.softmax(x_qk).matmul(x_v)\n",
    "    if visualization:\n",
    "      vis_softmax = nn.Softmax(dim=1)\n",
    "      for i in range(x_qk.shape[1]):\n",
    "        print(x_qk.shape)\n",
    "        plt.imshow(vis_softmax(x_qk[0][i]).clone().detach().cpu(), vmin=0, vmax=1)\n",
    "        plt.show()\n",
    "      print(self.softmax(x_qk).shape)\n",
    "    x = x.permute(0, 2, 1, 3)\n",
    "    x = x.reshape(x.shape[0], -1, self.num_dims)\n",
    "    x = self.preoutput_0(x)\n",
    "    x = self.preoutput_1(x)\n",
    "    x += id\n",
    "    x = self.normalize(x)\n",
    "    x = self.gelu(x)\n",
    "\n",
    "    return x\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class multi_layer_self_attention_encoder_decoder_scheme(nn.Module):\n",
    "  def __init__(self, device=\"cuda:0\", outs=[144,256,512,519], seq_length=12, dims=12, num_heads=4, emb_size=uniques):\n",
    "    super().__init__()\n",
    "    self.attention_w_enc = self_attention_w_pos_enc_pure(device=device, num_words=seq_length, num_dims=dims, out=outs[0], num_heads=num_heads, emb_size=emb_size)\n",
    "    self.attention_1 = self_attention_pure_encoder_qk_decoder_v(device=device, p_in=seq_length, dims=dims, mid=outs[0], num_heads=num_heads)\n",
    "    self.attention_2 = self_attention_pure_encoder_qk_decoder_v(device=device, p_in=seq_length, dims=dims, mid=outs[1], num_heads=num_heads)\n",
    "    self.attention_3 = self_attention_pure(device=device, p_in=seq_length, dims=dims, mid=outs[2], num_heads=num_heads)\n",
    "\n",
    "    self.attention_w_enc_encoder = self_attention_w_pos_enc_pure(device=device, num_words=sequence_length_enc, num_dims=dims, out=outs[0], num_heads=num_heads, emb_size=emb_size)\n",
    "    self.attention_1_encoder = self_attention_pure(device=device, p_in=sequence_length_enc, dims=dims, mid=outs[0], num_heads=num_heads)\n",
    "    self.attention_2_encoder = self_attention_pure(device=device, p_in=sequence_length_enc, dims=dims, mid=outs[1], num_heads=num_heads)\n",
    "    self.attention_3_encoder = self_attention_pure(device=device, p_in=sequence_length_enc, dims=dims, mid=outs[2], num_heads=num_heads)\n",
    "    self.dims= dims\n",
    "    self.linear_1 = nn.Linear(seq_length*dims, outs[0]).to(device)\n",
    "    self.linear_2 = nn.Linear(outs[0], outs[1]).to(device)\n",
    "    self.flatten = nn.Flatten(start_dim=1)\n",
    "    self.output = nn.Linear(outs[1], outs[3]).to(device)\n",
    "    self.device=device\n",
    "\n",
    "  def forward(self, x_enc, x_dec, visualization=False):\n",
    "    #Gjør først venstre side av figure 1 for å få query og key\n",
    "    x_enc = x_enc.to(self.device)\n",
    "    x_enc = self.attention_w_enc_encoder(x_enc, visualization=visualization)\n",
    "    x_enc = self.attention_1_encoder(x_enc, visualization=visualization)\n",
    "    x_enc = self.attention_2_encoder(x_enc, visualization=visualization)\n",
    "    x_enc = self.attention_3_encoder(x_enc, visualization=visualization)\n",
    "\n",
    "    x_enc = self.flatten(x_enc)\n",
    "    querykey_for_dec = x_enc\n",
    "    querykey_for_dec = querykey_for_dec.reshape(x_enc.shape[0], -1, self.dims)\n",
    "\n",
    "    x_dec = x_dec.to(self.device)\n",
    "    x_dec = self.attention_w_enc(x_dec, visualization=visualization)\n",
    "    x_dec = self.attention_1(x_dec, querykey_for_dec, visualization=visualization)\n",
    "    x_dec = self.attention_2(x_dec, querykey_for_dec, visualization=visualization)\n",
    "    x_dec = self.attention_3(x_dec, visualization=visualization)\n",
    "\n",
    "    x_dec = self.flatten(x_dec)\n",
    "    x_dec = self.linear_1(x_dec)\n",
    "    x_dec = self.linear_2(x_dec)\n",
    "    x_dec = self.output(x_dec)\n",
    "\n",
    "    return x_dec\n",
    "\n",
    "\n",
    "seq_length = sequence_length\n",
    "self_atten = multi_layer_self_attention_encoder_decoder_scheme(seq_length=seq_length, dims=256, num_heads=8, outs=[768,1024,512,uniques], emb_size=uniques)\n",
    "#p = self_atten(numdata[:12])\n",
    "\n",
    "optimizer = torch.optim.Adam(self_atten.parameters(), lr = 0.0001, weight_decay=0.000005)\n",
    "loss = nn.CrossEntropyLoss()\n",
    "schedule = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=2, cooldown=3, factor=0.333)\n",
    "\n",
    "\n",
    "\n",
    "for k in range(25):\n",
    "  loss_n = 0\n",
    "  acc = 0\n",
    "  for i, batch in enumerate(dataloader):\n",
    "    optimizer.zero_grad()\n",
    "    x, y = batch[0].view((-1, sequence_length_enc)), batch[1].view((-1,sequence_length))\n",
    "    z = torch.zeros((batch[0].shape[0], uniques))\n",
    "    z = z.to(\"cuda:0\")\n",
    "    z[:, batch[2].view((-1, 1))[:]] = 1\n",
    "    \n",
    "    y_p = self_atten(x, y, visualization=False)\n",
    "    l = loss(y_p, z)\n",
    "    l.backward()\n",
    "    optimizer.step()\n",
    "    loss_n += l/batch[0].shape[0]\n",
    "    acc += torch.eq(z.argmax(dim=1), y_p.argmax(dim=1)).sum().item()/batch[0].shape[0]\n",
    "    print(f\"\\rit: {i+1}/{len(dataloader)}, loss: {loss_n/(i+1)}, acc: {acc/(i+1)}\", end=\"\")\n",
    "  loss_n /= len(dataloader)\n",
    "  acc /= len(dataloader)\n",
    "  schedule.step(loss_n)\n",
    "  print(\"epoch:\", k ,\"loss:\", loss_n, \"acc:\",acc)\n",
    "\n",
    "\n",
    "#p = self_atten(numdata[:12])\n",
    "\n",
    "#print(p, words[p.argmax()], data[13])\n",
    "\n",
    "accuracy = 0\n",
    "for i, batch in enumerate(dataloader):\n",
    "  x, y = batch[0].view((-1, sequence_length_enc)), batch[1].view((-1,sequence_length))\n",
    "  z = torch.zeros((batch[2].shape[0], uniques))\n",
    "  z[batch[2].view((-1,1))] = 1\n",
    "  y_p = self_atten(x, y, visualization=False)\n",
    "  y_p = y_p.to(\"cpu\")\n",
    "  #print(y, words[y_p.argmax()])\n",
    "  accuracy += torch.eq(z.argmax(dim=1), y_p.argmax(dim=1)).sum().item()/batch[0].shape[0]\n",
    "\n",
    "accuracy /= len(dataloader)-seq_length-1\n",
    "print(\"accuracy:\", accuracy)\n",
    "\n",
    "print(dataset.tokens_to_text(dataset[128][0]), dataset.tokens_to_text(torch.Tensor([dataset[128][1]])),dataset.tokens_to_text(torch.Tensor([dataset[128][2]])))\n",
    "print(dataset[128][0].shape)\n",
    "x = torch.Tensor(dataset[128][0]).unsqueeze(0)\n",
    "y = torch.Tensor(dataset[128][1]).unsqueeze(0)\n",
    "\n",
    "self_atten(x, y, visualization=True)\n",
    "#x = torch.Tensor(dataset[129][0]).unsqueeze(0)\n",
    "#self_atten(x, visualization=True)"
   ],
   "metadata": {
    "collapsed": false,
    "is_executing": true,
    "ExecuteTime": {
     "start_time": "2023-10-24T14:52:14.340199127Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-10-24T14:33:25.139919255Z",
     "start_time": "2023-10-24T14:33:25.138710881Z"
    }
   }
  }
 ]
}