deep-cross-attention 0.0.2__tar.gz

deep_cross_attention-0.0.2/.github/workflows/python-publish.yml

@@ -0,0 +1,36 @@
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build
+    - name: Build package
+      run: python -m build
+    - name: Publish package
+      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      with:
+        user: __token__
+        password: ${{ secrets.PYPI_API_TOKEN }}

deep_cross_attention-0.0.2/.gitignore

@@ -0,0 +1,171 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# PyPI configuration file
+.pypirc

deep_cross_attention-0.0.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Phil Wang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

deep_cross_attention-0.0.2/PKG-INFO

@@ -0,0 +1,103 @@
+Metadata-Version: 2.4
+Name: deep-cross-attention
+Version: 0.0.2
+Summary: Deep Cross Attention Language Model
+Project-URL: Homepage, https://pypi.org/project/deep-cross-attention/
+Project-URL: Repository, https://github.com/lucidrains/deep-cross-attention
+Author-email: Phil Wang <lucidrains@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2025 Phil Wang
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: artificial intelligence,deep learning,residuals,transformers
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Requires-Dist: einops>=0.8.0
+Requires-Dist: rotary-embedding-torch
+Requires-Dist: torch>=2.3
+Provides-Extra: examples
+Requires-Dist: tqdm; extra == 'examples'
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Description-Content-Type: text/markdown
+
+<img src="./fig4.png" width="400px"></img>
+
+## Deep Cross Attention
+
+Implementation of the proposed [DeepCrossAttention](https://arxiv.org/abs/2502.06785) by [Mike Heddes](https://www.mikeheddes.nl/) while at Google research, in Pytorch
+
+My analysis is although I still prefer [Hyper Connections](https://arxiv.org/abs/2409.19606), they have an important idea here that I have been trying concurrently. Mainly the queries, keys, values can be [routed from different layers](https://github.com/lucidrains/x-transformers/blob/main/x_transformers/x_transformers.py#L1226) of the past
+
+## Install
+
+```bash
+$ pip install deep-cross-attention
+```
+
+## Usage
+
+```python
+import torch
+from deep_cross_attention import DCAGPT
+
+gpt = DCAGPT(
+    num_tokens = 256,
+    dim = 512,
+    depth = 6,
+    heads = 8,
+    dim_head = 64,
+    past_layers_k = 2
+)
+
+ids = torch.randint(0, 256, (2, 4096))
+
+logits = gpt(ids) # (2, 4096, 256)
+```
+
+## Example
+
+First
+
+```bash
+$ pip install .[examples]
+```
+
+Next
+
+```bash
+$ python train.py
+```
+
+## Citations
+
+```bibtex
+@inproceedings{Heddes2025DeepCrossAttentionST,
+    title   = {DeepCrossAttention: Supercharging Transformer Residual Connections},
+    author  = {Mike Heddes and Adel Javanmard and Kyriakos Axiotis and Gang Fu and MohammadHossein Bateni and Vahab S. Mirrokni},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:276250576}
+}
+```

deep_cross_attention-0.0.2/README.md

@@ -0,0 +1,58 @@
+<img src="./fig4.png" width="400px"></img>
+
+## Deep Cross Attention
+
+Implementation of the proposed [DeepCrossAttention](https://arxiv.org/abs/2502.06785) by [Mike Heddes](https://www.mikeheddes.nl/) while at Google research, in Pytorch
+
+My analysis is although I still prefer [Hyper Connections](https://arxiv.org/abs/2409.19606), they have an important idea here that I have been trying concurrently. Mainly the queries, keys, values can be [routed from different layers](https://github.com/lucidrains/x-transformers/blob/main/x_transformers/x_transformers.py#L1226) of the past
+
+## Install
+
+```bash
+$ pip install deep-cross-attention
+```
+
+## Usage
+
+```python
+import torch
+from deep_cross_attention import DCAGPT
+
+gpt = DCAGPT(
+    num_tokens = 256,
+    dim = 512,
+    depth = 6,
+    heads = 8,
+    dim_head = 64,
+    past_layers_k = 2
+)
+
+ids = torch.randint(0, 256, (2, 4096))
+
+logits = gpt(ids) # (2, 4096, 256)
+```
+
+## Example
+
+First
+
+```bash
+$ pip install .[examples]
+```
+
+Next
+
+```bash
+$ python train.py
+```
+
+## Citations
+
+```bibtex
+@inproceedings{Heddes2025DeepCrossAttentionST,
+    title   = {DeepCrossAttention: Supercharging Transformer Residual Connections},
+    author  = {Mike Heddes and Adel Javanmard and Kyriakos Axiotis and Gang Fu and MohammadHossein Bateni and Vahab S. Mirrokni},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:276250576}
+}
+```

deep_cross_attention-0.0.2/data/README.md

@@ -0,0 +1,3 @@
+# Data source
+
+The enwik8 data was downloaded from the Hutter prize page: http://prize.hutter1.net/

deep_cross_attention-0.0.2/deep_cross_attention/__init__.py

@@ -0,0 +1,2 @@
+from deep_cross_attention.gpt import GPT
+from deep_cross_attention.dca_gpt import DCAGPT

deep_cross_attention-0.0.2/deep_cross_attention/dca_gpt.py

@@ -0,0 +1,230 @@
+import torch
+from torch import nn, cat, stack
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList, Linear, RMSNorm
+
+from einops import rearrange, einsum
+from einops.layers.torch import Rearrange
+
+from rotary_embedding_torch import RotaryEmbedding
+
+# functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+# attention
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        dim_head = 64,
+        heads = 8
+    ):
+        super().__init__()
+        self.norm = RMSNorm(dim)
+
+        self.heads = heads
+        dim_inner = heads * dim_head
+
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        self.to_q = nn.Sequential(RMSNorm(dim), nn.Linear(dim, dim_inner, bias = False))
+        self.to_k = nn.Sequential(RMSNorm(dim), nn.Linear(dim, dim_inner, bias = False))
+        self.to_v = nn.Sequential(RMSNorm(dim), nn.Linear(dim, dim_inner, bias = False))
+
+        self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
+        self.merge_heads = Rearrange('b h n d -> b n (h d)')
+
+        self.to_out = nn.Linear(dim_inner, dim, bias = False)
+
+    def forward(
+        self,
+        q_input,
+        k_input,
+        v_input
+    ):
+
+        q = self.to_q(q_input)
+        k = self.to_k(k_input)
+        v = self.to_v(v_input)
+
+        q, k, v = map(self.split_heads, (q, k, v))
+
+        # relative positions
+
+        q, k = self.rotary_embed.rotate_queries_with_cached_keys(q, k)
+
+        # attention branch
+
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal = True
+        )
+
+        out = self.merge_heads(out)
+
+        return self.to_out(out)
+
+# feedforward
+
+def FeedForward(dim, expansion_factor = 4.):
+    dim_hidden = int(dim * expansion_factor)
+
+    return nn.Sequential(
+        Linear(dim, dim_hidden),
+        nn.GELU(),
+        Linear(dim_hidden, dim)
+    )
+
+# GRNv3
+# the input dependent one lines up with all the literature, and the winning solution for hyper connections (dynamic)
+
+class GRN(Module):
+    def __init__(
+        self,
+        dim,
+    ):
+        super().__init__()
+
+        self.to_aggregate = nn.Sequential(
+            RMSNorm(dim),
+            Linear(dim, 1),
+            nn.ReLU(),
+        )
+
+        nn.init.zeros_(self.to_aggregate[-2].weight)
+
+    def forward(
+        self,
+        tokens_across_depth # Float['depth b n d']
+    ):
+        aggregate = self.to_aggregate(tokens_across_depth)
+
+        return (tokens_across_depth * aggregate).sum(dim = 0)
+
+# DCA Decoder Block
+
+class DCABlock(Module):
+    def __init__(
+        self,
+        dim,
+        dim_head = 64,
+        heads = 8,
+        ff_expansion_factor = 4.
+    ):
+        super().__init__()
+
+        self.q_grn = GRN(dim)
+        self.k_grn = GRN(dim)
+        self.v_grn = GRN(dim)
+
+        self.attn = Attention(dim = dim, dim_head = dim_head, heads = heads)
+
+        self.pre_ff_norm = RMSNorm(dim)
+
+        self.ff = FeedForward(dim = dim, expansion_factor = ff_expansion_factor)
+
+    def forward(
+        self,
+        tokens_across_depth # Float['depth b n d']
+    ):
+        q_input, k_input, v_input = self.q_grn(tokens_across_depth), self.k_grn(tokens_across_depth), self.v_grn(tokens_across_depth)
+
+        residual = q_input
+
+        attn_out = self.attn(q_input, k_input, v_input)
+
+        ff_input = self.pre_ff_norm(attn_out + residual)
+
+        ff_out = self.ff(ff_input)
+
+        return ff_out + attn_out
+
+# classes
+
+class DCAGPT(Module):
+    def __init__(
+        self,
+        num_tokens,
+        dim,
+        depth,
+        past_layers_k = 2,
+        dim_head = 64,
+        heads = 8,
+        ff_expansion_factor = 4.
+    ):
+        super().__init__()
+        self.token_emb = nn.Embedding(num_tokens, dim)
+
+        # the `k` hyperparameter, which seems to refer to sub sampling of which layers to include for efficiency
+        # but weirdly, they not only do last k layers, but also the first k? also some mention about intermediate layers being pooled? just go with first and last for now
+
+        self.past_layers_k = past_layers_k
+
+        # the proposed DCA blocks
+
+        dca_blocks = []
+        for _ in range(depth):
+            dca = DCABlock(dim = dim, dim_head = dim_head, heads = heads, ff_expansion_factor = ff_expansion_factor)
+
+            dca_blocks.append(dca)
+
+        self.dca_blocks = ModuleList(dca_blocks)
+
+        # norm and logits
+
+        self.final_grn = GRN(dim)
+
+        self.norm = RMSNorm(dim)
+        self.to_logits = Linear(dim, num_tokens, bias = False)
+ 
+    def forward(
+        self,
+        ids,
+        return_loss = False
+    ):
+        k = self.past_layers_k # k in paper
+
+        if return_loss:
+            ids, labels = ids[:, :-1], ids[:, 1:]
+
+        tokens = self.token_emb(ids)
+
+        all_tokens = [tokens]
+
+        for dca_block in self.dca_blocks:
+
+            all_tokens_stacked = stack(all_tokens)
+            num_layers = all_tokens_stacked.shape[0]
+
+            # determine which layers to include
+
+            if num_layers < (k * 2):
+                dca_block_input = all_tokens_stacked
+            else:
+                dca_block_input = cat((
+                    all_tokens_stacked[:k], # first k layers
+                    all_tokens_stacked[-k:] # last k layers
+                ))
+
+            dca_out = dca_block(dca_block_input)
+
+            # append dca output for next iteration
+
+            all_tokens.append(dca_out)
+
+        pooled_tokens = self.final_grn(stack(all_tokens))
+
+        embed = self.norm(pooled_tokens)
+
+        logits = self.to_logits(embed)
+
+        if not return_loss:
+            return logits
+
+        return F.cross_entropy(rearrange(logits, 'b n l -> b l n'), labels)

deep_cross_attention-0.0.2/pyproject.toml

@@ -0,0 +1,61 @@
+[project]
+name = "deep-cross-attention"
+version = "0.0.2"
+description = "Deep Cross Attention Language Model"
+authors = [
+    { name = "Phil Wang", email = "lucidrains@gmail.com" }
+]
+readme = "README.md"
+requires-python = ">= 3.9"
+license = { file = "LICENSE" }
+keywords = [
+    'artificial intelligence',
+    'deep learning',
+    'transformers',
+    'residuals'
+]
+
+classifiers=[
+    'Development Status :: 4 - Beta',
+    'Intended Audience :: Developers',
+    'Topic :: Scientific/Engineering :: Artificial Intelligence',
+    'License :: OSI Approved :: MIT License',
+    'Programming Language :: Python :: 3.9',
+]
+
+dependencies = [
+    "einops>=0.8.0",
+    "rotary-embedding-torch",
+    "torch>=2.3",
+]
+
+[project.urls]
+Homepage = "https://pypi.org/project/deep-cross-attention/"
+Repository = "https://github.com/lucidrains/deep-cross-attention"
+
+[project.optional-dependencies]
+examples = [
+    "tqdm"
+]
+test = [
+    "pytest"
+]
+
+[tool.pytest.ini_options]
+pythonpath = [
+  "."
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.rye]
+managed = true
+dev-dependencies = []
+
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.hatch.build.targets.wheel]
+packages = ["deep_cross_attention"]

deep_cross_attention-0.0.2/train.py

@@ -0,0 +1,179 @@
+import math
+import gzip
+import random
+import tqdm
+import numpy as np
+
+import torch
+from torch.optim import Adam
+from torch import Tensor
+from torch.utils.data import DataLoader, Dataset
+
+from deep_cross_attention import (
+    GPT,
+    DCAGPT
+)
+
+# constants
+
+NUM_BATCHES = int(1e5)
+BATCH_SIZE = 4
+GRAD_ACCUM_EVERY = 4
+LEARNING_RATE = 1e-4
+VALIDATE_EVERY = 100
+PRIME_LENGTH = 128
+GENERATE_EVERY = 500
+GENERATE_LENGTH = 512
+SEQ_LEN = 512
+
+USE_DCA = True
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def cycle(loader):
+    while True:
+        for data in loader:
+            yield data
+
+def decode_token(token):
+    return str(chr(max(32, token)))
+
+def decode_tokens(tokens):
+    return "".join(list(map(decode_token, tokens)))
+
+# sampling helpers
+
+def log(t, eps = 1e-20):
+    return torch.log(t.clamp(min = eps))
+
+def gumbel_noise(t):
+    noise = torch.zeros_like(t).uniform_(0, 1)
+    return -log(-log(noise))
+
+def gumbel_sample(t, temperature = 1., dim = -1, keepdim = True):
+    return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim = dim, keepdim = keepdim)
+
+def top_k(logits, thres = 0.9):
+    k = math.ceil((1 - thres) * logits.shape[-1])
+    val, ind = torch.topk(logits, k)
+    probs = torch.full_like(logits, float('-inf'))
+    probs.scatter_(-1, ind, val)
+    return probs
+
+def base_decoding(
+    net,
+    prompt: Tensor,
+    seq_len: int,
+    temperature = 1.,
+    filter_thres = 0.9,
+):
+    prompt_seq_len, out = prompt.shape[-1], prompt.clone()
+    sample_num_times = max(0, seq_len - prompt_seq_len)
+
+    for _ in range(sample_num_times):
+        logits = net(out)
+        logits = logits[:, -1]
+        logits = top_k(logits, thres = filter_thres)
+        sample = gumbel_sample(logits, temperature = temperature, dim = -1)
+
+        out = torch.cat((out, sample), dim = -1)
+
+    return out[..., prompt_seq_len:]
+
+# model
+
+if USE_DCA:
+    model = DCAGPT(
+        num_tokens = 256,
+        dim = 512,
+        depth = 6,
+        past_layers_k = 2  # the `k` value in their paper, which refers to how many layers in the past it looks, a la Denseformer    
+    )
+else:
+    model = GPT(
+        num_tokens = 256,
+        dim = 512,
+        depth = 6
+    )
+
+model = model.cuda()
+
+# prepare enwik8 data
+
+with gzip.open('./data/enwik8.gz') as file:
+    data = np.frombuffer(file.read(int(95e6)), dtype=np.uint8).copy()
+    np_train, np_valid = np.split(data, [int(90e6)])
+    data_train, data_val = torch.from_numpy(np_train), torch.from_numpy(np_valid)
+
+class TextSamplerDataset(Dataset):
+    def __init__(self, data, seq_len):
+        super().__init__()
+        self.data = data
+        self.seq_len = seq_len
+
+    def __len__(self):
+        return self.data.size(0) // self.seq_len
+
+    def __getitem__(self, index):
+        rand_start = torch.randint(0, self.data.size(0) - self.seq_len, (1,))
+        full_seq = self.data[rand_start : rand_start + self.seq_len + 1].long()
+        return full_seq.cuda()
+
+train_dataset = TextSamplerDataset(data_train, SEQ_LEN)
+val_dataset = TextSamplerDataset(data_val, SEQ_LEN)
+train_loader = DataLoader(train_dataset, batch_size = BATCH_SIZE)
+val_loader = DataLoader(val_dataset, batch_size = BATCH_SIZE)
+
+# optimizer
+
+optim = Adam(model.parameters(), lr = LEARNING_RATE)
+
+train_loader = cycle(train_loader)
+val_loader = cycle(val_loader)
+
+# training
+
+for i in tqdm.tqdm(range(NUM_BATCHES), mininterval = 10.0, desc = "training"):
+    model.train()
+
+    for _ in range(GRAD_ACCUM_EVERY):
+        data = next(train_loader)
+
+        loss = model(data, return_loss = True)
+
+        (loss / GRAD_ACCUM_EVERY).backward()
+
+    print(f"training loss: {loss.item():.3f}")
+
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+
+    optim.step()
+    optim.zero_grad()
+
+    if i % VALIDATE_EVERY == 0:
+        model.eval()
+        with torch.no_grad():
+            valid_data = next(val_loader)
+
+            loss = model(valid_data, return_loss = True)
+            print(f"validation loss: {loss.item():.3f}")
+
+    if i % GENERATE_EVERY == 0:
+        model.eval()
+
+        inp = random.choice(val_dataset)[:PRIME_LENGTH]
+        inp = inp.cuda()
+
+        prime = decode_tokens(inp)
+        print(f"\n{prime}\n")
+
+        prompt = inp[None, ...]
+
+        sampled = base_decoding(model, prompt, GENERATE_LENGTH)
+
+        base_decode_output = decode_tokens(sampled[0])
+
+        print(f"\n{base_decode_output}\n")