HyenaDNA_training_&_inference_example_(Public).ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/martinaryee/hyena-dna/blob/main/HyenaDNA_training_%26_inference_example_(Public).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Gkte68aIhTy"
      },
      "source": [
        "Welcome!\n",
        "\n",
        "So glad you made it and thanks for checking out HyenaDNA! We're \\**very*\\* excited to get this in the hands of researchers.\n",
        "\n",
        "Let's walk through how to setup a simple HyenaDNA model to:\n",
        "\n",
        "*    Train on the GenomicBenchmarks dataset from scratch or fine-tune (sequence level classification tasks)\n",
        "*    Load from HuggingFace :)\n",
        "*    Do inference on up to 1M tokens!\n",
        "\n",
        "--\n",
        "\n",
        "After this tutorial, you can check out more advanced training (more configurable and optimized for speed) through the HyenaDNA repo, which uses Pytorch Lightning and Hydra for managing all the experimental configs and settings.\n",
        "\n",
        "Luckily, the code for classification here is 95% the same as the pretraining code!\n",
        "\n",
        "- [arxiv](https://arxiv.org/abs/2306.15794)  \n",
        "- [blog](https://hazyresearch.stanford.edu/blog/2023-06-29-hyena-dna)\n",
        "- [github](https://github.com/HazyResearch/hyena-dna)  \n",
        "- [huggingface](https://huggingface.co/LongSafari) (pretrained weights)  \n",
        "\n",
        "Good luck and enjoy!!!\n",
        "\n",
        "\\\\\n",
        "Credit: much of the code is forked and extended from [S4](https://github.com/HazyResearch/state-spaces) and [Safari](https://github.com/HazyResearch/safari)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "JWUc70H7Iiwd",
        "outputId": "d415de3a-3463-40a9-f21a-55ce384490ba"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting einops\n",
            "  Downloading einops-0.6.1-py3-none-any.whl (42 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.2/42.2 kB\u001b[0m \u001b[31m2.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: einops\n",
            "Successfully installed einops-0.6.1\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.10/dist-packages (0.15.2+cu118)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from torchvision) (1.22.4)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from torchvision) (2.27.1)\n",
            "Requirement already satisfied: torch==2.0.1 in /usr/local/lib/python3.10/dist-packages (from torchvision) (2.0.1+cu118)\n",
            "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.10/dist-packages (from torchvision) (8.4.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch==2.0.1->torchvision) (3.12.2)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch==2.0.1->torchvision) (4.6.3)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch==2.0.1->torchvision) (1.11.1)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch==2.0.1->torchvision) (3.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch==2.0.1->torchvision) (3.1.2)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch==2.0.1->torchvision) (2.0.0)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch==2.0.1->torchvision) (3.25.2)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch==2.0.1->torchvision) (16.0.6)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision) (1.26.16)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision) (2023.5.7)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision) (3.4)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch==2.0.1->torchvision) (2.1.3)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch==2.0.1->torchvision) (1.3.0)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting transformers\n",
            "  Downloading transformers-4.30.2-py3-none-any.whl (7.2 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.2/7.2 MB\u001b[0m \u001b[31m45.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.12.2)\n",
            "Collecting huggingface-hub<1.0,>=0.14.1 (from transformers)\n",
            "  Downloading huggingface_hub-0.15.1-py3-none-any.whl (236 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m236.8/236.8 kB\u001b[0m \u001b[31m19.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.22.4)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (23.1)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2022.10.31)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.27.1)\n",
            "Collecting tokenizers!=0.11.3,<0.14,>=0.11.1 (from transformers)\n",
            "  Downloading tokenizers-0.13.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (7.8 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m57.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting safetensors>=0.3.1 (from transformers)\n",
            "  Downloading safetensors-0.3.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m32.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.65.0)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.14.1->transformers) (2023.6.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.14.1->transformers) (4.6.3)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (1.26.16)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2023.5.7)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.4)\n",
            "Installing collected packages: tokenizers, safetensors, huggingface-hub, transformers\n",
            "Successfully installed huggingface-hub-0.15.1 safetensors-0.3.1 tokenizers-0.13.3 transformers-4.30.2\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting genomic-benchmarks\n",
            "  Downloading genomic_benchmarks-0.0.9.tar.gz (21 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting biopython>=1.79 (from genomic-benchmarks)\n",
            "  Downloading biopython-1.81-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.1 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.1/3.1 MB\u001b[0m \u001b[31m18.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: requests>=2.23.0 in /usr/local/lib/python3.10/dist-packages (from genomic-benchmarks) (2.27.1)\n",
            "Requirement already satisfied: pip>=20.0.1 in /usr/local/lib/python3.10/dist-packages (from genomic-benchmarks) (23.1.2)\n",
            "Requirement already satisfied: numpy>=1.17.0 in /usr/local/lib/python3.10/dist-packages (from genomic-benchmarks) (1.22.4)\n",
            "Requirement already satisfied: pandas>=1.1.4 in /usr/local/lib/python3.10/dist-packages (from genomic-benchmarks) (1.5.3)\n",
            "Requirement already satisfied: tqdm>=4.41.1 in /usr/local/lib/python3.10/dist-packages (from genomic-benchmarks) (4.65.0)\n",
            "Requirement already satisfied: pyyaml>=5.3.1 in /usr/local/lib/python3.10/dist-packages (from genomic-benchmarks) (6.0)\n",
            "Requirement already satisfied: gdown>=4.2.0 in /usr/local/lib/python3.10/dist-packages (from genomic-benchmarks) (4.6.6)\n",
            "Requirement already satisfied: yarl in /usr/local/lib/python3.10/dist-packages (from genomic-benchmarks) (1.9.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from gdown>=4.2.0->genomic-benchmarks) (3.12.2)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.10/dist-packages (from gdown>=4.2.0->genomic-benchmarks) (1.16.0)\n",
            "Requirement already satisfied: beautifulsoup4 in /usr/local/lib/python3.10/dist-packages (from gdown>=4.2.0->genomic-benchmarks) (4.11.2)\n",
            "Requirement already satisfied: python-dateutil>=2.8.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.1.4->genomic-benchmarks) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.1.4->genomic-benchmarks) (2022.7.1)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.23.0->genomic-benchmarks) (1.26.16)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.23.0->genomic-benchmarks) (2023.5.7)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests>=2.23.0->genomic-benchmarks) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.23.0->genomic-benchmarks) (3.4)\n",
            "Requirement already satisfied: multidict>=4.0 in /usr/local/lib/python3.10/dist-packages (from yarl->genomic-benchmarks) (6.0.4)\n",
            "Requirement already satisfied: soupsieve>1.2 in /usr/local/lib/python3.10/dist-packages (from beautifulsoup4->gdown>=4.2.0->genomic-benchmarks) (2.4.1)\n",
            "Requirement already satisfied: PySocks!=1.5.7,>=1.5.6 in /usr/local/lib/python3.10/dist-packages (from requests>=2.23.0->genomic-benchmarks) (1.7.1)\n",
            "Building wheels for collected packages: genomic-benchmarks\n",
            "  Building wheel for genomic-benchmarks (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for genomic-benchmarks: filename=genomic_benchmarks-0.0.9-py3-none-any.whl size=22506 sha256=9d8b2ccd8e2c0102afd660f31ab821452fa1a78d738c7dc9bb480f7bc85fdac8\n",
            "  Stored in directory: /root/.cache/pip/wheels/86/78/72/de9d75be800b6d679b8814b9663b3f610c34188b443b8b9b77\n",
            "Successfully built genomic-benchmarks\n",
            "Installing collected packages: biopython, genomic-benchmarks\n",
            "Successfully installed biopython-1.81 genomic-benchmarks-0.0.9\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting OmegaConf\n",
            "  Downloading omegaconf-2.3.0-py3-none-any.whl (79 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m79.5/79.5 kB\u001b[0m \u001b[31m4.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting antlr4-python3-runtime==4.9.* (from OmegaConf)\n",
            "  Downloading antlr4-python3-runtime-4.9.3.tar.gz (117 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m117.0/117.0 kB\u001b[0m \u001b[31m10.9 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: PyYAML>=5.1.0 in /usr/local/lib/python3.10/dist-packages (from OmegaConf) (6.0)\n",
            "Building wheels for collected packages: antlr4-python3-runtime\n",
            "  Building wheel for antlr4-python3-runtime (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for antlr4-python3-runtime: filename=antlr4_python3_runtime-4.9.3-py3-none-any.whl size=144554 sha256=cf564d6d9233b423f31ffb0813dd242c50ef64d3e920d989573c98de3668ae3c\n",
            "  Stored in directory: /root/.cache/pip/wheels/12/93/dd/1f6a127edc45659556564c5730f6d4e300888f4bca2d4c5a88\n",
            "Successfully built antlr4-python3-runtime\n",
            "Installing collected packages: antlr4-python3-runtime, OmegaConf\n",
            "Successfully installed OmegaConf-2.3.0 antlr4-python3-runtime-4.9.3\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "pydevd_plugins"
                ]
              }
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "#@title Installs\n",
        "!pip install einops\n",
        "!pip install torchvision\n",
        "!pip install transformers==4.26.1\n",
        "!pip install genomic-benchmarks\n",
        "!pip install OmegaConf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-WIH4-WEI4in"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "# for HyenaDNA specifically\n",
        "import torch\n",
        "import math\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from functools import partial\n",
        "from einops import rearrange\n",
        "from typing import Optional\n",
        "from functools import partial\n",
        "from torch import Tensor\n",
        "from torchvision.ops import StochasticDepth\n",
        "from collections import namedtuple"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nu3UYdFqKzTs"
      },
      "source": [
        "# HyenaDNA\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "pWJZc3volbL6"
      },
      "outputs": [],
      "source": [
        "#@title Hyena layer\n",
        "\n",
        "\n",
        "def fftconv(u, k, D):\n",
        "    \"\"\"\n",
        "    We apply a convolution through the fourier domain (from the Convolution Theorem)\n",
        "\n",
        "    \"\"\"\n",
        "    seqlen = u.shape[-1]\n",
        "    fft_size = 2 * seqlen\n",
        "\n",
        "    k_f = torch.fft.rfft(k, n=fft_size) / fft_size\n",
        "    u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size)\n",
        "\n",
        "    if len(u.shape) > 3: k_f = k_f.unsqueeze(1)\n",
        "    y = torch.fft.irfft(u_f * k_f, n=fft_size, norm='forward')[..., :seqlen]\n",
        "\n",
        "    out = y + u * D.unsqueeze(-1)\n",
        "    return out.to(dtype=u.dtype)\n",
        "\n",
        "\n",
        "@torch.jit.script\n",
        "def mul_sum(q, y):\n",
        "    return (q * y).sum(dim=1)\n",
        "\n",
        "class OptimModule(nn.Module):\n",
        "    \"\"\" Interface for Module that allows registering buffers/parameters with configurable optimizer hyperparameters \"\"\"\n",
        "\n",
        "    def register(self, name, tensor, lr=None, wd=0.0):\n",
        "        \"\"\"Register a tensor with a configurable learning rate and 0 weight decay\"\"\"\n",
        "\n",
        "        if lr == 0.0:\n",
        "            self.register_buffer(name, tensor)\n",
        "        else:\n",
        "            self.register_parameter(name, nn.Parameter(tensor))\n",
        "\n",
        "            optim = {}\n",
        "            if lr is not None: optim[\"lr\"] = lr\n",
        "            if wd is not None: optim[\"weight_decay\"] = wd\n",
        "            setattr(getattr(self, name), \"_optim\", optim)\n",
        "\n",
        "\n",
        "class Sin(nn.Module):\n",
        "    \"\"\"The Sin activation function for the Hyena Filter function.\"\"\"\n",
        "    def __init__(self, dim, w=10, train_freq=True):\n",
        "        super().__init__()\n",
        "        self.freq = nn.Parameter(w * torch.ones(1, dim)) if train_freq else w * torch.ones(1, dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return torch.sin(self.freq * x)\n",
        "\n",
        "\n",
        "class PositionalEmbedding(OptimModule):\n",
        "    def __init__(self, emb_dim: int, seq_len: int, lr_pos_emb: float=1e-5, **kwargs):\n",
        "        \"\"\"Complex exponential positional embeddings for Hyena filters.\"\"\"\n",
        "        super().__init__()\n",
        "\n",
        "        self.seq_len = seq_len\n",
        "        # The time embedding fed to the filteres is normalized so that t_f = 1\n",
        "        t = torch.linspace(0, 1, self.seq_len)[None, :, None] # 1, L, 1\n",
        "\n",
        "        if emb_dim > 1:\n",
        "            bands = (emb_dim - 1) // 2\n",
        "        # To compute the right embeddings we use the \"proper\" linspace\n",
        "        t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None]\n",
        "        w = 2 * math.pi * t_rescaled / seq_len # 1, L, 1\n",
        "\n",
        "        f = torch.linspace(1e-4, bands - 1, bands)[None, None]\n",
        "        z = torch.exp(-1j * f * w)\n",
        "        z = torch.cat([t, z.real, z.imag], dim=-1)\n",
        "        self.register(\"z\", z, lr=lr_pos_emb)\n",
        "        self.register(\"t\", t, lr=0.0)\n",
        "\n",
        "    def forward(self, L):\n",
        "        return self.z[:, :L], self.t[:, :L]\n",
        "\n",
        "\n",
        "class ExponentialModulation(OptimModule):\n",
        "    \"\"\"The window function applied to the output of the (MLP) filter function.\"\"\"\n",
        "    def __init__(\n",
        "        self,\n",
        "        d_model,\n",
        "        fast_decay_pct=0.3,\n",
        "        slow_decay_pct=1.5,\n",
        "        target=1e-2,\n",
        "        modulation_lr=0.0,\n",
        "        modulate: bool=True,\n",
        "        shift: float = 0.05,\n",
        "        **kwargs\n",
        "    ):\n",
        "        super().__init__()\n",
        "        self.modulate = modulate\n",
        "        self.shift = shift\n",
        "        max_decay = math.log(target) / fast_decay_pct\n",
        "        min_decay = math.log(target) / slow_decay_pct\n",
        "        deltas = torch.linspace(min_decay, max_decay, d_model)[None, None]\n",
        "        self.register(\"deltas\", deltas, lr=modulation_lr)\n",
        "\n",
        "    def forward(self, t, x):\n",
        "        if self.modulate:\n",
        "            decay = torch.exp(-t * self.deltas.abs())\n",
        "            x = x * (decay + self.shift)\n",
        "        return x\n",
        "\n",
        "\n",
        "class HyenaFilter(OptimModule):\n",
        "    def __init__(\n",
        "            self,\n",
        "            d_model,\n",
        "            emb_dim=3, # dim of input to MLP, augments with positional encoding\n",
        "            order=16, # width of the implicit MLP\n",
        "            fused_fft_conv=False,\n",
        "            seq_len=1024,\n",
        "            lr=1e-3,\n",
        "            lr_pos_emb=1e-5,\n",
        "            dropout=0.0,\n",
        "            w=1, # frequency of periodic activations\n",
        "            wd=0, # weight decay of kernel parameters\n",
        "            bias=True,\n",
        "            num_inner_mlps=2,\n",
        "            normalized=False,\n",
        "            **kwargs\n",
        "        ):\n",
        "        \"\"\"\n",
        "        Implicit long filter with modulation.\n",
        "\n",
        "        Args:\n",
        "            d_model: number of channels in the input\n",
        "            emb_dim: dimension of the positional encoding (`emb_dim` - 1) // 2 is the number of bands\n",
        "            order: width of the FFN\n",
        "            num_inner_mlps: number of inner linear layers inside filter MLP\n",
        "\n",
        "        Note:\n",
        "            filter_dropout is not implemented\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "\n",
        "        self.d_model = d_model\n",
        "        self.use_bias = bias\n",
        "        self.fused_fft_conv = fused_fft_conv\n",
        "        self.bias = nn.Parameter(torch.randn(self.d_model))\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "        act = Sin(dim=order, w=w)\n",
        "        self.emb_dim = emb_dim\n",
        "        assert emb_dim % 2 != 0 and emb_dim >= 3, \"emb_dim must be odd and greater or equal to 3 (time, sine and cosine)\"\n",
        "        self.seq_len = seq_len\n",
        "\n",
        "        self.pos_emb = PositionalEmbedding(emb_dim, seq_len, lr_pos_emb)\n",
        "\n",
        "        self.implicit_filter = nn.Sequential(\n",
        "            nn.Linear(emb_dim, order),\n",
        "            act,\n",
        "        )\n",
        "        for i in range(num_inner_mlps):\n",
        "            self.implicit_filter.append(nn.Linear(order, order))\n",
        "            self.implicit_filter.append(act)\n",
        "\n",
        "        self.implicit_filter.append(nn.Linear(order, d_model, bias=False))\n",
        "\n",
        "        self.modulation = ExponentialModulation(d_model, **kwargs)\n",
        "\n",
        "        self.normalized = normalized\n",
        "        for c in self.implicit_filter.children():\n",
        "            for name, v in c.state_dict().items():\n",
        "                optim = {\"weight_decay\": wd, \"lr\": lr}\n",
        "                setattr(getattr(c, name), \"_optim\", optim)\n",
        "\n",
        "    def filter(self, L, *args, **kwargs):\n",
        "        z, t = self.pos_emb(L)\n",
        "        h = self.implicit_filter(z)\n",
        "        h = self.modulation(t, h)\n",
        "        return h\n",
        "\n",
        "    def forward(self, x, L, k=None, bias=None, *args, **kwargs):\n",
        "        if k is None: k = self.filter(L)\n",
        "\n",
        "        # Ensure compatibility with filters that return a tuple\n",
        "        k = k[0] if type(k) is tuple else k\n",
        "\n",
        "        y = fftconv(x, k, bias)\n",
        "        return y\n",
        "\n",
        "\n",
        "class HyenaOperator(nn.Module):\n",
        "    def __init__(\n",
        "            self,\n",
        "            d_model,\n",
        "            l_max,\n",
        "            order=2,\n",
        "            filter_order=64,\n",
        "            dropout=0.0,\n",
        "            filter_dropout=0.0,\n",
        "            **filter_args,\n",
        "        ):\n",
        "        r\"\"\"\n",
        "        Hyena operator described in the paper https://arxiv.org/pdf/2302.10866.pdf\n",
        "\n",
        "        Args:\n",
        "            d_model (int): Dimension of the input and output embeddings (width of the layer)\n",
        "            l_max: (int): Maximum input sequence length. Defaults to None\n",
        "            order: (int): Depth of the Hyena recurrence. Defaults to 2\n",
        "            dropout: (float): Dropout probability. Defaults to 0.0\n",
        "            filter_dropout: (float): Dropout probability for the filter. Defaults to 0.0\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "\n",
        "        self.d_model = d_model\n",
        "        self.l_max = l_max\n",
        "        self.order = order\n",
        "        inner_width = d_model * (order + 1)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.in_proj = nn.Linear(d_model, inner_width)\n",
        "        self.out_proj = nn.Linear(d_model, d_model)\n",
        "\n",
        "        self.short_filter = nn.Conv1d(\n",
        "            inner_width,\n",
        "            inner_width,\n",
        "            3,\n",
        "            padding=2,\n",
        "            groups=inner_width\n",
        "        )\n",
        "        self.filter_fn = HyenaFilter(\n",
        "            d_model * (order - 1),\n",
        "            order=filter_order,\n",
        "            seq_len=l_max,\n",
        "            channels=1,\n",
        "            dropout=filter_dropout,\n",
        "            **filter_args\n",
        "        )\n",
        "\n",
        "    def forward(self, u, *args, **kwargs):\n",
        "        l = u.size(-2)\n",
        "        l_filter = min(l, self.l_max)\n",
        "        u = self.in_proj(u)\n",
        "        u = rearrange(u, 'b l d -> b d l')\n",
        "\n",
        "        uc = self.short_filter(u)[...,:l_filter]\n",
        "        *x, v = uc.split(self.d_model, dim=1)\n",
        "\n",
        "        k = self.filter_fn.filter(l_filter)[0]\n",
        "        k = rearrange(k, 'l (o d) -> o d l', o=self.order - 1)\n",
        "        bias = rearrange(self.filter_fn.bias, '(o d) -> o d', o=self.order - 1)\n",
        "\n",
        "        for o, x_i in enumerate(reversed(x[1:])):\n",
        "            v = self.dropout(v * x_i)\n",
        "            v = self.filter_fn(v, l_filter, k=k[o], bias=bias[o])\n",
        "\n",
        "        y = rearrange(v * x[0], 'b d l -> b l d')\n",
        "\n",
        "        y = self.out_proj(y)\n",
        "        return y\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2I7LjzVvr9fp"
      },
      "outputs": [],
      "source": [
        "#@title Self-Attention (alternative)\n",
        "\n",
        "\"\"\"\n",
        "If you'd like to try the HyenaDNA model using attention instead, you can. ie,\n",
        "use a regular decoder only Transformer.\n",
        "\n",
        "Borrowed from the FlashAttention library by Tri Dao.\n",
        "\"\"\"\n",
        "\n",
        "class SelfAttention(nn.Module):\n",
        "    \"\"\"Implement the scaled dot product attention with softmax.\n",
        "    Arguments\n",
        "    ---------\n",
        "        softmax_scale: The temperature to use for the softmax attention.\n",
        "                      (default: 1/sqrt(d_keys) where d_keys is computed at\n",
        "                      runtime)\n",
        "        attention_dropout: The dropout rate to apply to the attention\n",
        "                           (default: 0.0)\n",
        "    \"\"\"\n",
        "    def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0):\n",
        "        super().__init__()\n",
        "        self.causal = causal\n",
        "        self.softmax_scale = softmax_scale\n",
        "        self.dropout_p = attention_dropout\n",
        "\n",
        "    def forward(self, qkv, causal=None, key_padding_mask=None):\n",
        "        \"\"\"Implements the multihead softmax attention.\n",
        "        Arguments\n",
        "        ---------\n",
        "            qkv: The tensor containing the query, key, and value. (B, S, 3, H, D)\n",
        "            causal: if passed, will override self.causal\n",
        "            key_padding_mask: boolean mask to apply to the attention weights. True means to keep,\n",
        "                False means to mask out. (B, S)\n",
        "        \"\"\"\n",
        "        batch_size, seqlen = qkv.shape[0], qkv.shape[1]\n",
        "        causal = self.causal if causal is None else causal\n",
        "        q, k, v = qkv.unbind(dim=2)\n",
        "        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])\n",
        "        scores = torch.einsum('bthd,bshd->bhts', q, k * softmax_scale)\n",
        "        if key_padding_mask is not None:\n",
        "            padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype,\n",
        "                                      device=scores.device)\n",
        "            padding_mask.masked_fill_(key_padding_mask, 0.0)\n",
        "            scores = scores + rearrange(padding_mask, 'b s -> b 1 1 s')\n",
        "        if causal:\n",
        "            # \"triu_tril_cuda_template\" not implemented for 'BFloat16'\n",
        "            # So we have to construct the mask in float\n",
        "            causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1)\n",
        "            scores = scores + causal_mask.to(dtype=scores.dtype)\n",
        "        attention = torch.softmax(scores, dim=-1, dtype=v.dtype)\n",
        "        attention_drop = F.dropout(attention, self.dropout_p if self.training else 0.0)\n",
        "        output = torch.einsum('bhts,bshd->bthd', attention_drop, v)\n",
        "        return output\n",
        "\n",
        "class MHA(nn.Module):\n",
        "    \"\"\"Multi-head self-attention and cross-attention\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, embed_dim, num_heads, bias=True, dropout=0.0,\n",
        "                 softmax_scale=None, causal=False, layer_idx=None, dwconv=False,return_residual=False,device=None, dtype=None) -> None:\n",
        "        \"\"\"\n",
        "            return_residual: whether to return the input x along with the output. This is for\n",
        "                performance reason: for post-norm architecture, returning the input allows us\n",
        "                to fuse the backward of nn.Linear with the residual connection.\n",
        "        \"\"\"\n",
        "        factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "        super().__init__()\n",
        "        self.embed_dim = embed_dim\n",
        "        self.causal = causal\n",
        "        self.layer_idx = layer_idx\n",
        "        self.dwconv = dwconv\n",
        "        self.return_residual = return_residual\n",
        "\n",
        "        self.num_heads = num_heads\n",
        "        assert self.embed_dim % num_heads == 0, \"self.kdim must be divisible by num_heads\"\n",
        "        self.head_dim = self.embed_dim // num_heads\n",
        "\n",
        "        linear_cls = nn.Linear\n",
        "        linear_resid_cls = LinearResidual\n",
        "        inner_attn_cls =  SelfAttention\n",
        "\n",
        "        if not self.return_residual:\n",
        "            self.Wqkv = linear_cls(embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs)\n",
        "        else:\n",
        "            self.Wqkv = linear_resid_cls(embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs)\n",
        "        if self.dwconv:\n",
        "            self.dwconv_qkv = nn.Conv1d(3 * embed_dim, 3 * embed_dim, kernel_size=3, padding=2,\n",
        "                                        groups=3 * embed_dim)\n",
        "\n",
        "        self.inner_attn = inner_attn_cls(causal=causal, softmax_scale=softmax_scale,\n",
        "                                         attention_dropout=dropout)\n",
        "\n",
        "        # output projection always have the bias (for now)\n",
        "        self.out_proj = linear_cls(embed_dim, embed_dim, **factory_kwargs)\n",
        "\n",
        "    def forward(self, x, key_padding_mask=None, **kwargs):\n",
        "        \"\"\"\n",
        "        Arguments:\n",
        "            x: (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim) if\n",
        "                cu_seqlens is None and max_seqlen is None, else (total, hidden_dim) where total\n",
        "                is the is the sum of the sequence lengths in the batch.\n",
        "            cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n",
        "                of the sequences in the batch, used to index into x. Only applicable when using\n",
        "                FlashAttention.\n",
        "            max_seqlen: int. Maximum sequence length in the batch.\n",
        "            key_padding_mask: boolean mask, True means to keep, False means to mask out.\n",
        "                (batch, seqlen). Only applicable when not using FlashAttention.\n",
        "            mixer_subset: for cross-attention only. If not None, will take a subset of x\n",
        "                before applying the query projection. Useful for e.g., ViT where we only care\n",
        "                about the CLS token in the last layer.\n",
        "            inference_params: for generation. Adapted from Megatron-LM (and Apex)\n",
        "            https://github.com/NVIDIA/apex/blob/3ff1a10f72ec07067c4e44759442329804ac5162/apex/transformer/testing/standalone_transformer_lm.py#L470\n",
        "        \"\"\"\n",
        "\n",
        "        kwargs = ({'key_padding_mask': key_padding_mask, **kwargs})\n",
        "\n",
        "        if not self.return_residual:\n",
        "            qkv = self.Wqkv(x)\n",
        "        else:\n",
        "            qkv, x = self.Wqkv(x)\n",
        "        if self.dwconv:\n",
        "            qkv = rearrange(self.dwconv_qkv(rearrange(qkv, 'b s d -> b d s'))[..., :-2],\n",
        "                            'b d s -> b s d').contiguous()\n",
        "        qkv = rearrange(qkv, '... (three h d) -> ... three h d', three=3, d=self.head_dim)\n",
        "\n",
        "        context = self.inner_attn(qkv, **kwargs)\n",
        "\n",
        "        out = self.out_proj(rearrange(context, '... h d -> ... (h d)'))\n",
        "        return out if not self.return_residual else (out, x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "dgkZggOetkbR"
      },
      "outputs": [],
      "source": [
        "#@title MLP layer\n",
        "\n",
        "\"\"\"\n",
        "The MLP layer after the mixer layer (HyenaOperator).\n",
        "\"\"\"\n",
        "\n",
        "class Mlp(nn.Module):\n",
        "\n",
        "    def __init__(self, in_features, hidden_features=None, out_features=None, activation=F.gelu,\n",
        "                 return_residual=False, device=None, dtype=None):\n",
        "        \"\"\"\n",
        "        From https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/modules/mlp.py\n",
        "        \"\"\"\n",
        "        factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "        super().__init__()\n",
        "        out_features = out_features or in_features\n",
        "        hidden_features = hidden_features or in_features\n",
        "        self.return_residual = return_residual\n",
        "        self.fc1 = nn.Linear(in_features, hidden_features, **factory_kwargs)\n",
        "        self.activation = activation\n",
        "        self.fc2 = nn.Linear(hidden_features, out_features, **factory_kwargs)\n",
        "\n",
        "    def forward(self, x):\n",
        "        y = self.fc1(x)\n",
        "        y = self.activation(y)\n",
        "        y = self.fc2(y)\n",
        "        return y if not self.return_residual else (y, x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xAC-fIFFstei"
      },
      "outputs": [],
      "source": [
        "#@title Block layer (Hyena + MLP layers)\n",
        "\n",
        "\"\"\"\n",
        "A block consists of a Mixer layer (Hyena or attention), and a MLP layer.\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "class LinearResidual(nn.Linear):\n",
        "    \"\"\"Wrap nn.Linear to return the residual as well. For compatibility with FusedDense.\n",
        "    \"\"\"\n",
        "\n",
        "    def forward(self, input: torch.Tensor) -> torch.Tensor:\n",
        "        return super().forward(input), input\n",
        "\n",
        "class Block(nn.Module):\n",
        "\n",
        "    def __init__(self, dim, mixer_cls=None, mlp_cls=None, norm_cls=nn.LayerNorm,\n",
        "                 dropout_cls=nn.Dropout, prenorm=True, resid_dropout1=0., resid_dropout2=0.,\n",
        "                 drop_path1=0., drop_path2=0.,\n",
        "                 return_residual=False,\n",
        "                 residual_in_fp32=False):\n",
        "        \"\"\"\n",
        "        From https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/modules/block.py\n",
        "        For prenorm=True, this Block has a slightly different structure compared to a regular\n",
        "        prenorm Transformer block.\n",
        "        The standard block is: LN -> MHA -> Dropout -> Add -> LN -> MLP -> Dropout -> Add.\n",
        "        [Ref: https://arxiv.org/abs/2002.04745]\n",
        "        Here we have: Dropout -> Add -> LN -> MHA -> Dropout -> Add -> LN -> MLP, returning both\n",
        "        the hidden_states (output of the MLP) and the residual.\n",
        "        This is for performance reasons, as we can fuse the dropout, add and LayerNorm.\n",
        "        The residual needs to be provided (except for the very first block).\n",
        "        For prenorm=False, this Block has the same structure as a regular postnorm Transformer\n",
        "        block: MHA -> Dropout -> Add -> LN -> MLP -> Dropout -> Add -> LN.\n",
        "        return_residual: whether each of the sub-layers (mixer and mlp) will return the residual.\n",
        "        This is for performance reason: for post-norm architecture, returning the input allows us\n",
        "        to fuse the backward of nn.Linear with the residual connection.\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.prenorm = prenorm\n",
        "        self.return_residual = return_residual\n",
        "        self.residual_in_fp32 = residual_in_fp32\n",
        "        if self.residual_in_fp32:\n",
        "            assert self.prenorm, 'residual_in_fp32 is only compatible with prenorm=True'\n",
        "        if mixer_cls is None:\n",
        "            mixer_cls = partial(MHA, num_heads=dim // 64)\n",
        "        if mlp_cls is None:\n",
        "            mlp_cls = partial(Mlp, hidden_features=4 * dim)\n",
        "        self.mixer = mixer_cls()\n",
        "        self.dropout1 = dropout_cls(resid_dropout1)\n",
        "        self.drop_path1 = StochasticDepth(drop_path1, mode='row')\n",
        "        self.norm1 = norm_cls(dim)\n",
        "        self.mlp = mlp_cls(dim)\n",
        "        if not isinstance(self.mlp, nn.Identity):\n",
        "            self.dropout2 = dropout_cls(resid_dropout2)\n",
        "            self.drop_path2 = StochasticDepth(drop_path2, mode='row')\n",
        "            self.norm2 = norm_cls(dim)\n",
        "\n",
        "    def forward(self, hidden_states, residual = None,\n",
        "                mixer_subset=None, mixer_kwargs=None):\n",
        "        r\"\"\"Pass the input through the encoder layer.\n",
        "        Args:\n",
        "            hidden_states: the sequence to the encoder layer (required).\n",
        "            residual: if postnorm, residual=None, If prenorm, hidden_states = Attn/MLP(LN(residual))\n",
        "            mixer_subset: for cross-attention only. If not None, will take a subset of x\n",
        "                before applying the query projection. Useful for e.g., ViT where we only care\n",
        "                about the CLS token in the last layer.\n",
        "        \"\"\"\n",
        "        if self.prenorm:\n",
        "            dropped = self.drop_path1(self.dropout1(hidden_states))\n",
        "            residual = (dropped + residual) if residual is not None else dropped\n",
        "            hidden_states = self.norm1(residual.to(dtype=self.norm1.weight.dtype))\n",
        "            if self.residual_in_fp32:\n",
        "                residual = residual.to(torch.float32)\n",
        "            if mixer_kwargs is None:\n",
        "                mixer_kwargs = {}\n",
        "            if mixer_subset is not None:\n",
        "                mixer_kwargs['mixer_subset'] = mixer_subset\n",
        "            hidden_states = self.mixer(hidden_states, **mixer_kwargs)\n",
        "            if mixer_subset is not None:\n",
        "                residual = residual[:, mixer_subset]\n",
        "            if not isinstance(self.mlp, nn.Identity):\n",
        "                dropped = self.drop_path2(self.dropout2(hidden_states))\n",
        "                residual = (dropped + residual) if residual is not None else dropped\n",
        "                hidden_states = self.norm2(residual.to(dtype=self.norm2.weight.dtype))\n",
        "                if self.residual_in_fp32:\n",
        "                    residual = residual.to(torch.float32)\n",
        "\n",
        "                hidden_states = self.mlp(hidden_states)\n",
        "            return hidden_states, residual\n",
        "        else:\n",
        "            assert residual is None\n",
        "            mixer_out = self.mixer(\n",
        "                hidden_states, **(mixer_kwargs if mixer_kwargs is not None else {})\n",
        "            )\n",
        "            if self.return_residual:  # mixer out is actually a pair here\n",
        "                mixer_out, hidden_states = mixer_out\n",
        "\n",
        "            hidden_states = self.norm1((self.drop_path1(self.dropout1(mixer_out))\n",
        "                                        + hidden_states).to(dtype=self.norm1.weight.dtype))\n",
        "\n",
        "            if not isinstance(self.mlp, nn.Identity):\n",
        "                mlp_out = self.mlp(hidden_states)\n",
        "                if self.return_residual:  # mlp out is actually a pair here\n",
        "                    mlp_out, hidden_states = mlp_out\n",
        "\n",
        "                hidden_states = self.norm2((self.drop_path2(self.dropout2(mlp_out))\n",
        "                                            + hidden_states).to(dtype=self.norm2.weight.dtype))\n",
        "\n",
        "            return hidden_states\n",
        "\n",
        "def create_mixer_cls(layer=None,\n",
        "                     attn_layer_idx=None, attn_cfg=None, layer_idx=None,\n",
        "                     device=None, dtype=None):\n",
        "    factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "    if attn_layer_idx is not None and layer_idx in attn_layer_idx:\n",
        "        causal = True if attn_cfg is None else attn_cfg.pop('causal', True)\n",
        "\n",
        "        mha_cls = MHA\n",
        "\n",
        "        mixer_cls = partial(mha_cls, causal=causal, layer_idx=layer_idx,\n",
        "                            **(attn_cfg if attn_cfg is not None else {}),**factory_kwargs)\n",
        "    else:\n",
        "        # mixer_cls = instantiate(registry.layer, layer, partial=True, layer_idx=layer_idx, **factory_kwargs)\n",
        "\n",
        "        mixer_cls = partial(HyenaOperator, **layer)\n",
        "\n",
        "    return mixer_cls\n",
        "\n",
        "def create_mlp_cls(d_model, d_inner=None, device=None, dtype=None):\n",
        "    factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "    inner_dim = d_inner if d_inner is not None else 4 * d_model\n",
        "\n",
        "    mlp_cls = partial(Mlp, hidden_features=inner_dim,\n",
        "                          activation=partial(F.gelu, approximate='tanh'), **factory_kwargs)\n",
        "\n",
        "    return mlp_cls\n",
        "\n",
        "\n",
        "def create_block(d_model, d_inner=None,\n",
        "                 layer=None, attn_layer_idx=None,\n",
        "                 attn_cfg=None, layer_norm_epsilon=1e-5,\n",
        "                 resid_dropout1=0.0, resid_dropout2=0.0, residual_in_fp32=False,\n",
        "                 layer_idx=None,\n",
        "                 device=None, dtype=None):\n",
        "    factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "    mixer_cls = create_mixer_cls(layer=layer,\n",
        "                                 attn_layer_idx=attn_layer_idx,\n",
        "                                 attn_cfg=attn_cfg, layer_idx=layer_idx,\n",
        "                                 **factory_kwargs)\n",
        "    mlp_cls = create_mlp_cls(d_model, d_inner=d_inner,\n",
        "                             **factory_kwargs)\n",
        "    norm_cls = partial(nn.LayerNorm, eps=layer_norm_epsilon, **factory_kwargs)\n",
        "    block = Block(d_model, mixer_cls, mlp_cls, norm_cls=norm_cls,\n",
        "                  prenorm=True, resid_dropout1=resid_dropout1, resid_dropout2=resid_dropout2,residual_in_fp32=residual_in_fp32)\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(module, n_layer, initializer_range=0.02, rescale_prenorm_residual=True,\n",
        "                  glu_act=False):\n",
        "    if isinstance(module, nn.Linear):\n",
        "        nn.init.normal_(module.weight, std=initializer_range)\n",
        "        if module.bias is not None:\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",
        "                nn.init.normal_(p, mean=0.0, std=initializer_range / math.sqrt(2 * n_layer))\n",
        "            # If using GLU activation for now, we scale the std by 2\n",
        "            elif name in [\"output_linear.0.weight\"]:\n",
        "                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block\n",
        "                if not glu_act:\n",
        "                    nn.init.normal_(p, mean=0.0, std=initializer_range / math.sqrt(2 * n_layer))\n",
        "                else:\n",
        "                    out_features = p.shape[0]\n",
        "                    # Multiplying the first half of the matrix by 2 since sigmoid scales it down by 0.5\n",
        "                    # on average.\n",
        "                    nn.init.normal_(p[:out_features // 2], mean=0.0, std=initializer_range / math.sqrt(2 * n_layer) * 2)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OJHCegJvtASc"
      },
      "outputs": [],
      "source": [
        "#@title Backbone model (stack of blocks)\n",
        "\n",
        "\"\"\"\n",
        "A backbone model consists of a stack of blocks. If you use attention, then\n",
        "positional embeddings are included. When using Hyena, then the pos emb\n",
        "revert to doing nothing.\n",
        "\"\"\"\n",
        "\n",
        "class GPT2Embeddings(nn.Module):\n",
        "\n",
        "    def __init__(self, embed_dim, vocab_size, max_position_embeddings, padding_idx=None,\n",
        "                 word_embed_proj_dim=None, device=None, dtype=None):\n",
        "        \"\"\"\n",
        "            If max_position_embeddings <= 0, there's no position embeddings\n",
        "            If word_embe_proj_dim is not None (e.g., OPT-350m), we embed to that dimension\n",
        "                the project up to embed_dim\n",
        "        \"\"\"\n",
        "        factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "        super().__init__()\n",
        "        if word_embed_proj_dim is None:\n",
        "            self.word_embeddings = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx,\n",
        "                                                **factory_kwargs)\n",
        "            self.project_in = None\n",
        "        else:\n",
        "            self.word_embeddings = nn.Embedding(vocab_size, word_embed_proj_dim,\n",
        "                                                padding_idx=padding_idx, **factory_kwargs)\n",
        "            self.project_in = nn.Linear(word_embed_proj_dim, embed_dim, bias=False,\n",
        "                                        **factory_kwargs)\n",
        "        self.max_position_embeddings = max_position_embeddings\n",
        "        if self.max_position_embeddings > 0:\n",
        "            self.position_embeddings = nn.Embedding(max_position_embeddings, embed_dim,\n",
        "                                                    **factory_kwargs)\n",
        "\n",
        "    def forward(self, input_ids, position_ids=None):\n",
        "        \"\"\"\n",
        "            input_ids: (batch, seqlen)\n",
        "            position_ids: (batch, seqlen)\n",
        "        \"\"\"\n",
        "        batch_size, seqlen = input_ids.shape\n",
        "        embeddings = self.word_embeddings(input_ids)\n",
        "        if self.project_in is not None:\n",
        "            embeddings = self.project_in(embeddings)\n",
        "        if self.max_position_embeddings > 0:\n",
        "            if position_ids is None:\n",
        "                position_ids = torch.arange(seqlen, dtype=torch.long, device=input_ids.device)\n",
        "            position_embeddings = self.position_embeddings(position_ids)\n",
        "            embeddings = embeddings + position_embeddings\n",
        "        return embeddings\n",
        "\n",
        "class LMBackbone(nn.Module):\n",
        "\n",
        "    def __init__(self, d_model: int, n_layer: int, d_inner: int, vocab_size: int,\n",
        "                 process_group=None, layer=None,\n",
        "                 attn_layer_idx=None, attn_cfg=None, max_position_embeddings=0,\n",
        "                 resid_dropout: float = 0.0, embed_dropout: float = 0.1,\n",
        "                 layer_norm_epsilon: float = 1e-5, initializer_cfg=None,residual_in_fp32=False,\n",
        "                 device=None, dtype=None, **kwargs) -> None:\n",
        "        factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "        super().__init__()\n",
        "        self.process_group = process_group\n",
        "        self.residual_in_fp32 = residual_in_fp32\n",
        "        # note max_position_embeddings is 0 for Hyena, and therefore isn't used\n",
        "        self.embeddings = GPT2Embeddings(d_model, vocab_size, max_position_embeddings,\n",
        "                                             **factory_kwargs)\n",
        "\n",
        "        self.layers = nn.ModuleList([create_block(\n",
        "            d_model, d_inner=d_inner,\n",
        "            layer=layer, attn_layer_idx=attn_layer_idx,\n",
        "            attn_cfg=attn_cfg, layer_norm_epsilon=layer_norm_epsilon,\n",
        "            resid_dropout1=embed_dropout if i == 0 else resid_dropout,\n",
        "            resid_dropout2=resid_dropout, residual_in_fp32=residual_in_fp32,layer_idx=i,\n",
        "            **factory_kwargs,\n",
        "        ) for i in range(n_layer)])\n",
        "\n",
        "        self.drop_f = nn.Dropout(resid_dropout)\n",
        "        self.ln_f = nn.LayerNorm(d_model, eps=layer_norm_epsilon, **factory_kwargs)\n",
        "\n",
        "        self.apply(partial(_init_weights, n_layer=n_layer,\n",
        "                           **(initializer_cfg if initializer_cfg is not None else {})))\n",
        "\n",
        "    def forward(self, input_ids, position_ids=None):\n",
        "        hidden_states = self.embeddings(input_ids, position_ids=position_ids,)\n",
        "        residual = None\n",
        "\n",
        "        for layer in self.layers:\n",
        "            hidden_states, residual = layer(hidden_states, residual)\n",
        "\n",
        "        dropped = self.drop_f(hidden_states)\n",
        "        residual = (dropped + residual) if residual is not None else dropped\n",
        "        hidden_states = self.ln_f(residual.to(dtype=self.ln_f.weight.dtype))\n",
        "\n",
        "        return hidden_states"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CR0wFk-CawPB"
      },
      "outputs": [],
      "source": [
        "#@title Decoder head layer\n",
        "\n",
        "\"\"\"\n",
        "A simple decoder head (using MLP) to predict a sequence level classification.\n",
        "You have the option to average across all the tokens in a sequence or using the\n",
        "\"last\" token to classify.  At least, those 2 worked best for us, but we provide\n",
        "other \"modes\" as well.\n",
        "\n",
        "We only need this for classification.  Otherwise we'll use the hidden\n",
        "states of the backbone as embeddings.\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "class SequenceDecoder(nn.Module):\n",
        "    def __init__(\n",
        "        self, d_model, d_output=None, l_output=None, use_lengths=False, mode=\"last\"\n",
        "    ):\n",
        "        super().__init__()\n",
        "\n",
        "        self.output_transform = nn.Identity() if d_output is None else nn.Linear(d_model, d_output)\n",
        "\n",
        "        if l_output is None:\n",
        "            self.l_output = None\n",
        "            self.squeeze = False\n",
        "        elif l_output == 0:\n",
        "            # Equivalent to getting an output of length 1 and then squeezing\n",
        "            self.l_output = 1\n",
        "            self.squeeze = True\n",
        "        else:\n",
        "            assert l_output > 0\n",
        "            self.l_output = l_output\n",
        "            self.squeeze = False\n",
        "\n",
        "        self.use_lengths = use_lengths\n",
        "        self.mode = mode\n",
        "\n",
        "        if mode == 'ragged':\n",
        "            assert not use_lengths\n",
        "\n",
        "    def forward(self, x, state=None, lengths=None, l_output=None):\n",
        "        \"\"\"\n",
        "        x: (n_batch, l_seq, d_model)\n",
        "        Returns: (n_batch, l_output, d_output)\n",
        "        \"\"\"\n",
        "\n",
        "        if self.l_output is None:\n",
        "            if l_output is not None:\n",
        "                assert isinstance(l_output, int)  # Override by pass in\n",
        "            else:\n",
        "                # Grab entire output\n",
        "                l_output = x.size(-2)\n",
        "            squeeze = False\n",
        "        else:\n",
        "            l_output = self.l_output\n",
        "            squeeze = self.squeeze\n",
        "\n",
        "        if self.mode == \"last\":\n",
        "            restrict = lambda x: x[..., -l_output:, :]\n",
        "        elif self.mode == \"first\":\n",
        "            restrict = lambda x: x[..., :l_output, :]\n",
        "        elif self.mode == \"pool\":\n",
        "            restrict = lambda x: (\n",
        "                torch.cumsum(x, dim=-2)\n",
        "                / torch.arange(\n",
        "                    1, 1 + x.size(-2), device=x.device, dtype=x.dtype\n",
        "                ).unsqueeze(-1)\n",
        "            )[..., -l_output:, :]\n",
        "\n",
        "            def restrict(x):\n",
        "                L = x.size(-2)\n",
        "                s = x.sum(dim=-2, keepdim=True)\n",
        "                if l_output > 1:\n",
        "                    c = torch.cumsum(x[..., -(l_output - 1) :, :].flip(-2), dim=-2)\n",
        "                    c = F.pad(c, (0, 0, 1, 0))\n",
        "                    s = s - c  # (B, l_output, D)\n",
        "                    s = s.flip(-2)\n",
        "                denom = torch.arange(\n",
        "                    L - l_output + 1, L + 1, dtype=x.dtype, device=x.device\n",
        "                )\n",
        "                s = s / denom\n",
        "                return s\n",
        "\n",
        "        elif self.mode == \"sum\":\n",
        "            restrict = lambda x: torch.cumsum(x, dim=-2)[..., -l_output:, :]\n",
        "            # TODO use same restrict function as pool case\n",
        "        elif self.mode == 'ragged':\n",
        "            assert lengths is not None, \"lengths must be provided for ragged mode\"\n",
        "            # remove any additional padding (beyond max length of any sequence in the batch)\n",
        "            restrict = lambda x: x[..., : max(lengths), :]\n",
        "        else:\n",
        "            raise NotImplementedError(\n",
        "                \"Mode must be ['last' | 'first' | 'pool' | 'sum']\"\n",
        "            )\n",
        "\n",
        "        # Restrict to actual length of sequence\n",
        "        if self.use_lengths:\n",
        "            assert lengths is not None\n",
        "            x = torch.stack(\n",
        "                [\n",
        "                    restrict(out[..., :length, :])\n",
        "                    for out, length in zip(torch.unbind(x, dim=0), lengths)\n",
        "                ],\n",
        "                dim=0,\n",
        "            )\n",
        "        else:\n",
        "            x = restrict(x)\n",
        "\n",
        "        if squeeze:\n",
        "            assert x.size(-2) == 1\n",
        "            x = x.squeeze(-2)\n",
        "\n",
        "        x = self.output_transform(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "    def step(self, x, state=None):\n",
        "        # Ignore all length logic\n",
        "        return self.output_transform(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZM3g4ImZtKNu"
      },
      "outputs": [],
      "source": [
        "#@title Model (backbone + head)\n",
        "\n",
        "\"\"\"\n",
        "Putting it all together, the model consists of a backbone model\n",
        "and a decoder head (you can turn off head for embeddings only too).\n",
        "\n",
        "Here we use a simple head to do multi-classification, but\n",
        "can also swap the head to do next token prediction too.  We defer to the main\n",
        "HyenaDNA for that code, since pretraining with next token prediction isn't quite\n",
        "feasible on colab.\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "class HyenaDNAModel(nn.Module):\n",
        "\n",
        "    def __init__(self, d_model: int, n_layer: int, d_inner: int, vocab_size: int,\n",
        "                 layer=None, attn_layer_idx=None, attn_cfg=None, max_position_embeddings=0,\n",
        "                 resid_dropout: float = 0.0, embed_dropout: float = 0.1,\n",
        "                 layer_norm_epsilon: float = 1e-5, initializer_cfg=None,residual_in_fp32=False,\n",
        "                 pad_vocab_size_multiple: int = 1, use_head=False, n_classes: int = 2,\n",
        "                 device=None, dtype=None, **kwargs) -> None:\n",
        "        factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "        super().__init__()\n",
        "        if vocab_size % pad_vocab_size_multiple != 0:\n",
        "            vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)\n",
        "\n",
        "        self.use_head = use_head\n",
        "\n",
        "        # check if layer (config) has d_model (HF code differs from main Safari code)\n",
        "        if 'd_model' not in layer:\n",
        "            layer['d_model'] = d_model\n",
        "\n",
        "        self.backbone = LMBackbone(\n",
        "            d_model=d_model, n_layer=n_layer, d_inner=d_inner, vocab_size=vocab_size,\n",
        "            layer=layer, attn_layer_idx=attn_layer_idx, attn_cfg=attn_cfg,\n",
        "            max_position_embeddings=max_position_embeddings,\n",
        "            resid_dropout=resid_dropout, embed_dropout=embed_dropout,\n",
        "            layer_norm_epsilon=layer_norm_epsilon,\n",
        "            initializer_cfg=initializer_cfg, residual_in_fp32=residual_in_fp32,\n",
        "            **factory_kwargs, **kwargs\n",
        "        )\n",
        "\n",
        "        # we only need a head if doing classification, otherwise we'll use the\n",
        "        # hidden states as embeddings\n",
        "        if self.use_head:\n",
        "            self.head = SequenceDecoder(d_model=d_model, d_output=n_classes, l_output=0, mode='pool')\n",
        "\n",
        "        # Initialize weights and apply final processing\n",
        "        self.apply(partial(_init_weights, n_layer=n_layer,\n",
        "                           **(initializer_cfg if initializer_cfg is not None else {})))\n",
        "\n",
        "        # if self.use_head:\n",
        "        #     self.tie_weights()\n",
        "\n",
        "    # def tie_weights(self):\n",
        "    #     self.head.weight = self.backbone.embeddings.word_embeddings.weight\n",
        "\n",
        "    def forward(self, input_ids, position_ids=None, state=None): # state for the repo interface\n",
        "        hidden_states = self.backbone(input_ids, position_ids=position_ids)\n",
        "\n",
        "        if self.use_head:\n",
        "            return self.head(hidden_states)\n",
        "        else:\n",
        "            return hidden_states"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3gGiI05HX-AU"
      },
      "outputs": [],
      "source": [
        "#@title Huggingface Pretrained Wrapper\n",
        "# for Huggingface integration, we use a wrapper class around the model\n",
        "# to load weights\n",
        "import json\n",
        "import os\n",
        "import subprocess\n",
        "import transformers\n",
        "from transformers import PreTrainedModel, AutoModelForCausalLM, PretrainedConfig\n",
        "import re\n",
        "\n",
        "def inject_substring(orig_str):\n",
        "    \"\"\"Hack to handle matching keys between models trained with and without\n",
        "    gradient checkpointing.\"\"\"\n",
        "\n",
        "    # modify for mixer keys\n",
        "    pattern = r\"\\.mixer\"\n",
        "    injection = \".mixer.layer\"\n",
        "\n",
        "    modified_string = re.sub(pattern, injection, orig_str)\n",
        "\n",
        "    # modify for mlp keys\n",
        "    pattern = r\"\\.mlp\"\n",
        "    injection = \".mlp.layer\"\n",
        "\n",
        "    modified_string = re.sub(pattern, injection, modified_string)\n",
        "\n",
        "    return modified_string\n",
        "\n",
        "def load_weights(scratch_dict, pretrained_dict, checkpointing=False):\n",
        "    \"\"\"Loads pretrained (backbone only) weights into the scratch state dict.\n",
        "\n",
        "    scratch_dict: dict, a state dict from a newly initialized HyenaDNA model\n",
        "    pretrained_dict: dict, a state dict from the pretrained ckpt\n",
        "    checkpointing: bool, whether the gradient checkpoint flag was used in the\n",
        "    pretrained model ckpt. This slightly changes state dict keys, so we patch\n",
        "    that if used.\n",
        "\n",
        "    return:\n",
        "    dict, a state dict with the pretrained weights loaded (head is scratch)\n",
        "\n",
        "    # loop thru state dict of scratch\n",
        "    # find the corresponding weights in the loaded model, and set it\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "    # need to do some state dict \"surgery\"\n",
        "    for key, value in scratch_dict.items():\n",
        "        if 'backbone' in key:\n",
        "            # the state dicts differ by one prefix, '.model', so we add that\n",
        "            key_loaded = 'model.' + key\n",
        "            # breakpoint()\n",
        "            # need to add an extra \".layer\" in key\n",
        "            if checkpointing:\n",
        "                key_loaded = inject_substring(key_loaded)\n",
        "            try:\n",
        "                scratch_dict[key] = pretrained_dict[key_loaded]\n",
        "            except:\n",
        "                raise Exception('key mismatch in the state dicts!')\n",
        "\n",
        "    # scratch_dict has been updated\n",
        "    return scratch_dict\n",
        "\n",
        "class HyenaDNAPreTrainedModel(PreTrainedModel):\n",
        "    \"\"\"\n",
        "    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained\n",
        "    models.\n",
        "    \"\"\"\n",
        "    base_model_prefix = \"hyenadna\"\n",
        "\n",
        "    def __init__(self, config):\n",
        "        pass\n",
        "\n",
        "    def forward(self, input_ids, **kwargs):\n",
        "        return self.model(input_ids, **kwargs)\n",
        "\n",
        "    @classmethod\n",
        "    def from_pretrained(cls,\n",
        "                        path,\n",
        "                        model_name,\n",
        "                        download=False,\n",
        "                        config=None,\n",
        "                        device='cpu',\n",
        "                        use_head=False,\n",
        "                        n_classes=2,\n",
        "                      ):\n",
        "        # first check if it is a local path\n",
        "        pretrained_model_name_or_path = os.path.join(path, model_name)\n",
        "        if os.path.isdir(pretrained_model_name_or_path) and download == False:\n",
        "            if config is None:\n",
        "                config = json.load(open(os.path.join(pretrained_model_name_or_path, 'config.json')))\n",
        "        else:\n",
        "            hf_url = f'https://huggingface.co/LongSafari/{model_name}'\n",
        "\n",
        "            subprocess.run(f'rm -rf {pretrained_model_name_or_path}', shell=True)\n",
        "            command = f'mkdir -p {path} && cd {path} && git lfs install && git clone {hf_url}'\n",
        "            subprocess.run(command, shell=True)\n",
        "\n",
        "            if config is None:\n",
        "                config = json.load(open(os.path.join(pretrained_model_name_or_path, 'config.json')))\n",
        "\n",
        "        scratch_model = HyenaDNAModel(**config, use_head=use_head, n_classes=n_classes)  # the new model format\n",
        "        loaded_ckpt = torch.load(\n",
        "            os.path.join(pretrained_model_name_or_path, 'weights.ckpt'),\n",
        "            map_location=torch.device(device)\n",
        "        )\n",
        "\n",
        "        # need to load weights slightly different if using gradient checkpointing\n",
        "        if config.get(\"checkpoint_mixer\", False):\n",
        "            checkpointing = config[\"checkpoint_mixer\"] == True or config[\"checkpoint_mixer\"] == True\n",
        "        else:\n",
        "            checkpointing = False\n",
        "\n",
        "        # grab state dict from both and load weights\n",
        "        state_dict = load_weights(scratch_model.state_dict(), loaded_ckpt['state_dict'], checkpointing=checkpointing)\n",
        "\n",
        "        # scratch model has now been updated\n",
        "        scratch_model.load_state_dict(state_dict)\n",
        "        print(\"Loaded pretrained weights ok!\")\n",
        "        return scratch_model\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qsRya6hEryfg"
      },
      "source": [
        "# Data pipeline\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "vsL1Afn5flo0"
      },
      "outputs": [],
      "source": [
        "#@title Tokenizer\n",
        "\n",
        "\"\"\"\n",
        "Just a simple character level tokenizer.\n",
        "\n",
        "From: https://github.com/dariush-bahrami/character-tokenizer/blob/master/charactertokenizer/core.py\n",
        "\n",
        "CharacterTokenzier for Hugging Face Transformers.\n",
        "This is heavily inspired from CanineTokenizer in transformers package.\n",
        "\"\"\"\n",
        "import json\n",
        "import os\n",
        "from pathlib import Path\n",
        "from typing import Dict, List, Optional, Sequence, Union\n",
        "\n",
        "from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer\n",
        "\n",
        "\n",
        "class CharacterTokenizer(PreTrainedTokenizer):\n",
        "    def __init__(self, characters: Sequence[str], model_max_length: int, padding_side: str='left', **kwargs):\n",
        "        \"\"\"Character tokenizer for Hugging Face transformers.\n",
        "        Args:\n",
        "            characters (Sequence[str]): List of desired characters. Any character which\n",
        "                is not included in this list will be replaced by a special token called\n",
        "                [UNK] with id=6. Following are list of all of the special tokens with\n",
        "                their corresponding ids:\n",
        "                    \"[CLS]\": 0\n",
        "                    \"[SEP]\": 1\n",
        "                    \"[BOS]\": 2\n",
        "                    \"[MASK]\": 3\n",
        "                    \"[PAD]\": 4\n",
        "                    \"[RESERVED]\": 5\n",
        "                    \"[UNK]\": 6\n",
        "                an id (starting at 7) will be assigned to each character.\n",
        "            model_max_length (int): Model maximum sequence length.\n",
        "        \"\"\"\n",
        "        self.characters = characters\n",
        "        self.model_max_length = model_max_length\n",
        "        bos_token = AddedToken(\"[BOS]\", lstrip=False, rstrip=False)\n",
        "        eos_token = AddedToken(\"[SEP]\", lstrip=False, rstrip=False)\n",
        "        sep_token = AddedToken(\"[SEP]\", lstrip=False, rstrip=False)\n",
        "        cls_token = AddedToken(\"[CLS]\", lstrip=False, rstrip=False)\n",
        "        pad_token = AddedToken(\"[PAD]\", lstrip=False, rstrip=False)\n",
        "        unk_token = AddedToken(\"[UNK]\", lstrip=False, rstrip=False)\n",
        "\n",
        "        mask_token = AddedToken(\"[MASK]\", lstrip=True, rstrip=False)\n",
        "\n",
        "        super().__init__(\n",
        "            bos_token=bos_token,\n",
        "            eos_token=sep_token,\n",
        "            sep_token=sep_token,\n",
        "            cls_token=cls_token,\n",
        "            pad_token=pad_token,\n",
        "            mask_token=mask_token,\n",
        "            unk_token=unk_token,\n",
        "            add_prefix_space=False,\n",
        "            model_max_length=model_max_length,\n",
        "            padding_side=padding_side,\n",
        "            **kwargs,\n",
        "        )\n",
        "\n",
        "        self._vocab_str_to_int = {\n",
        "            \"[CLS]\": 0,\n",
        "            \"[SEP]\": 1,\n",
        "            \"[BOS]\": 2,\n",
        "            \"[MASK]\": 3,\n",
        "            \"[PAD]\": 4,\n",
        "            \"[RESERVED]\": 5,\n",
        "            \"[UNK]\": 6,\n",
        "            **{ch: i + 7 for i, ch in enumerate(characters)},\n",
        "        }\n",
        "        self._vocab_int_to_str = {v: k for k, v in self._vocab_str_to_int.items()}\n",
        "\n",
        "    @property\n",
        "    def vocab_size(self) -> int:\n",
        "        return len(self._vocab_str_to_int)\n",
        "\n",
        "    def _tokenize(self, text: str) -> List[str]:\n",
        "        return list(text)\n",
        "\n",
        "    def _convert_token_to_id(self, token: str) -> int:\n",
        "        return self._vocab_str_to_int.get(token, self._vocab_str_to_int[\"[UNK]\"])\n",
        "\n",
        "    def _convert_id_to_token(self, index: int) -> str:\n",
        "        return self._vocab_int_to_str[index]\n",
        "\n",
        "    def convert_tokens_to_string(self, tokens):\n",
        "        return \"\".join(tokens)\n",
        "\n",
        "    def build_inputs_with_special_tokens(\n",
        "        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None\n",
        "    ) -> List[int]:\n",
        "        sep = [self.sep_token_id]\n",
        "        cls = [self.cls_token_id]\n",
        "        result = cls + token_ids_0 + sep\n",
        "        if token_ids_1 is not None:\n",
        "            result += token_ids_1 + sep\n",
        "        return result\n",
        "\n",
        "    def get_special_tokens_mask(\n",
        "        self,\n",
        "        token_ids_0: List[int],\n",
        "        token_ids_1: Optional[List[int]] = None,\n",
        "        already_has_special_tokens: bool = False,\n",
        "    ) -> List[int]:\n",
        "        if already_has_special_tokens:\n",
        "            return super().get_special_tokens_mask(\n",
        "                token_ids_0=token_ids_0,\n",
        "                token_ids_1=token_ids_1,\n",
        "                already_has_special_tokens=True,\n",
        "            )\n",
        "\n",
        "        result = [1] + ([0] * len(token_ids_0)) + [1]\n",
        "        if token_ids_1 is not None:\n",
        "            result += ([0] * len(token_ids_1)) + [1]\n",
        "        return result\n",
        "\n",
        "    def create_token_type_ids_from_sequences(\n",
        "        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None\n",
        "    ) -> List[int]:\n",
        "        sep = [self.sep_token_id]\n",
        "        cls = [self.cls_token_id]\n",
        "\n",
        "        result = len(cls + token_ids_0 + sep) * [0]\n",
        "        if token_ids_1 is not None:\n",
        "            result += len(token_ids_1 + sep) * [1]\n",
        "        return result\n",
        "\n",
        "    def get_config(self) -> Dict:\n",
        "        return {\n",
        "            \"char_ords\": [ord(ch) for ch in self.characters],\n",
        "            \"model_max_length\": self.model_max_length,\n",
        "        }\n",
        "\n",
        "    @classmethod\n",
        "    def from_config(cls, config: Dict) -> \"CharacterTokenizer\":\n",
        "        cfg = {}\n",
        "        cfg[\"characters\"] = [chr(i) for i in config[\"char_ords\"]]\n",
        "        cfg[\"model_max_length\"] = config[\"model_max_length\"]\n",
        "        return cls(**cfg)\n",
        "\n",
        "    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):\n",
        "        cfg_file = Path(save_directory) / \"tokenizer_config.json\"\n",
        "        cfg = self.get_config()\n",
        "        with open(cfg_file, \"w\") as f:\n",
        "            json.dump(cfg, f, indent=4)\n",
        "\n",
        "    @classmethod\n",
        "    def from_pretrained(cls, save_directory: Union[str, os.PathLike], **kwargs):\n",
        "        cfg_file = Path(save_directory) / \"tokenizer_config.json\"\n",
        "        with open(cfg_file) as f:\n",
        "            cfg = json.load(f)\n",
        "        return cls.from_config(cfg)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "U2_Uw320e6mk"
      },
      "outputs": [],
      "source": [
        "#@title GenomicBenchmark dataset\n",
        "\n",
        "\"\"\"\n",
        "The GenomicBenchmarks dataset will automatically download to /contents on colab.\n",
        "There are 8 datasets to choose from.\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "from random import random\n",
        "import numpy as np\n",
        "from pathlib import Path\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "from genomic_benchmarks.loc2seq import download_dataset\n",
        "from genomic_benchmarks.data_check import is_downloaded\n",
        "\n",
        "\n",
        "# helper functions\n",
        "def exists(val):\n",
        "    return val is not None\n",
        "\n",
        "def coin_flip():\n",
        "    return random() > 0.5\n",
        "\n",
        "\n",
        "string_complement_map = {'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A', 'a': 't', 'c': 'g', 'g': 'c', 't': 'a'}\n",
        "# augmentation\n",
        "def string_reverse_complement(seq):\n",
        "    rev_comp = ''\n",
        "    for base in seq[::-1]:\n",
        "        if base in string_complement_map:\n",
        "            rev_comp += string_complement_map[base]\n",
        "        # if bp not complement map, use the same bp\n",
        "        else:\n",
        "            rev_comp += base\n",
        "    return rev_comp\n",
        "\n",
        "\n",
        "class GenomicBenchmarkDataset(torch.utils.data.Dataset):\n",
        "\n",
        "    '''\n",
        "    Loop thru bed file, retrieve (chr, start, end), query fasta file for sequence.\n",
        "    Returns a generator that retrieves the sequence.\n",
        "\n",
        "    Genomic Benchmarks Dataset, from:\n",
        "    https://github.com/ML-Bioinfo-CEITEC/genomic_benchmarks\n",
        "\n",
        "\n",
        "    '''\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        split,\n",
        "        max_length,\n",
        "        dataset_name='human_enhancers_cohn',\n",
        "        d_output=2, # default binary classification\n",
        "        dest_path=\"/content\", # default for colab\n",
        "        tokenizer=None,\n",
        "        tokenizer_name=None,\n",
        "        use_padding=None,\n",
        "        add_eos=False,\n",
        "        rc_aug=False,\n",
        "        return_augs=False,\n",
        "    ):\n",
        "\n",
        "        self.max_length = max_length\n",
        "        self.use_padding = use_padding\n",
        "        self.tokenizer_name = tokenizer_name\n",
        "        self.tokenizer = tokenizer\n",
        "        self.return_augs = return_augs\n",
        "        self.add_eos = add_eos\n",
        "        self.d_output = d_output  # needed for decoder to grab\n",
        "        self.rc_aug = rc_aug\n",
        "\n",
        "        if not is_downloaded(dataset_name, cache_path=dest_path):\n",
        "            print(\"downloading {} to {}\".format(dataset_name, dest_path))\n",
        "            download_dataset(dataset_name, version=0, dest_path=dest_path)\n",
        "        else:\n",
        "            print(\"already downloaded {}-{}\".format(split, dataset_name))\n",
        "\n",
        "        # use Path object\n",
        "        base_path = Path(dest_path) / dataset_name / split\n",
        "\n",
        "        self.all_paths = []\n",
        "        self.all_labels = []\n",
        "        label_mapper = {}\n",
        "\n",
        "        for i, x in enumerate(base_path.iterdir()):\n",
        "            label_mapper[x.stem] = i\n",
        "\n",
        "        for label_type in label_mapper.keys():\n",
        "            for x in (base_path / label_type).iterdir():\n",
        "                self.all_paths.append(x)\n",
        "                self.all_labels.append(label_mapper[label_type])\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.all_paths)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        txt_path = self.all_paths[idx]\n",
        "        with open(txt_path, \"r\") as f:\n",
        "            content = f.read()\n",
        "        x = content\n",
        "        y = self.all_labels[idx]\n",
        "\n",
        "        # apply rc_aug here if using\n",
        "        if self.rc_aug and coin_flip():\n",
        "            x = string_reverse_complement(x)\n",
        "\n",
        "        seq = self.tokenizer(x,\n",
        "            add_special_tokens=False,\n",
        "            padding=\"max_length\" if self.use_padding else None,\n",
        "            max_length=self.max_length,\n",
        "            truncation=True,\n",
        "        )  # add cls and eos token (+2)\n",
        "        seq = seq[\"input_ids\"]  # get input_ids\n",
        "\n",
        "        # need to handle eos here\n",
        "        if self.add_eos:\n",
        "            # append list seems to be faster than append tensor\n",
        "            seq.append(self.tokenizer.sep_token_id)\n",
        "\n",
        "        # convert to tensor\n",
        "        seq = torch.LongTensor(seq)\n",
        "\n",
        "        # need to wrap in list\n",
        "        target = torch.LongTensor([y])\n",
        "\n",
        "        return seq, target"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5vMJu0wQWUBN"
      },
      "source": [
        "# Training! (and fine-tuning)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dMNRAa4gKJOK"
      },
      "outputs": [],
      "source": [
        "import torch.optim as optim\n",
        "\n",
        "\"\"\"\n",
        "We provide simple training code for the GenomicBenchmark datasets.\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "def train(model, device, train_loader, optimizer, epoch, loss_fn, log_interval=10):\n",
        "    \"\"\"Training loop.\"\"\"\n",
        "    model.train()\n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        data, target = data.to(device), target.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        output = model(data)\n",
        "        loss = loss_fn(output, target.squeeze())\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        if batch_idx % log_interval == 0:\n",
        "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
        "                epoch, batch_idx * len(data), len(train_loader.dataset),\n",
        "                100. * batch_idx / len(train_loader), loss.item()))\n",
        "\n",
        "def test(model, device, test_loader, loss_fn):\n",
        "    \"\"\"Test loop.\"\"\"\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    with torch.no_grad():\n",
        "        for data, target in test_loader:\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            output = model(data)\n",
        "            test_loss += loss_fn(output, target.squeeze()).item()  # sum up batch loss\n",
        "            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability\n",
        "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
        "\n",
        "    test_loss /= len(test_loader.dataset)\n",
        "\n",
        "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\\n'.format(\n",
        "        test_loss, correct, len(test_loader.dataset),\n",
        "        100. * correct / len(test_loader.dataset)))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z6sguA32WZE3"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "import subprocess\n",
        "import transformers\n",
        "from transformers import PreTrainedModel, AutoModelForCausalLM, PretrainedConfig\n",
        "\n",
        "def run_train():\n",
        "\n",
        "    '''\n",
        "    Main entry point for training.  Select the dataset name and metadata, as\n",
        "    well as model and training args, and you're off to the genomic races!\n",
        "\n",
        "    ### GenomicBenchmarks Metadata\n",
        "    # there are 8 datasets in this suite, choose 1 at a time, with their corresponding settings\n",
        "    # name                                num_seqs        num_classes     median len    std\n",
        "    # dummy_mouse_enhancers_ensembl       1210            2               2381          984.4\n",
        "    # demo_coding_vs_intergenomic_seqs    100_000         2               200           0\n",
        "    # demo_human_or_worm                  100_000         2               200           0\n",
        "    # human_enhancers_cohn                27791           2               500           0\n",
        "    # human_enhancers_ensembl             154842          2               269           122.6\n",
        "    # human_ensembl_regulatory            289061          3               401           184.3\n",
        "    # human_nontata_promoters             36131           2               251           0\n",
        "    # human_ocr_ensembl                   174756          2               315           108.1\n",
        "\n",
        "    '''\n",
        "    # experiment settings:\n",
        "    num_epochs = 100  # ~100 seems fine\n",
        "    max_length = 500  # max len of sequence of dataset (of what you want)\n",
        "    use_padding = True\n",
        "    dataset_name = 'human_enhancers_cohn'\n",
        "    batch_size = 256\n",
        "    learning_rate = 6e-4  # good default for Hyena\n",
        "    rc_aug = True  # reverse complement augmentation\n",
        "    add_eos = False  # add end of sentence token\n",
        "    weight_decay = 0.1\n",
        "\n",
        "    # for fine-tuning, only the 'tiny' model can fit on colab\n",
        "    pretrained_model_name = 'hyenadna-tiny-1k-seqlen'  # use None if training from scratch\n",
        "\n",
        "    # we need these for the decoder head, if using\n",
        "    use_head = True\n",
        "    n_classes = 2\n",
        "\n",
        "    # you can override with your own backbone config here if you want,\n",
        "    # otherwise we'll load the HF one by default\n",
        "    backbone_cfg = None\n",
        "\n",
        "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "    print(\"Using device:\", device)\n",
        "\n",
        "    # instantiate the model (pretrained here)\n",
        "    if pretrained_model_name in ['hyenadna-tiny-1k-seqlen']:\n",
        "        # use the pretrained Huggingface wrapper instead\n",
        "        model = HyenaDNAPreTrainedModel.from_pretrained(\n",
        "            './checkpoints',\n",
        "            pretrained_model_name,\n",
        "            download=True,\n",
        "            config=backbone_cfg,\n",
        "            device=device,\n",
        "            use_head=use_head,\n",
        "            n_classes=n_classes,\n",
        "        )\n",
        "\n",
        "    # from scratch\n",
        "    else:\n",
        "        model = HyenaDNAModel(**backbone_cfg, use_head=use_head, n_classes=n_classes)\n",
        "\n",
        "    # create tokenizer\n",
        "    tokenizer = CharacterTokenizer(\n",
        "        characters=['A', 'C', 'G', 'T', 'N'],  # add DNA characters, N is uncertain\n",
        "        model_max_length=max_length + 2,  # to account for special tokens, like EOS\n",
        "        add_special_tokens=False,  # we handle special tokens elsewhere\n",
        "        padding_side='left', # since HyenaDNA is causal, we pad on the left\n",
        "    )\n",
        "\n",
        "    # create datasets\n",
        "    ds_train = GenomicBenchmarkDataset(\n",
        "        max_length = max_length,\n",
        "        use_padding = use_padding,\n",
        "        split = 'train',\n",
        "        tokenizer=tokenizer,\n",
        "        dataset_name=dataset_name,\n",
        "        rc_aug=rc_aug,\n",
        "        add_eos=add_eos,\n",
        "    )\n",
        "\n",
        "    ds_test = GenomicBenchmarkDataset(\n",
        "        max_length = max_length,\n",
        "        use_padding = use_padding,\n",
        "        split = 'test',\n",
        "        tokenizer=tokenizer,\n",
        "        dataset_name=dataset_name,\n",
        "        rc_aug=rc_aug,\n",
        "        add_eos=add_eos,\n",
        "    )\n",
        "\n",
        "    train_loader = DataLoader(ds_train, batch_size=batch_size, shuffle=True)\n",
        "    test_loader = DataLoader(ds_test, batch_size=batch_size, shuffle=False)\n",
        "\n",
        "    # loss function\n",
        "    loss_fn = nn.CrossEntropyLoss()\n",
        "\n",
        "    # create optimizer\n",
        "    optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)\n",
        "\n",
        "    model.to(device)\n",
        "\n",
        "    for epoch in range(num_epochs):\n",
        "        train(model, device, train_loader, optimizer, epoch, loss_fn)\n",
        "        test(model, device, test_loader, loss_fn)\n",
        "        optimizer.step()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GhYYO_NP-Au4"
      },
      "outputs": [],
      "source": [
        "# launch it!\n",
        "# run_train()  # uncomment to run"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OMgEPV0J92Xj"
      },
      "source": [
        "# Inference (450k to 1M tokens)!\n",
        "\n",
        "If all you're interested in is getting embeddings on long DNA sequences\n",
        "(inference), then we can do that right here in Colab!\n",
        "\n",
        "\n",
        "*   We provide an example how to load the weights from Huggingface.\n",
        "*   On the free tier, which uses a\n",
        "T4 GPU w/16GB of memory, we can process 450k tokens / nucleotides.\n",
        "*   For processing 1M tokens, you'll need an A100, which Colab offers as a paid tier.\n",
        "*   (Don't forget to run the entire notebook above too)\n",
        "\n",
        "--\n",
        "\n",
        "To pretrain or fine-tune the 1M long sequence model (8 layers, d_model=256),\n",
        "you'll need 8 A100s 80GB, and all that code is in the main [github repo](https://github.com/HazyResearch/hyena-dna)!\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-1wq2uwUctPV"
      },
      "outputs": [],
      "source": [
        "#@title Single example\n",
        "import json\n",
        "import os\n",
        "import subprocess\n",
        "import transformers\n",
        "from transformers import PreTrainedModel, AutoModelForCausalLM, PretrainedConfig\n",
        "\n",
        "def inference_single():\n",
        "\n",
        "    '''\n",
        "    this selects which backbone to use, and grabs weights/ config from HF\n",
        "    4 options:\n",
        "      'hyenadna-tiny-1k-seqlen'   # fine-tune on colab ok\n",
        "      'hyenadna-small-32k-seqlen'\n",
        "      'hyenadna-medium-160k-seqlen'  # inference only on colab\n",
        "      'hyenadna-medium-450k-seqlen'  # inference only on colab\n",
        "      'hyenadna-large-1m-seqlen'  # inference only on colab\n",
        "    '''\n",
        "\n",
        "    # you only need to select which model to use here, we'll do the rest!\n",
        "    pretrained_model_name = 'hyenadna-small-32k-seqlen'\n",
        "\n",
        "    max_lengths = {\n",
        "        'hyenadna-tiny-1k-seqlen': 1024,\n",
        "        'hyenadna-small-32k-seqlen': 32768,\n",
        "        'hyenadna-medium-160k-seqlen': 160000,\n",
        "        'hyenadna-medium-450k-seqlen': 450000,  # T4 up to here\n",
        "        'hyenadna-large-1m-seqlen': 1_000_000,  # only A100 (paid tier)\n",
        "    }\n",
        "\n",
        "    max_length = max_lengths[pretrained_model_name]  # auto selects\n",
        "\n",
        "    # data settings:\n",
        "    use_padding = True\n",
        "    rc_aug = False  # reverse complement augmentation\n",
        "    add_eos = False  # add end of sentence token\n",
        "\n",
        "    # we need these for the decoder head, if using\n",
        "    use_head = False\n",
        "    n_classes = 2  # not used for embeddings only\n",
        "\n",
        "    # you can override with your own backbone config here if you want,\n",
        "    # otherwise we'll load the HF one in None\n",
        "    backbone_cfg = None\n",
        "\n",
        "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "    print(\"Using device:\", device)\n",
        "\n",
        "    # instantiate the model (pretrained here)\n",
        "    if pretrained_model_name in ['hyenadna-tiny-1k-seqlen',\n",
        "                                 'hyenadna-small-32k-seqlen',\n",
        "                                 'hyenadna-medium-160k-seqlen',\n",
        "                                 'hyenadna-medium-450k-seqlen',\n",
        "                                 'hyenadna-large-1m-seqlen']:\n",
        "        # use the pretrained Huggingface wrapper instead\n",
        "        model = HyenaDNAPreTrainedModel.from_pretrained(\n",
        "            './checkpoints',\n",
        "            pretrained_model_name,\n",
        "            download=True,\n",
        "            config=backbone_cfg,\n",
        "            device=device,\n",
        "            use_head=use_head,\n",
        "            n_classes=n_classes,\n",
        "        )\n",
        "\n",
        "    # from scratch\n",
        "    elif pretrained_model_name is None:\n",
        "        model = HyenaDNAModel(**backbone_cfg, use_head=use_head, n_classes=n_classes)\n",
        "\n",
        "    # create tokenizer\n",
        "    tokenizer = CharacterTokenizer(\n",
        "        characters=['A', 'C', 'G', 'T', 'N'],  # add DNA characters, N is uncertain\n",
        "        model_max_length=max_length + 2,  # to account for special tokens, like EOS\n",
        "        add_special_tokens=False,  # we handle special tokens elsewhere\n",
        "        padding_side='left', # since HyenaDNA is causal, we pad on the left\n",
        "    )\n",
        "\n",
        "    #### Single embedding example ####\n",
        "\n",
        "    # create a sample 450k long, prepare\n",
        "    sequence = 'ACTG' * int(max_length/4)\n",
        "    tok_seq = tokenizer(sequence)\n",
        "    tok_seq = tok_seq[\"input_ids\"]  # grab ids\n",
        "\n",
        "    # place on device, convert to tensor\n",
        "    tok_seq = torch.LongTensor(tok_seq).unsqueeze(0)  # unsqueeze for batch dim\n",
        "    tok_seq = tok_seq.to(device)\n",
        "\n",
        "    # prep model and forward\n",
        "    model.to(device)\n",
        "    model.eval()\n",
        "    with torch.inference_mode():\n",
        "        embeddings = model(tok_seq)\n",
        "\n",
        "    print(embeddings.shape)  # embeddings here!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2BYJx_XSiXZr",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5bc33934-6a5c-4a67-8b96-c6c71db75627"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using device: cpu\n",
            "Loaded pretrained weights ok!\n",
            "torch.Size([1, 32770, 256])\n"
          ]
        }
      ],
      "source": [
        "# # uncomment to run! (to get embeddings)\n",
        "inference_single()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "arcIbXNzHmyj",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Batch example\n",
        "\"\"\"\n",
        "Let's say you want to do inference on a dataset to grab a lot of embeddings,\n",
        "you can just loop thru a dataloader like this.\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "import json\n",
        "import os\n",
        "import subprocess\n",
        "import transformers\n",
        "from transformers import PreTrainedModel, AutoModelForCausalLM, PretrainedConfig\n",
        "\n",
        "# example, if you want to do look thru a dataloader and get embeddings\n",
        "\n",
        "def infer_loop(model, device, test_loader):\n",
        "    \"\"\"inference loop.\"\"\"\n",
        "    # model.eval()\n",
        "\n",
        "    with torch.inference_mode():\n",
        "        for data, target in test_loader:\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            output = model(data)\n",
        "\n",
        "            # print embedding shape, or do something else here, like save.\n",
        "            print('embeddings shape', output.shape)\n",
        "\n",
        "def inference():\n",
        "\n",
        "    '''\n",
        "    this selects which backbone to use, and grabs weights/ config from HF\n",
        "    4 options:\n",
        "      'hyenadna-tiny-1k-seqlen'   # fine-tune on colab ok\n",
        "      'hyenadna-small-32k-seqlen'\n",
        "      'hyenadna-medium-160k-seqlen'  # inference only on colab\n",
        "      'hyenadna-medium-450k-seqlen'  # inference only on colab\n",
        "      'hyenadna-large-1m-seqlen'  # inference only on colab\n",
        "    '''\n",
        "\n",
        "    # select model\n",
        "    pretrained_model_name = 'hyenadna-medium-160k-seqlen'  # use None if training from scratch\n",
        "\n",
        "    max_lengths = {\n",
        "        'hyenadna-tiny-1k-seqlen': 1024,\n",
        "        'hyenadna-small-32k-seqlen': 32768,\n",
        "        'hyenadna-medium-160k-seqlen': 160000,\n",
        "        'hyenadna-medium-450k-seqlen': 450000,  # T4 up to here\n",
        "        'hyenadna-large-1m-seqlen': 1_000_000,  # only A100 (paid tier)\n",
        "    }\n",
        "\n",
        "    # let's fix the max_length (to reduce the padding amount, conserves memory)\n",
        "    max_length = 500\n",
        "\n",
        "    # data settings:\n",
        "    # we need to choose the dataset and batch size to loop thru\n",
        "    dataset_name = 'human_enhancers_cohn'\n",
        "    batch_size = 4\n",
        "    use_padding = True\n",
        "    rc_aug = False  # reverse complement augmentation\n",
        "    add_eos = False  # add end of sentence token\n",
        "\n",
        "    # we need these for the decoder head, if using\n",
        "    use_head = False\n",
        "    n_classes = 2  # not used for embeddings only\n",
        "\n",
        "    # you can override with your own backbone config here if you want,\n",
        "    # otherwise we'll load the HF one in None\n",
        "    backbone_cfg = None\n",
        "\n",
        "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "    print(\"Using device:\", device)\n",
        "\n",
        "    # instantiate the model (pretrained here)\n",
        "    if pretrained_model_name in ['hyenadna-tiny-1k-seqlen',\n",
        "                                 'hyenadna-small-32k-seqlen',\n",
        "                                 'hyenadna-medium-160k-seqlen',\n",
        "                                 'hyenadna-medium-450k-seqlen',\n",
        "                                 'hyenadna-large-1m-seqlen']:\n",
        "        # use the pretrained Huggingface wrapper instead\n",
        "        model = HyenaDNAPreTrainedModel.from_pretrained(\n",
        "            './checkpoints',\n",
        "            pretrained_model_name,\n",
        "            download=True,\n",
        "            config=backbone_cfg,\n",
        "            device=device,\n",
        "            use_head=use_head,\n",
        "            n_classes=n_classes,\n",
        "        )\n",
        "\n",
        "    # from scratch\n",
        "    elif pretrained_model_name is None:\n",
        "        model = HyenaDNAModel(**backbone_cfg, use_head=use_head, n_classes=n_classes)\n",
        "\n",
        "    # create tokenizer\n",
        "    tokenizer = CharacterTokenizer(\n",
        "        characters=['A', 'C', 'G', 'T', 'N'],  # add DNA characters, N is uncertain\n",
        "        model_max_length=max_length + 2,  # to account for special tokens, like EOS\n",
        "        add_special_tokens=False,  # we handle special tokens elsewhere\n",
        "        padding_side='left', # since HyenaDNA is causal, we pad on the left\n",
        "    )\n",
        "\n",
        "    # Sample small dataloader w/ GenomicBenchmarks\n",
        "\n",
        "    ds_test = GenomicBenchmarkDataset(\n",
        "        max_length = max_length,\n",
        "        use_padding = use_padding,\n",
        "        split = 'test',\n",
        "        tokenizer=tokenizer,\n",
        "        dataset_name=dataset_name,\n",
        "        rc_aug=rc_aug,\n",
        "        add_eos=add_eos,\n",
        "    )\n",
        "\n",
        "    ds_test = GenomicBenchmarkDataset(\n",
        "        max_length = max_length,\n",
        "        use_padding = use_padding,\n",
        "        split = 'test',\n",
        "        tokenizer=tokenizer,\n",
        "        dataset_name=dataset_name,\n",
        "        rc_aug=rc_aug,\n",
        "        add_eos=add_eos,\n",
        "    )\n",
        "\n",
        "    test_loader = DataLoader(ds_test, batch_size=batch_size, shuffle=False)\n",
        "\n",
        "    model.to(device)\n",
        "    model.eval()\n",
        "\n",
        "    infer_loop(model, device, test_loader)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DSkwVPfAHnAG"
      },
      "outputs": [],
      "source": [
        "# uncomment to run!\n",
        "# inference()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DdWb4ti0iXkB"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-1rs0t9RIWKt"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "5vMJu0wQWUBN"
      ],
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}