Refactored_SimplerMambaSSM.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/kaballas/AutoGPT/blob/master/Refactored_SimplerMambaSSM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "boUz4Cj6N1ZM",
        "outputId": "4ee4e1fb-0538-402d-e300-7ba18fce8d5d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting mamba-ssm\n",
            "  Downloading mamba_ssm-1.1.1.tar.gz (34 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting causal-conv1d\n",
            "  Downloading causal_conv1d-1.1.1.tar.gz (6.6 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (2.1.0+cu121)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from mamba-ssm) (23.2)\n",
            "Collecting ninja (from mamba-ssm)\n",
            "  Downloading ninja-1.11.1.1-py2.py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl (307 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m307.2/307.2 kB\u001b[0m \u001b[31m6.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting einops (from mamba-ssm)\n",
            "  Downloading einops-0.7.0-py3-none-any.whl (44 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.6/44.6 kB\u001b[0m \u001b[31m5.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: triton in /usr/local/lib/python3.10/dist-packages (from mamba-ssm) (2.1.0)\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (from mamba-ssm) (4.35.2)\n",
            "Collecting buildtools (from causal-conv1d)\n",
            "  Downloading buildtools-1.0.6.tar.gz (446 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m446.5/446.5 kB\u001b[0m \u001b[31m41.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch) (3.13.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch) (3.1.2)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch) (2023.6.0)\n",
            "Requirement already satisfied: sqlalchemy in /usr/local/lib/python3.10/dist-packages (from buildtools->causal-conv1d) (2.0.24)\n",
            "Collecting argparse (from buildtools->causal-conv1d)\n",
            "  Downloading argparse-1.4.0-py2.py3-none-any.whl (23 kB)\n",
            "Collecting twisted (from buildtools->causal-conv1d)\n",
            "  Downloading twisted-23.10.0-py3-none-any.whl (3.2 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.2/3.2 MB\u001b[0m \u001b[31m64.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting simplejson (from buildtools->causal-conv1d)\n",
            "  Downloading simplejson-3.19.2-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (137 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m137.9/137.9 kB\u001b[0m \u001b[31m20.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting furl (from buildtools->causal-conv1d)\n",
            "  Downloading furl-2.1.3-py2.py3-none-any.whl (20 kB)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from buildtools->causal-conv1d) (2.31.0)\n",
            "Collecting docopt (from buildtools->causal-conv1d)\n",
            "  Downloading docopt-0.6.2.tar.gz (25 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: python-dateutil in /usr/local/lib/python3.10/dist-packages (from buildtools->causal-conv1d) (2.8.2)\n",
            "Collecting redo (from buildtools->causal-conv1d)\n",
            "  Downloading redo-2.0.4-py2.py3-none-any.whl (12 kB)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch) (2.1.3)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch) (1.3.0)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (0.20.2)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (1.23.5)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (6.0.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (2023.6.3)\n",
            "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (0.15.0)\n",
            "Requirement already satisfied: safetensors>=0.3.1 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (0.4.1)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (4.66.1)\n",
            "Requirement already satisfied: six>=1.8.0 in /usr/local/lib/python3.10/dist-packages (from furl->buildtools->causal-conv1d) (1.16.0)\n",
            "Collecting orderedmultidict>=1.0.1 (from furl->buildtools->causal-conv1d)\n",
            "  Downloading orderedmultidict-1.0.1-py2.py3-none-any.whl (11 kB)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->buildtools->causal-conv1d) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->buildtools->causal-conv1d) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->buildtools->causal-conv1d) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->buildtools->causal-conv1d) (2023.11.17)\n",
            "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from sqlalchemy->buildtools->causal-conv1d) (3.0.3)\n",
            "Requirement already satisfied: attrs>=21.3.0 in /usr/local/lib/python3.10/dist-packages (from twisted->buildtools->causal-conv1d) (23.2.0)\n",
            "Collecting automat>=0.8.0 (from twisted->buildtools->causal-conv1d)\n",
            "  Downloading Automat-22.10.0-py2.py3-none-any.whl (26 kB)\n",
            "Collecting constantly>=15.1 (from twisted->buildtools->causal-conv1d)\n",
            "  Downloading constantly-23.10.4-py3-none-any.whl (13 kB)\n",
            "Collecting hyperlink>=17.1.1 (from twisted->buildtools->causal-conv1d)\n",
            "  Downloading hyperlink-21.0.0-py2.py3-none-any.whl (74 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m74.6/74.6 kB\u001b[0m \u001b[31m12.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting incremental>=22.10.0 (from twisted->buildtools->causal-conv1d)\n",
            "  Downloading incremental-22.10.0-py2.py3-none-any.whl (16 kB)\n",
            "Collecting zope-interface>=5 (from twisted->buildtools->causal-conv1d)\n",
            "  Downloading zope.interface-6.1-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (247 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m247.1/247.1 kB\u001b[0m \u001b[31m18.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from zope-interface>=5->twisted->buildtools->causal-conv1d) (67.7.2)\n",
            "Building wheels for collected packages: mamba-ssm, causal-conv1d, buildtools, docopt\n",
            "  Building wheel for mamba-ssm (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for mamba-ssm: filename=mamba_ssm-1.1.1-cp310-cp310-linux_x86_64.whl size=137574741 sha256=998cf941c62596ab13766c8144a1a1b269915f501d95e600b28b85627fdf6dee\n",
            "  Stored in directory: /root/.cache/pip/wheels/22/02/37/2f190644b184e7e3db96b25084fe313637a06764a8c66a9515\n",
            "  Building wheel for causal-conv1d (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for causal-conv1d: filename=causal_conv1d-1.1.1-cp310-cp310-linux_x86_64.whl size=13747478 sha256=21624d55bd83e013a1df9ef514b6690aabbce45c0dd1d1b978a772edbd6c1c85\n",
            "  Stored in directory: /root/.cache/pip/wheels/4f/9e/69/9060a1871f461dfca87b667350f5728d44c555f98dacaf04ff\n",
            "  Building wheel for buildtools (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for buildtools: filename=buildtools-1.0.6-py3-none-any.whl size=512343 sha256=538b019456dd53da8471037b023efc5e8d67bbca9986269cba5e95aad4b59f6d\n",
            "  Stored in directory: /root/.cache/pip/wheels/90/e9/2a/625d99dffa430d0b4293d3d386f63e0eb8edeeb54f3f29d208\n",
            "  Building wheel for docopt (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for docopt: filename=docopt-0.6.2-py2.py3-none-any.whl size=13706 sha256=d13cc64641b2751906c791b119204ca0cec6ec401682d073d37a35ce4d0e0996\n",
            "  Stored in directory: /root/.cache/pip/wheels/fc/ab/d4/5da2067ac95b36618c629a5f93f809425700506f72c9732fac\n",
            "Successfully built mamba-ssm causal-conv1d buildtools docopt\n",
            "Installing collected packages: redo, ninja, incremental, docopt, argparse, zope-interface, simplejson, orderedmultidict, hyperlink, einops, constantly, automat, twisted, furl, buildtools, causal-conv1d, mamba-ssm\n",
            "Successfully installed argparse-1.4.0 automat-22.10.0 buildtools-1.0.6 causal-conv1d-1.1.1 constantly-23.10.4 docopt-0.6.2 einops-0.7.0 furl-2.1.3 hyperlink-21.0.0 incremental-22.10.0 mamba-ssm-1.1.1 ninja-1.11.1.1 orderedmultidict-1.0.1 redo-2.0.4 simplejson-3.19.2 twisted-23.10.0 zope-interface-6.1\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "argparse"
                ]
              }
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "!pip install mamba-ssm causal-conv1d torch"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
      ],
      "metadata": {
        "id": "jIYu0grlX8sD",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b3187e1e-e521-44d0-aba4-f433fcf76b9f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2024-01-12 22:44:59--  https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.111.133, 185.199.109.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1115394 (1.1M) [text/plain]\n",
            "Saving to: ‘input.txt’\n",
            "\n",
            "input.txt           100%[===================>]   1.06M  --.-KB/s    in 0.008s  \n",
            "\n",
            "2024-01-12 22:45:00 (141 MB/s) - ‘input.txt’ saved [1115394/1115394]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir differentattention"
      ],
      "metadata": {
        "id": "b4hR736M5eNT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "from tqdm import tqdm\n",
        "from mamba_ssm import Mamba\n",
        "import time\n",
        "\n",
        "# Hyperparameters and device configuration\n",
        "config = {\n",
        "    \"device\": \"cuda\" if torch.cuda.is_available() else \"cpu\",\n",
        "    \"train_data_fraction\": 0.99,\n",
        "    \"epochs\": 1,\n",
        "    \"lr\": 1e-3,\n",
        "    \"batch_size\": 64,\n",
        "    \"block_size\": 256,\n",
        "    \"max_iters\": 10000,\n",
        "    \"print_iters\": 100,\n",
        "    \"eval_iters\": 10,\n",
        "    \"eval_interval\": 300,\n",
        "    \"n_embed\": 384,\n",
        "    \"n_heads\": 6,\n",
        "    \"n_layers\": 6,\n",
        "    \"dropout\": 0.2,\n",
        "}\n",
        "\n",
        "\n",
        "# Model definition\n",
        "class FeedForward(nn.Module):\n",
        "    \"\"\"\n",
        "    Feed-forward network used in each transformer block.\n",
        "    \"\"\"\n",
        "    def __init__(self, embed_dim, expansion_factor=4):\n",
        "        super(FeedForward, self).__init__()\n",
        "        self.ffn = nn.Sequential(\n",
        "            nn.Linear(embed_dim, expansion_factor * embed_dim),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(expansion_factor * embed_dim, embed_dim),\n",
        "            nn.Dropout(config[\"dropout\"])\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.ffn(x)\n",
        "\n",
        "\n",
        "class Block(nn.Module):\n",
        "    \"\"\"\n",
        "    Basic building block of the transformer consisting of Mamba and feed-forward network.\n",
        "    \"\"\"\n",
        "    def __init__(self, embed_dim, num_heads):\n",
        "        super(Block, self).__init__()\n",
        "        # self.attention = MultiHeadAttention(num_heads, embed_dim)\n",
        "        self.attention = Mamba(\n",
        "            # This module uses roughly 3 * expand * d_model^2 parameters\n",
        "            d_model=embed_dim, # Model dimension d_model\n",
        "            d_state=16,  # SSM state expansion factor\n",
        "            d_conv=4,    # Local convolution width\n",
        "            expand=2,    # Block expansion factor\n",
        "        )\n",
        "        self.feed_forward = FeedForward(embed_dim)\n",
        "        self.norm1 = nn.LayerNorm(embed_dim)\n",
        "        self.norm2 = nn.LayerNorm(embed_dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        attention_output = self.attention(self.norm1(x))\n",
        "        x = x + attention_output\n",
        "        feed_forward_output = self.feed_forward(self.norm2(x))\n",
        "        x = x + feed_forward_output\n",
        "        return x\n",
        "\n",
        "\n",
        "class BigramNeuralNetwork(nn.Module):\n",
        "    \"\"\"\n",
        "    Defines the overall model.\n",
        "    \"\"\"\n",
        "    def __init__(self, vocab_size, embed_dim, num_heads, num_layers):\n",
        "        super(BigramNeuralNetwork, self).__init__()\n",
        "        self.token_embedding = nn.Embedding(vocab_size, embed_dim)\n",
        "        self.position_embedding = nn.Embedding(config[\"block_size\"], embed_dim)\n",
        "        self.blocks = nn.Sequential(*[Block(embed_dim, num_heads) for _ in range(num_layers)])\n",
        "        self.lm_head = nn.Linear(embed_dim, vocab_size)\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        token_embeddings = self.token_embedding(idx)\n",
        "        position_embeddings = self.position_embedding(torch.arange(idx.size(1), device=idx.device))\n",
        "        x = token_embeddings + position_embeddings\n",
        "        x = self.blocks(x)\n",
        "\n",
        "        logits = self.lm_head(x)\n",
        "        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) if targets is not None else None\n",
        "\n",
        "        return logits, loss\n",
        "\n",
        "    def generate(self, idx, max_new_tokens):\n",
        "        for _ in range(max_new_tokens):\n",
        "            idx_cond = idx[:, -config[\"block_size\"]:]\n",
        "            logits, _ = self(idx_cond)\n",
        "            next_token_probs = F.softmax(logits[:, -1, :], dim=-1)\n",
        "            next_token = torch.multinomial(next_token_probs, 1)\n",
        "            idx = torch.cat([idx, next_token], dim=-1)\n",
        "\n",
        "        return idx\n",
        "\n",
        "\n",
        "class TextDataset(Dataset):\n",
        "    \"\"\"\n",
        "    Loads, transforms, and serves the data\n",
        "    \"\"\"\n",
        "    def __init__(self, file_path, block_size):\n",
        "        with open(file_path, \"r\") as f:\n",
        "            text = f.read()\n",
        "        self.chars = sorted(list(set(text)))\n",
        "        self.stoi = {ch: i for i, ch in enumerate(self.chars)}\n",
        "        self.itos = {i: ch for i, ch in enumerate(self.chars)}\n",
        "        self.data = torch.tensor([self.stoi[ch] for ch in text], dtype=torch.long)\n",
        "        self.block_size = block_size\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data) - self.block_size\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        chunk = self.data[idx:idx + self.block_size + 1]\n",
        "        return chunk[:-1], chunk[1:]\n",
        "\n",
        "\n",
        "# Load data\n",
        "dataset = TextDataset(file_path=\"input.txt\",\n",
        "                      block_size=config[\"block_size\"])\n",
        "train_size = int(config[\"train_data_fraction\"] * len(dataset))\n",
        "config[\"vocab_size\"]=len(dataset.chars)\n",
        "print(\"Train size:\", train_size)\n",
        "print(\"Val size:\", len(dataset)-train_size)\n",
        "print(\"Vocab size:\", config[\"vocab_size\"])\n",
        "\n",
        "train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, len(dataset) - train_size])\n",
        "train_loader = DataLoader(train_dataset, batch_size=config[\"batch_size\"], shuffle=True)\n",
        "val_loader = DataLoader(val_dataset, batch_size=config[\"batch_size\"], shuffle=True)\n",
        "\n",
        "\n",
        "# Initialize Mamba model\n",
        "mamba_model = BigramNeuralNetwork(\n",
        "    vocab_size=config[\"vocab_size\"],\n",
        "    embed_dim=config[\"n_embed\"],\n",
        "    num_heads=config[\"n_heads\"],\n",
        "    num_layers=config[\"n_layers\"]\n",
        ").to(config[\"device\"])\n",
        "\n",
        "optimizer = torch.optim.AdamW(mamba_model.parameters(), lr=config[\"lr\"])\n",
        "\n",
        "\n",
        "def train_model(model, train_loader, val_loader, optimizer, config, early_stop=None):\n",
        "    print(\"\\nStarting train\")\n",
        "    for epoch in range(config[\"epochs\"]):\n",
        "        # Training Phase\n",
        "        start_time = time.time()\n",
        "        model.train()\n",
        "\n",
        "        for i, (x_batch, y_batch) in enumerate(train_loader):\n",
        "            x_batch, y_batch = x_batch.to(config[\"device\"]), y_batch.to(config[\"device\"])\n",
        "\n",
        "            optimizer.zero_grad()  # Clear gradients from the previous step\n",
        "            _, train_loss = model(x_batch, y_batch)  # Forward pass and loss computation\n",
        "            train_loss.backward()  # Backpropagation\n",
        "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # Gradient clipping\n",
        "            optimizer.step()  # Update model parameters\n",
        "\n",
        "            if i % config[\"print_iters\"] == 0:\n",
        "                # Validation Phase\n",
        "                model.eval()  # Set the model to evaluation mode\n",
        "                total_val_loss = 0\n",
        "                with torch.no_grad():  # Disable gradient computation\n",
        "                    for x_batch, y_batch in val_loader:\n",
        "                        x_batch, y_batch = x_batch.to(config[\"device\"]), y_batch.to(config[\"device\"])\n",
        "                        _, val_loss = model(x_batch, y_batch)\n",
        "                        total_val_loss += val_loss.item()\n",
        "\n",
        "                average_val_loss = total_val_loss / len(val_loader)\n",
        "                print(f\"Epoch [{epoch+1}/{config['epochs']}], Step [{i+1}/{len(train_loader)}], Train Loss: {train_loss.item():.4f}, Validation Loss: {average_val_loss:.4f}\")\n",
        "                model.train()\n",
        "\n",
        "                # Optional: Save Model Checkpoint\n",
        "                # torch.save(model.state_dict(), f\"./model_checkpoint_epoch_{epoch+1}.pt\")\n",
        "\n",
        "            if i == early_stop:\n",
        "                break # Early stop\n",
        "\n",
        "    end_time = time.time()\n",
        "    train_duration = end_time - start_time\n",
        "    print(f\"Training completed in {train_duration:.2f} seconds\")\n",
        "\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4Lxp6vSFW5WH",
        "outputId": "50d063cb-d18b-4b9c-abf0-536b9abe66e3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train size: 1103986\n",
            "Val size: 11152\n",
            "Vocab size: 65\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_model(\n",
        "    model=mamba_model,\n",
        "    train_loader=train_loader,\n",
        "    val_loader=val_loader,\n",
        "    optimizer=optimizer,\n",
        "    config=config,\n",
        "    early_stop=500\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GZnJ8ZYM00SP",
        "outputId": "879f1d03-b17b-49d9-cee0-adabe6513ae7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Starting train\n",
            "Epoch [1/1], Step [1/17250], Train Loss: 4.7023, Validation Loss: 9.3851\n",
            "Epoch [1/1], Step [101/17250], Train Loss: 1.5919, Validation Loss: 1.5709\n",
            "Epoch [1/1], Step [201/17250], Train Loss: 1.3830, Validation Loss: 1.3720\n",
            "Epoch [1/1], Step [301/17250], Train Loss: 1.2471, Validation Loss: 1.2713\n",
            "Epoch [1/1], Step [401/17250], Train Loss: 1.2207, Validation Loss: 1.2019\n",
            "Epoch [1/1], Step [501/17250], Train Loss: 1.1363, Validation Loss: 1.1431\n",
            "Training completed in 393.98 seconds\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class SelfAttentionHead(nn.Module):\n",
        "    \"\"\"\n",
        "    Implements a self-attention mechanism.\n",
        "    \"\"\"\n",
        "    def __init__(self, embed_dim, head_size):\n",
        "        super(SelfAttentionHead, self).__init__()\n",
        "        self.key = nn.Linear(embed_dim, head_size)\n",
        "        self.query = nn.Linear(embed_dim, head_size)\n",
        "        self.value = nn.Linear(embed_dim, head_size)\n",
        "        self.scale = head_size ** -0.5\n",
        "        self.dropout = nn.Dropout(config[\"dropout\"])\n",
        "\n",
        "    def forward(self, x):\n",
        "        batch_size, seq_length, _ = x.shape\n",
        "        keys = self.key(x)\n",
        "        queries = self.query(x)\n",
        "        values = self.value(x)\n",
        "\n",
        "        weights = (keys @ queries.transpose(-2, -1)) * self.scale\n",
        "        weights = torch.softmax(weights, dim=-1)\n",
        "        weights = self.dropout(weights)\n",
        "\n",
        "        output = weights @ values\n",
        "        return output\n",
        "\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    \"\"\"\n",
        "    Multi-head attention mechanism.\n",
        "    \"\"\"\n",
        "    def __init__(self, num_heads, embed_dim):\n",
        "        super(MultiHeadAttention, self).__init__()\n",
        "        head_size = embed_dim // num_heads\n",
        "        self.heads = nn.ModuleList([SelfAttentionHead(embed_dim, head_size) for _ in range(num_heads)])\n",
        "        self.linear = nn.Linear(embed_dim, embed_dim)\n",
        "        self.dropout = nn.Dropout(config[\"dropout\"])\n",
        "\n",
        "    def forward(self, x):\n",
        "        concatenated_heads = torch.cat([head(x) for head in self.heads], dim=-1)\n",
        "        output = self.linear(concatenated_heads)\n",
        "        output = self.dropout(output)\n",
        "        return output\n",
        "\n",
        "class Block(nn.Module):\n",
        "    \"\"\"\n",
        "    Redo the Block class to use attention instead of Mamba\n",
        "    \"\"\"\n",
        "    def __init__(self, embed_dim, num_heads):\n",
        "        super(Block, self).__init__()\n",
        "        self.attention = MultiHeadAttention(num_heads, embed_dim)\n",
        "        #self.attention = Mamba(\n",
        "        #    # This module uses roughly 3 * expand * d_model^2 parameters\n",
        "        #    d_model=embed_dim, # Model dimension d_model\n",
        "        #    d_state=16,  # SSM state expansion factor\n",
        "        #    d_conv=4,    # Local convolution width\n",
        "        #    expand=2,    # Block expansion factor\n",
        "        #)\n",
        "        self.feed_forward = FeedForward(embed_dim)\n",
        "        self.norm1 = nn.LayerNorm(embed_dim)\n",
        "        self.norm2 = nn.LayerNorm(embed_dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        attention_output = self.attention(self.norm1(x))\n",
        "        x = x + attention_output\n",
        "        feed_forward_output = self.feed_forward(self.norm2(x))\n",
        "        x = x + feed_forward_output\n",
        "        return x\n",
        "\n",
        "\n",
        "# Initialize Transformer model\n",
        "transformer_model = BigramNeuralNetwork(\n",
        "    vocab_size=len(dataset.chars),\n",
        "    embed_dim=config[\"n_embed\"],\n",
        "    num_heads=config[\"n_heads\"],\n",
        "    num_layers=config[\"n_layers\"]\n",
        ").to(config[\"device\"])\n",
        "\n",
        "optimizer = torch.optim.AdamW(transformer_model.parameters(), lr=config[\"lr\"])\n"
      ],
      "metadata": {
        "id": "m6pczvm9Xr80"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_model(\n",
        "    model=transformer_model,\n",
        "    train_loader=train_loader,\n",
        "    val_loader=val_loader,\n",
        "    optimizer=optimizer,\n",
        "    config=config,\n",
        "    early_stop=500\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UW6jvstL04Mo",
        "outputId": "156addfe-1c18-4865-a22a-f8219921458a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Starting train\n",
            "Epoch [1/1], Step [1/17250], Train Loss: 4.6324, Validation Loss: 9.3575\n",
            "Epoch [1/1], Step [101/17250], Train Loss: 0.2796, Validation Loss: 0.0974\n",
            "Epoch [1/1], Step [201/17250], Train Loss: 0.0187, Validation Loss: 0.0107\n",
            "Epoch [1/1], Step [301/17250], Train Loss: 0.0144, Validation Loss: 0.0102\n",
            "Epoch [1/1], Step [401/17250], Train Loss: 0.0125, Validation Loss: 0.0102\n",
            "Epoch [1/1], Step [501/17250], Train Loss: 0.0106, Validation Loss: 0.0102\n",
            "Training completed in 390.31 seconds\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "import math\n",
        "from functools import partial\n",
        "import json\n",
        "import os\n",
        "\n",
        "from collections import namedtuple\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "from mamba_ssm.models.config_mamba import MambaConfig\n",
        "from mamba_ssm.modules.mamba_simple import Mamba, Block\n",
        "from mamba_ssm.utils.generation import GenerationMixin\n",
        "from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf\n",
        "\n",
        "try:\n",
        "    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn\n",
        "except ImportError:\n",
        "    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None\n",
        "\n",
        "\n",
        "def create_block(\n",
        "    d_model,\n",
        "    ssm_cfg=None,\n",
        "    norm_epsilon=1e-5,\n",
        "    rms_norm=False,\n",
        "    residual_in_fp32=False,\n",
        "    fused_add_norm=False,\n",
        "    layer_idx=None,\n",
        "    device=None,\n",
        "    dtype=None,\n",
        "):\n",
        "    if ssm_cfg is None:\n",
        "        ssm_cfg = {}\n",
        "    factory_kwargs = {\"device\": device, \"dtype\": dtype}\n",
        "    mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)\n",
        "    norm_cls = partial(\n",
        "        nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs\n",
        "    )\n",
        "    block = Block(\n",
        "        d_model,\n",
        "        mixer_cls,\n",
        "        norm_cls=norm_cls,\n",
        "        fused_add_norm=fused_add_norm,\n",
        "        residual_in_fp32=residual_in_fp32,\n",
        "    )\n",
        "    block.layer_idx = layer_idx\n",
        "    return block\n",
        "\n",
        "\n",
        "# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454\n",
        "def _init_weights(\n",
        "    module,\n",
        "    n_layer,\n",
        "    initializer_range=0.02,  # Now only used for embedding layer.\n",
        "    rescale_prenorm_residual=True,\n",
        "    n_residuals_per_layer=1,  # Change to 2 if we have MLP\n",
        "):\n",
        "    if isinstance(module, nn.Linear):\n",
        "        if module.bias is not None:\n",
        "            if not getattr(module.bias, \"_no_reinit\", False):\n",
        "                nn.init.zeros_(module.bias)\n",
        "    elif isinstance(module, nn.Embedding):\n",
        "        nn.init.normal_(module.weight, std=initializer_range)\n",
        "\n",
        "    if rescale_prenorm_residual:\n",
        "        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:\n",
        "        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale\n",
        "        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.\n",
        "        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/\n",
        "        #\n",
        "        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py\n",
        "        for name, p in module.named_parameters():\n",
        "            if name in [\"out_proj.weight\", \"fc2.weight\"]:\n",
        "                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block\n",
        "                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)\n",
        "                # We need to reinit p since this code could be called multiple times\n",
        "                # Having just p *= scale would repeatedly scale it down\n",
        "                nn.init.kaiming_uniform_(p, a=math.sqrt(5))\n",
        "                with torch.no_grad():\n",
        "                    p /= math.sqrt(n_residuals_per_layer * n_layer)\n",
        "\n",
        "\n",
        "class MixerModel(nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        d_model: int,\n",
        "        n_layer: int,\n",
        "        vocab_size: int,\n",
        "        ssm_cfg=None,\n",
        "        norm_epsilon: float = 1e-5,\n",
        "        rms_norm: bool = False,\n",
        "        initializer_cfg=None,\n",
        "        fused_add_norm=False,\n",
        "        residual_in_fp32=False,\n",
        "        device=None,\n",
        "        dtype=None,\n",
        "    ) -> None:\n",
        "        factory_kwargs = {\"device\": device, \"dtype\": dtype}\n",
        "        super().__init__()\n",
        "        self.residual_in_fp32 = residual_in_fp32\n",
        "\n",
        "        self.embedding = nn.Embedding(vocab_size, d_model, **factory_kwargs)\n",
        "        self.dropout = nn.Dropout(p=0.2)\n",
        "\n",
        "        # We change the order of residual and layer norm:\n",
        "        # Instead of LN -> Attn / MLP -> Add, we do:\n",
        "        # Add -> LN -> Attn / MLP / Mixer, returning both the residual branch (output of Add) and\n",
        "        # the main branch (output of MLP / Mixer). The model definition is unchanged.\n",
        "        # This is for performance reason: we can fuse add + layer_norm.\n",
        "        self.fused_add_norm = fused_add_norm\n",
        "        if self.fused_add_norm:\n",
        "            if layer_norm_fn is None or rms_norm_fn is None:\n",
        "                raise ImportError(\"Failed to import Triton LayerNorm / RMSNorm kernels\")\n",
        "\n",
        "        self.layers = nn.ModuleList(\n",
        "            [\n",
        "                create_block(\n",
        "                    d_model,\n",
        "                    ssm_cfg=ssm_cfg,\n",
        "                    norm_epsilon=norm_epsilon,\n",
        "                    rms_norm=rms_norm,\n",
        "                    residual_in_fp32=residual_in_fp32,\n",
        "                    fused_add_norm=fused_add_norm,\n",
        "                    layer_idx=i,\n",
        "                    **factory_kwargs,\n",
        "                )\n",
        "                for i in range(n_layer)\n",
        "            ]\n",
        "        )\n",
        "\n",
        "        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(\n",
        "            d_model, eps=norm_epsilon, **factory_kwargs\n",
        "        )\n",
        "\n",
        "        self.apply(\n",
        "            partial(\n",
        "                _init_weights,\n",
        "                n_layer=n_layer,\n",
        "                **(initializer_cfg if initializer_cfg is not None else {}),\n",
        "            )\n",
        "        )\n",
        "\n",
        "    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):\n",
        "        return {\n",
        "            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)\n",
        "            for i, layer in enumerate(self.layers)\n",
        "        }\n",
        "\n",
        "    def forward(self, input_ids, inference_params=None):\n",
        "        hidden_states = self.embedding(input_ids)\n",
        "        hidden_states = self.dropout(hidden_states)\n",
        "        residual = None\n",
        "        for layer in self.layers:\n",
        "            hidden_states, residual = layer(\n",
        "                hidden_states, residual, inference_params=inference_params\n",
        "            )\n",
        "        hidden_states = self.dropout(hidden_states)\n",
        "        if not self.fused_add_norm:\n",
        "            residual = (hidden_states + residual) if residual is not None else hidden_states\n",
        "            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))\n",
        "        else:\n",
        "            # Set prenorm=False here since we don't need the residual\n",
        "            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn\n",
        "            hidden_states = fused_add_norm_fn(\n",
        "                hidden_states,\n",
        "                self.norm_f.weight,\n",
        "                self.norm_f.bias,\n",
        "                eps=self.norm_f.eps,\n",
        "                residual=residual,\n",
        "                prenorm=False,\n",
        "                residual_in_fp32=self.residual_in_fp32,\n",
        "            )\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class MambaClassifier(nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        config: MambaConfig,\n",
        "        num_classes: int,\n",
        "        initializer_cfg=None,\n",
        "        device=None,\n",
        "        dtype=None,\n",
        "    ):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "        self.num_classes = num_classes\n",
        "\n",
        "        # Backbone model (MixerModel) setup\n",
        "        self.backbone = MixerModel(\n",
        "            d_model=config.d_model,\n",
        "            n_layer=config.n_layer,\n",
        "            vocab_size=config.vocab_size,\n",
        "            ssm_cfg=config.ssm_cfg,\n",
        "            rms_norm=config.rms_norm,\n",
        "            initializer_cfg=initializer_cfg,\n",
        "            fused_add_norm=config.fused_add_norm,\n",
        "            residual_in_fp32=config.residual_in_fp32,\n",
        "            device=device,\n",
        "            dtype=dtype,\n",
        "        )\n",
        "\n",
        "        # Output layer for multi-label classification\n",
        "        self.classifier = nn.Linear(config.d_model, num_classes, **{\"device\": device, \"dtype\": dtype})\n",
        "        #self.loss_fn = nn.BCEWithLogitsLoss()\n",
        "        self.loss_fn = nn.CrossEntropyLoss()\n",
        "        self.dropout = nn.Dropout(p=0.2)\n",
        "\n",
        "        # Initialize weights\n",
        "        self.apply(\n",
        "            partial(\n",
        "                _init_weights,\n",
        "                n_layer=config.n_layer,\n",
        "                **(initializer_cfg if initializer_cfg is not None else {}),\n",
        "            )\n",
        "        )\n",
        "\n",
        "    def forward(self, input_ids, labels=None):\n",
        "        # Passing the input through the backbone model\n",
        "        hidden_states = self.backbone(input_ids)\n",
        "        last_hidden_state = self.dropout(hidden_states[:, -1, :])\n",
        "\n",
        "        # Applying the classifier to get logits for each label\n",
        "        logits = self.classifier(last_hidden_state)  # Use the last hidden state\n",
        "\n",
        "        # Calculate loss if labels are provided\n",
        "        loss = self.loss_fn(logits, labels) if labels is not None else None\n",
        "\n",
        "        return (loss, logits) if labels is not None else loss\n",
        "\n"
      ],
      "metadata": {
        "id": "np6USQ2HyIL3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Usage of the class\n",
        "my_mamba_config = MambaConfig(d_model=config['n_embed'], n_layer=config['n_layers'], vocab_size=config['vocab_size'], ssm_cfg={})\n",
        "mamba_model_v2 = MambaClassifier(my_mamba_config, num_classes=config['vocab_size'])\n",
        "input_ids = train_dataset[0][0]\n",
        "logits = mamba_model_v2(input_ids)"
      ],
      "metadata": {
        "id": "AtIIDtpd1CQB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_model(\n",
        "    model=mamba_model_v2,\n",
        "    train_loader=train_loader,\n",
        "    val_loader=val_loader,\n",
        "    optimizer=optimizer,\n",
        "    config=config,\n",
        "    early_stop=500\n",
        ")"
      ],
      "metadata": {
        "id": "A2OzsRxx7MN7"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}