folk_mlm_test.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FolkMLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### imports and configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/master.zip -U\n",
    "# !pip3 install --force-reinstall torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\n",
    "# !pip install onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ignore warnings generate by pytorch and lightning\n",
    "import warnings\n",
    "warnings.simplefilter(\"ignore\")\n",
    "\n",
    "# # imports\n",
    "# import lightning.pytorch as pl\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import pickle\n",
    "import inspect\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import wandb\n",
    "\n",
    "from datetime import datetime\n",
    "\n",
    "from numpy.random import default_rng\n",
    "import math\n",
    "\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torch.utils.data import random_split\n",
    "\n",
    "import torchmetrics\n",
    "\n",
    "# from lightning.pytorch import LightningModule, LightningDataModule\n",
    "# from lightning.pytorch.loggers import WandbLogger\n",
    "# from lightning.pytorch.callbacks import StochasticWeightAveraging, RichProgressBar, RichModelSummary\n",
    "# from lightning.pytorch.strategies import ddp\n",
    "# from lightning.pytorch import Trainer\n",
    "\n",
    "from pytorch_lightning import LightningModule, LightningDataModule, Trainer\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "from pytorch_lightning.callbacks import StochasticWeightAveraging, RichProgressBar, RichModelSummary\n",
    "from pytorch_lightning.strategies import ddp\n",
    "\n",
    "from pathlib import Path\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class Config():\n",
    "#     def __init__(self) -> None:\n",
    "#         self.max_len = 256\n",
    "#         self.batch_size = 256\n",
    "#         self.epochs = 300\n",
    "#         self.weight_decay = 0.01\n",
    "#         self.lr = 1e-3\n",
    "#         self.d_model = 256\n",
    "#         self.d_hid = 4 * self.d_model\n",
    "#         self.nhead = 4 \n",
    "#         self.nlayers = 4\n",
    "#         self.dropout = 0.1\n",
    "#         self.lr_decay = True\n",
    "    \n",
    "#     def __repr__(self) -> str:\n",
    "#         for e in self.__dict__:\n",
    "#             print(str(e)+\": \"+str(self.__dict__[e]))\n",
    "#         return \"---\"\n",
    "    \n",
    "# args = Config()\n",
    "# args"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### datamodule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLMDataset(Dataset):\n",
    "    def __init__(self, dataset, block_size, TOKENS):\n",
    "        self.dataset = dataset      \n",
    "        self.stoi = { tk:i for i,tk in enumerate(TOKENS ) }\n",
    "        self.itos = { i:tk for i,tk in enumerate(TOKENS ) }\n",
    "        self.block_size = block_size\n",
    "        self.vocab_size = len(TOKENS)\n",
    "        self.IGNORE_TOKEN = -100 # as per pytorch crossentropy default\n",
    "        self.rng = default_rng()\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dataset) #- self.block_size\n",
    "\n",
    "\n",
    "    def mask_input(self, target, mask_size=0.15):\n",
    "        # the mask is as long as the block \n",
    "        # but only element before pad get masked\n",
    "        mask = np.zeros(self.block_size).astype(int)\n",
    "        seq_len = len(target[ target != self.stoi[\"<pad>\"]])\n",
    "        mask[ self.rng.choice(\n",
    "                np.arange(0,seq_len), \n",
    "                size=round(mask_size*seq_len), \n",
    "                replace=False)\n",
    "                ] = 1\n",
    "\n",
    "        # mask[target == self.stoi[\"<pad>\"]] = self.stoi[\"<pad>\"]\n",
    "        # always set mask for EOS so model can learn it better\n",
    "        # mask[target == self.stoi[\"</s>\"]] = 1\n",
    "        # also adding the first pad, if present\n",
    "        # if seq_len < 256: mask[seq_len] = 1\n",
    "        \n",
    "        # masking\n",
    "        input_seq = target.copy()\n",
    "        input_seq = np.where(mask==1, self.stoi[\"<mask>\"], input_seq)\n",
    "        # ignore unmasked\n",
    "        target = np.where(mask==0, self.IGNORE_TOKEN, target)\n",
    "\n",
    "        return input_seq, target        \n",
    "        \n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        target = np.array([self.stoi[s] for s in self.dataset[idx]])\n",
    "        # randomly sample from the masking schedule function gamma (See MASKGIT)\n",
    "        gamma = np.cos(np.random.uniform(0.01, np.pi/2))\n",
    "        input_seq, target = self.mask_input(target, mask_size=gamma)\n",
    "\n",
    "        input_seq = torch.tensor(input_seq, dtype=torch.long)\n",
    "        target = torch.tensor(target, dtype=torch.long)\n",
    "\n",
    "        return input_seq, target\n",
    "\n",
    "class MLMDataModule(LightningDataModule):\n",
    "    def __init__(self, data_dir: str = \"./datasets/\", batch_size: int = 32, max_len=256):\n",
    "        super().__init__()\n",
    "        self.data_dir = data_dir\n",
    "        self.batch_size = batch_size  \n",
    "        self.max_len = max_len\n",
    "        self.IGNORE_TOKEN = -100 # as per pytorch crossentropy default\n",
    "        self.rng = default_rng()\n",
    "\n",
    "    def load_tokens(self, tokens_path = \"./datasets/TOKENS_V4_arranged.pickle\"):\n",
    "        with open(tokens_path,\"rb\") as f:\n",
    "            print(\"Loading tokens:\", tokens_path)\n",
    "            tokens = pickle.load(f)\n",
    "        tokens = np.append(tokens,'<mask>')\n",
    "        # put <s> and </s> at the top and <pad> at the bottom\n",
    "        tokens = np.concatenate([tokens[2:3], tokens[0:2],tokens[3:]])\n",
    "        tokens = tokens[1:]\n",
    "        tokens = np.append(tokens,'<pad>')\n",
    "        return tokens\n",
    "\n",
    "    def load_dataframe(self, dataset_path = \"./datasets/df_v4.pickle\"):\n",
    "        print(\"Loading dataset:\", dataset_path)\n",
    "        tunes_df = pd.read_pickle(dataset_path).sort_values('length',ascending=False)\n",
    "        tunes_df[\"full_abc\"] = tunes_df['L'].map(str) + ' ' + tunes_df['M'].map(str) + ' ' + tunes_df['K'].map(str) + ' ' + tunes_df['abc'].map(str)\n",
    "        return tunes_df[tunes_df.length <= self.max_len-5] # adding <s> L M K </s>\n",
    "\n",
    "    def create_dataset(self, tunes_df):\n",
    "            df = tunes_df #[tunes_df.length <= self.max_len-5] # adding <s> L M K </s>    \n",
    "            strings = '<s> ' + df['L'].map(str) + '\\n' + df['M'].map(str) + '\\n' + df['K'].map(str) + '\\n' + df['abc'].map(str) + ' </s>'\n",
    "            strings = strings.apply(lambda x: x.split()[:])\n",
    "            strings = strings.values.reshape(-1,1)\n",
    "            dataset = np.asarray([self.padding(x) for x in strings[:]])\n",
    "            return dataset\n",
    "\n",
    "    #takes a numpy array as input\n",
    "    def padding(self, array):\n",
    "        array = array[0]\n",
    "        array = np.append(array,['<pad>']*(self.max_len-len(array) ))\n",
    "        assert len(array) == self.max_len\n",
    "        return np.array(array)\n",
    "        \n",
    "    def setup(self, stage):\n",
    "        print(\"setting up\", stage)\n",
    "        self.tunes_df = self.load_dataframe()\n",
    "        self.tokens = self.load_tokens()\n",
    "        self.stoi = { tk:i for i,tk in enumerate(self.tokens) }\n",
    "        self.itos = { i:tk for i,tk in enumerate(self.tokens) }\n",
    "        self.vocab_size = len(self.tokens)\n",
    "        self.train_df = self.tunes_df.sample(frac=0.9)\n",
    "        self.test_df = self.tunes_df.drop(self.train_df.index)\n",
    "        self.train_set = MLMDataset(self.create_dataset(self.train_df), self.max_len, self.tokens)\n",
    "        self.test_set = MLMDataset(self.create_dataset(self.test_df), self.max_len, self.tokens)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(self.train_set, batch_size=args.batch_size, num_workers=8)\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(self.test_set, batch_size=128, num_workers=8)\n",
    "\n",
    "#     def test_dataloader(self):\n",
    "#         return DataLoader(self.mnist_test, batch_size=self.batch_size)\n",
    "\n",
    "#     def predict_dataloader(self):\n",
    "#         return DataLoader(self.mnist_predict, batch_size=self.batch_size)\n",
    "\n",
    "#     def teardown(self, stage: str):\n",
    "#         # Used to clean-up when the run is finished\n",
    "#         ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "setting up train\n",
      "Loading dataset: ./datasets/df_v4.pickle\n",
      "Loading tokens: ./datasets/TOKENS_V4_arranged.pickle\n"
     ]
    }
   ],
   "source": [
    "datamodule = MLMDataModule()\n",
    "datamodule.setup(stage=\"train\")\n",
    "# datamodule.test_df.full_abc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'L:1/8M:6/8K:Cmaj|:E/2D/2|C2cBAG|ABcGFE|AFDGEC|B,CDD2E/2D/2|C2cBAG|ABcGFE|AFDGEC|DCB,C2:||:B,/2C/2|D3/2E/2DDEF|GABcGE|FGFEFE|CECB,AG|C2cBAG|ABcGFE|AFDGEC|DCB,C2:|'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "''.join(datamodule.train_df.loc[9731].full_abc.split(' '))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: '<s>',\n",
       " 1: '</s>',\n",
       " 2: 'L:1/8',\n",
       " 3: 'L:1/16',\n",
       " 4: 'M:12/8',\n",
       " 5: 'M:2/4',\n",
       " 6: 'M:3/2',\n",
       " 7: 'M:3/4',\n",
       " 8: 'M:4/4',\n",
       " 9: 'M:6/8',\n",
       " 10: 'M:9/8',\n",
       " 11: 'K:Cdor',\n",
       " 12: 'K:Cmaj',\n",
       " 13: 'K:Cmin',\n",
       " 14: 'K:Cmix',\n",
       " 15: '[',\n",
       " 16: ']',\n",
       " 17: '|',\n",
       " 18: '|1',\n",
       " 19: '|2',\n",
       " 20: '|:',\n",
       " 21: ':|',\n",
       " 22: '(3',\n",
       " 23: '/2',\n",
       " 24: '/4',\n",
       " 25: '/8',\n",
       " 26: '7/2',\n",
       " 27: '3/2',\n",
       " 28: '3/4',\n",
       " 29: '5/2',\n",
       " 30: '2',\n",
       " 31: '3',\n",
       " 32: '4',\n",
       " 33: '5',\n",
       " 34: '6',\n",
       " 35: '7',\n",
       " 36: '8',\n",
       " 37: '9',\n",
       " 38: '12',\n",
       " 39: '16',\n",
       " 40: '<',\n",
       " 41: '>',\n",
       " 42: 'A,',\n",
       " 43: 'B,',\n",
       " 44: 'C,',\n",
       " 45: 'D,',\n",
       " 46: 'E,',\n",
       " 47: 'F,',\n",
       " 48: 'G,',\n",
       " 49: 'A',\n",
       " 50: 'B',\n",
       " 51: 'C',\n",
       " 52: 'D',\n",
       " 53: 'E',\n",
       " 54: 'F',\n",
       " 55: 'G',\n",
       " 56: 'a',\n",
       " 57: 'b',\n",
       " 58: 'c',\n",
       " 59: 'd',\n",
       " 60: 'e',\n",
       " 61: 'f',\n",
       " 62: 'g',\n",
       " 63: \"a'\",\n",
       " 64: \"b'\",\n",
       " 65: \"c'\",\n",
       " 66: \"d'\",\n",
       " 67: \"e'\",\n",
       " 68: \"f'\",\n",
       " 69: \"g'\",\n",
       " 70: '=A,',\n",
       " 71: '=B,',\n",
       " 72: '=C,',\n",
       " 73: '=E,',\n",
       " 74: '=F,',\n",
       " 75: '=G,',\n",
       " 76: '=A',\n",
       " 77: '=B',\n",
       " 78: '=C',\n",
       " 79: '=D',\n",
       " 80: '=E',\n",
       " 81: '=F',\n",
       " 82: '=G',\n",
       " 83: '=a',\n",
       " 84: '=b',\n",
       " 85: '=c',\n",
       " 86: '=d',\n",
       " 87: '=e',\n",
       " 88: '=f',\n",
       " 89: '=g',\n",
       " 90: \"=c'\",\n",
       " 91: \"=e'\",\n",
       " 92: \"=f'\",\n",
       " 93: '^A,',\n",
       " 94: '^C,',\n",
       " 95: '^F,',\n",
       " 96: '^G,',\n",
       " 97: '^A',\n",
       " 98: '^C',\n",
       " 99: '^D',\n",
       " 100: '^F',\n",
       " 101: '^G',\n",
       " 102: '^a',\n",
       " 103: '^c',\n",
       " 104: '^d',\n",
       " 105: '^f',\n",
       " 106: '^g',\n",
       " 107: \"^c'\",\n",
       " 108: \"^f'\",\n",
       " 109: '_A,',\n",
       " 110: '_B,',\n",
       " 111: '_D',\n",
       " 112: '_E,',\n",
       " 113: '_A',\n",
       " 114: '_B',\n",
       " 115: '_C',\n",
       " 116: '_E',\n",
       " 117: '_G',\n",
       " 118: '_a',\n",
       " 119: '_b',\n",
       " 120: '_c',\n",
       " 121: '_d',\n",
       " 122: '_e',\n",
       " 123: '_g',\n",
       " 124: \"_d'\",\n",
       " 125: \"_e'\",\n",
       " 126: 'z',\n",
       " 127: '<mask>',\n",
       " 128: '<pad>'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datamodule.itos"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from pytorch lightning example on transformers\n",
    "# https://pytorch-lightning.readthedocs.io/en/stable/notebooks/course_UvA-DL/05-transformers-and-MH-attention.html\n",
    "class CosineWarmupScheduler(optim.lr_scheduler._LRScheduler):\n",
    "    def __init__(self, optimizer, warmup, max_iters, min_lr):\n",
    "        self.warmup = warmup\n",
    "        self.max_num_iters = max_iters\n",
    "        self.min_lr = min_lr\n",
    "        super().__init__(optimizer)\n",
    "\n",
    "    def get_lr(self):\n",
    "        lr_factor = self.get_lr_factor(epoch=self.last_epoch)\n",
    "        lrs = [base_lr * lr_factor for base_lr in self.base_lrs]\n",
    "        lrs = [lr if (lr >= self.min_lr) else self.min_lr for lr in lrs]\n",
    "        return lrs\n",
    "\n",
    "    def get_lr_factor(self, epoch):\n",
    "        lr_factor = 0.5 * (1 + torch.cos(torch.pi * epoch / self.max_num_iters))\n",
    "        if epoch <= self.warmup:\n",
    "            lr_factor *= epoch * 1.0 / self.warmup\n",
    "        if epoch > self.max_num_iters:\n",
    "            return 0.0\n",
    "        return lr_factor\n",
    "\n",
    "\n",
    "class SelfAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    A vanilla multi-head masked self-attention layer with a projection at the end.\n",
    "    It is possible to use torch.nn.MultiheadAttention here but I am including an\n",
    "    explicit implementation here to show that there is nothing too scary here.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, n_embd, n_head, attn_pdrop=0.0, resid_pdrop=0.0):\n",
    "        super().__init__()\n",
    "        assert n_embd % n_head == 0\n",
    "        # key, query, value projections for all heads\n",
    "        self.key = nn.Linear(n_embd, n_embd)\n",
    "        self.query = nn.Linear(n_embd, n_embd)\n",
    "        self.value = nn.Linear(n_embd, n_embd)\n",
    "        # regularization\n",
    "        self.attn_drop = nn.Dropout(attn_pdrop)\n",
    "        self.resid_drop = nn.Dropout(resid_pdrop)\n",
    "        # output projection\n",
    "        self.proj = nn.Linear(n_embd, n_embd)\n",
    "        \n",
    "        # # causal mask to ensure that attention is only applied to the left in the input sequence\n",
    "        # self.register_buffer(\"mask\", torch.tril(torch.ones(config.block_size, config.block_size))\n",
    "        #                              .view(1, 1, config.block_size, config.block_size))\n",
    "        \n",
    "        self.n_head = n_head\n",
    "\n",
    "    def forward(self, x, src_key_padding_mask):\n",
    "        # B is the batch size, \n",
    "        # T is the sequence length, \n",
    "        # C is the dimensionality of the embedding (n_embd).\n",
    "        B, T, C = x.size()\n",
    "\n",
    "        # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n",
    "        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
    "\n",
    "        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)\n",
    "        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n",
    "        att = att.masked_fill(src_key_padding_mask, float('-inf'))\n",
    "        att = F.softmax(att, dim=-1)\n",
    "        att = self.attn_drop(att)\n",
    "        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)\n",
    "        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n",
    "\n",
    "        # output projection\n",
    "        y = self.resid_drop(self.proj(y))\n",
    "        return y\n",
    "\n",
    "class Block(nn.Module):\n",
    "    \"\"\" an unassuming Transformer block \"\"\"\n",
    "\n",
    "    def __init__(self, n_embd, n_head, attn_pdrop=0.0, resid_pdrop=0.0):\n",
    "        super().__init__()\n",
    "        self.ln1 = nn.LayerNorm(n_embd,)\n",
    "        self.ln2 = nn.LayerNorm(n_embd, )\n",
    "        self.attn = SelfAttention(n_embd, n_head, attn_pdrop=attn_pdrop, resid_pdrop=resid_pdrop)\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(n_embd, 4 * n_embd),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(4 * n_embd, n_embd),\n",
    "            nn.Dropout(resid_pdrop),\n",
    "        )\n",
    "\n",
    "    def forward(self, x, src_key_padding_mask):\n",
    "        x = x + self.attn( self.ln1(x), src_key_padding_mask)\n",
    "        x = x + self.mlp( self.ln2(x) )\n",
    "        return x\n",
    "\n",
    "class MaskedLM(LightningModule):\n",
    "    def __init__(self, PAD_TOKEN, ntoken, IGNORE_TOKEN=-100, d_model=64, nhead=8, d_hid=128, nlayers=6, dropout=0.0, lr=1e-4, lr_sched=False, max_len=256, weight_decay=1e-2, custom_block=False):\n",
    "        super().__init__()\n",
    "        self.model_type = 'Transformer'\n",
    "        \n",
    "        self.d_model = d_model\n",
    "        self.PAD_TOKEN = PAD_TOKEN\n",
    "        self.IGNORE_TOKEN = IGNORE_TOKEN\n",
    "        self.lr = lr\n",
    "        self.lr_sched = lr_sched\n",
    "        self.max_len = max_len\n",
    "        self.weight_decay = weight_decay\n",
    "        \n",
    "        self.embedding = nn.Embedding(ntoken, d_model, padding_idx=PAD_TOKEN)\n",
    "        self.learned_pos = nn.Embedding(max_len, d_model)\n",
    "        \n",
    "        if custom_block:\n",
    "            self.transformer_encoder = nn.ModuleList([Block(n_embd=d_model, n_head=nhead) for _ in range(nlayers)])\n",
    "        else:\n",
    "            self.transformer_encoder = nn.TransformerEncoder(\n",
    "                nn.TransformerEncoderLayer(\n",
    "                    d_model=d_model, \n",
    "                    nhead=nhead, \n",
    "                    dim_feedforward=d_hid, \n",
    "                    dropout=dropout,\n",
    "                    batch_first=True, norm_first=True, \n",
    "                    activation='gelu',\n",
    "                ), \n",
    "                num_layers=nlayers\n",
    "            )\n",
    "        \n",
    "        self.ln_f = nn.LayerNorm(d_model) \n",
    "        self.decoder = nn.Linear(d_model, ntoken, bias=False)\n",
    "        self.dropout = nn.Dropout(p=0.0)\n",
    "\n",
    "        # weight tying see karphaty\n",
    "        self.embedding.weight = self.decoder.weight\n",
    "        \n",
    "        self.apply(self._init_weights)\n",
    "        for pn, p in self.named_parameters():\n",
    "            if pn.endswith('linear2.weight'):\n",
    "                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * nlayers))\n",
    "\n",
    "        \n",
    "        self.save_hyperparameters()\n",
    "        \n",
    "    def _init_weights(self, module):\n",
    "        if isinstance(module, nn.Linear):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "            if module.bias is not None:\n",
    "                torch.nn.init.zeros_(module.bias)\n",
    "        elif isinstance(module, nn.Embedding):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "    \n",
    "    def configure_optimizers(self, betas=[0.9,0.95]):\n",
    "            \"\"\"\n",
    "            This long function is unfortunately doing something very simple and is being very defensive:\n",
    "            We are separating out all parameters of the model into two buckets: those that will experience\n",
    "            weight decay for regularization and those that won't (biases, and layernorm/embedding weights).\n",
    "            We are then returning the PyTorch optimizer object.\n",
    "            \"\"\"\n",
    "\n",
    "            # separate out all parameters to those that will and won't experience regularizing weight decay\n",
    "            decay = set()\n",
    "            no_decay = set()\n",
    "            whitelist_weight_modules = (nn.Linear, nn.Parameter)\n",
    "            blacklist_weight_modules = (nn.LayerNorm, torch.nn.Embedding)\n",
    "            c = 0\n",
    "            for mn, m in self.named_modules():\n",
    "                for pn, p in m.named_parameters():\n",
    "                    c += 1\n",
    "                    fpn = '%s.%s' % (mn, pn) if mn else pn # full param name\n",
    "                    # random note: because named_modules and named_parameters are recursive\n",
    "                    # we will see the same tensors p many many times. but doing it this way\n",
    "                    # allows us to know which parent module any tensor p belongs to...\n",
    "                    if fpn.endswith('bias'):\n",
    "                        # all biases will not be decayed\n",
    "                        no_decay.add(fpn)\n",
    "                    elif fpn.endswith('weight') and isinstance(m, blacklist_weight_modules):\n",
    "                        # weights of blacklist modules will NOT be weight decayed\n",
    "                        no_decay.add(fpn)\n",
    "                    elif fpn.endswith('weight') and isinstance(m, whitelist_weight_modules):\n",
    "                        # weights of whitelist modules will be weight decayed\n",
    "                        decay.add(fpn)\n",
    "                    # a bit of an hack to make it work, I'll look into it more\n",
    "                    elif fpn.endswith('in_proj_weight'):\n",
    "                        decay.add(fpn)\n",
    "\n",
    "            # subtle: 'transformer.wte.weight' and 'lm_head.weight' are tied, so they\n",
    "            # will appear in the no_decay and decay sets respectively after the above.\n",
    "            # In addition, because named_parameters() doesn't return duplicates, it\n",
    "            # will only return the first occurence, key'd by 'transformer.wte.weight', below.\n",
    "            # so let's manually remove 'lm_head.weight' from decay set. This will include\n",
    "            # this tensor into optimization via transformer.wte.weight only, and not decayed.\n",
    "            decay.remove('decoder.weight')\n",
    "\n",
    "            # validate that we considered every parameter\n",
    "            param_dict = {pn: p for pn, p in self.named_parameters()}\n",
    "\n",
    "            inter_params = decay & no_decay\n",
    "            union_params = decay | no_decay\n",
    "            assert len(inter_params) == 0, \"parameters %s made it into both decay/no_decay sets!\" % (str(inter_params), )\n",
    "            assert len(param_dict.keys() - union_params) == 0, \"parameters %s were not separated into either decay/no_decay set!\" \\\n",
    "                                                        % (str(param_dict.keys() - union_params), )\n",
    "\n",
    "            # create the pytorch optimizer object\n",
    "            optim_groups = [\n",
    "                {\"params\": [param_dict[pn] for pn in sorted(list(decay))], \"weight_decay\": self.weight_decay},\n",
    "                {\"params\": [param_dict[pn] for pn in sorted(list(no_decay))], \"weight_decay\": 0.0},\n",
    "            ]\n",
    "\n",
    "            # new PyTorch nightly has a new 'fused' option for AdamW that is much faster\n",
    "            use_fused = (self.device == 'cuda') and ('fused' in inspect.signature(torch.optim.AdamW).parameters)\n",
    "            print(f\"using fused AdamW: {self.device} and {'fused' in inspect.signature(torch.optim.AdamW).parameters} !\")\n",
    "            extra_args = dict(fused=True) if use_fused else dict()\n",
    "            \n",
    "            optimizer = torch.optim.AdamW(optim_groups, lr=self.lr, fused=False, betas=betas, **extra_args)\n",
    "\n",
    "            # We don't return the lr scheduler because we need to apply it per iteration, not per epoch\n",
    "            if self.lr_sched:\n",
    "                self.lr_scheduler = CosineWarmupScheduler(\n",
    "                    optimizer, \n",
    "                    warmup=150, #self.hparams.warmup, \n",
    "                    max_iters=10000,\n",
    "                    min_lr=1e-4\n",
    "                )\n",
    "            \n",
    "            return optimizer\n",
    "\n",
    "    def optimizer_step(self, *args, **kwargs):\n",
    "            super().optimizer_step(*args, **kwargs)\n",
    "            if self.lr_sched: \n",
    "                self.lr_scheduler.step()  # Step per iteration\n",
    "\n",
    "    def forward(self, x):\n",
    "        attn_mask = (x ==  self.PAD_TOKEN) # True means no attention   \n",
    "\n",
    "        token_embeddings = self.embedding(x) # each index maps to a (learnable) vector\n",
    "        # from nanogpt\n",
    "        pos = torch.arange(0, self.max_len, device=x.device).view(1,self.max_len) # shape (1, t)\n",
    "        position_embeddings = self.learned_pos(pos)\n",
    "        x = self.dropout(token_embeddings + position_embeddings)\n",
    "\n",
    "        x = self.transformer_encoder(x, src_key_padding_mask = attn_mask)\n",
    "        x = self.decoder(self.ln_f(x))\n",
    "\n",
    "        return x\n",
    "    \n",
    "    def training_step(self, train_batch, batch_idx):\n",
    "        x, y = train_batch\n",
    "        o = self.forward(x)\n",
    "        # y = torch.where(y==self.PAD_TOKEN,self.IGNORE_TOKEN,y)\n",
    "        loss = F.cross_entropy(\n",
    "            o.permute(0,2,1), \n",
    "            y, \n",
    "            ignore_index=-100, \n",
    "            reduction=\"mean\", \n",
    "            # weight=TOKENS_WEIGHTS.to(self.device)\n",
    "            )\n",
    "        \n",
    "        self.log('train/loss', loss, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)\n",
    "        \n",
    "        pred = o.argmax(-1)\n",
    "        acc_top1 = torchmetrics.functional.accuracy(\n",
    "            o[torch.where(y!=self.IGNORE_TOKEN)], \n",
    "            y[torch.where(y!=self.IGNORE_TOKEN)], \n",
    "            task=\"multiclass\", \n",
    "            num_classes=129, \n",
    "            top_k=1\n",
    "        )\n",
    "        acc_top5 = torchmetrics.functional.accuracy(\n",
    "            o[torch.where(y!=self.IGNORE_TOKEN)], \n",
    "            y[torch.where(y!=self.IGNORE_TOKEN)], \n",
    "            task=\"multiclass\", \n",
    "            num_classes=129, \n",
    "            top_k=5\n",
    "        )\n",
    "\n",
    "        # accuracy = pred[torch.where(y!=self.IGNORE_TOKEN)] == y[torch.where(y!=self.IGNORE_TOKEN)]\n",
    "        # accuracy = torch.mean(accuracy, dim=-1, dtype=float)\n",
    "        \n",
    "        self.log(\"train/acc/top_1\",acc_top1, prog_bar=True, sync_dist=True)\n",
    "        self.log(\"train/acc/top_5\",acc_top5, prog_bar=True, sync_dist=True)\n",
    "\n",
    "        if self.lr_sched:\n",
    "            self.log(\"lr\",self.lr_scheduler.get_last_lr()[0] , prog_bar=True)\n",
    "        else:\n",
    "            self.log(\"lr\",self.lr)\n",
    "        \n",
    "        return {\"loss\": loss, \"pred\": pred, \"acc\":acc_top1}\n",
    "\n",
    "    def validation_step(self, val_batch, batch_idx):\n",
    "        x, y = val_batch\n",
    "        o = self.forward(x)\n",
    "        # y = torch.where(y==self.PAD_TOKEN,self.IGNORE_TOKEN,y)\n",
    "        loss = F.cross_entropy(\n",
    "            o.permute(0,2,1), \n",
    "            y, \n",
    "            ignore_index=-100, \n",
    "            # weight=TOKENS_WEIGHTS.to(self.device)\n",
    "            )\n",
    "        self.log('valid/loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)\n",
    "        pred = o.argmax(-1)\n",
    "        acc_top1 = torchmetrics.functional.accuracy(\n",
    "            o[torch.where(y!=self.IGNORE_TOKEN)], \n",
    "            y[torch.where(y!=self.IGNORE_TOKEN)], \n",
    "            task=\"multiclass\", \n",
    "            num_classes=129, \n",
    "            top_k=1\n",
    "        )\n",
    "        acc_top5 = torchmetrics.functional.accuracy(\n",
    "            o[torch.where(y!=self.IGNORE_TOKEN)], \n",
    "            y[torch.where(y!=self.IGNORE_TOKEN)], \n",
    "            task=\"multiclass\", \n",
    "            num_classes=129, \n",
    "            top_k=5\n",
    "        )\n",
    "\n",
    "        # accuracy = pred[torch.where(y!=self.IGNORE_TOKEN)] == y[torch.where(y!=self.IGNORE_TOKEN)]\n",
    "        # accuracy = torch.mean(accuracy, dim=-1, dtype=float)\n",
    "        \n",
    "        self.log(\"valid/acc/top_1\",acc_top1, prog_bar=True, sync_dist=True)\n",
    "        self.log(\"valid/acc/top_5\",acc_top5, prog_bar=True, sync_dist=True)\n",
    "        return {\"loss\": loss, \"pred\": pred, \"acc\":acc_top1}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### generation functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):\n",
    "    assert logits.dim() == 1  # batch size 1 for now - could be updated for more but the code would be less clear\n",
    "    top_k = min(top_k, logits.size(-1))  # Safety check\n",
    "    if top_k > 0:\n",
    "        # Remove all tokens with a probability less than the last token of the top-k\n",
    "        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]\n",
    "        logits[indices_to_remove] = filter_value\n",
    "\n",
    "    if top_p > 0.0:\n",
    "        sorted_logits, sorted_indices = torch.sort(logits, descending=True)\n",
    "        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)\n",
    "\n",
    "        # Remove tokens with cumulative probability above the threshold\n",
    "        sorted_indices_to_remove = cumulative_probs > top_p\n",
    "        # Shift the indices to the right to keep also the first token above the threshold\n",
    "        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()\n",
    "        sorted_indices_to_remove[..., 0] = 0\n",
    "\n",
    "        indices_to_remove = sorted_indices[sorted_indices_to_remove]\n",
    "        logits[indices_to_remove] = filter_value\n",
    "    return logits\n",
    "\n",
    "def generate_reels(datamodule, n=10, from_structure=True, sample_length=False, order=\"random\", verbose=False, top_k=None, top_p=0.9, savepath=None):\n",
    "    reels = []\n",
    "    for idx,e in datamodule.test_df[datamodule.test_df[\"M\"] == \"M:4/4\"].sample(n).iterrows():\n",
    "        if from_structure:\n",
    "            input_seq = '<s> ' + e.full_abc + ' </s>' #add SOS and EOS\n",
    "            input_seq = input_seq.split() \n",
    "            input_seq = np.array([ t if (datamodule.stoi[t] < 42 or datamodule.stoi[t] >= 126)  else \"<mask>\" for t in input_seq ]) # encode non-structural tokens\n",
    "        else:\n",
    "            M = ['M:4/4'] #list(tunes_df[\"M\"].sample(1).values)\n",
    "            K = [np.random.choice(['K:Cmaj','K:Cmin','K:Cmix','K:Cdor'])] # list(tunes_df[\"K\"].sample(1).values)\n",
    "            if sample_length:\n",
    "                input_seq = [\"<s>\", \"L:1/8\"] + M + K + [\"<mask>\"]*e.length  + [\"</s>\"]\n",
    "            else:\n",
    "                input_seq = [\"<s>\", \"L:1/8\"] + M + K + [\"<mask>\"]*252\n",
    "        \n",
    "        seq_len = len(input_seq)\n",
    "        # transform sequence to indices\n",
    "        input_seq = np.array([datamodule.stoi[t] for t in input_seq])\n",
    "        input_seq = np.pad(input_seq,(0,256-seq_len),\"constant\", constant_values=datamodule.stoi[\"<pad>\"]).reshape(1,-1) # pad\n",
    "        masked = list(np.where(input_seq[0]==datamodule.stoi[\"<mask>\"])[0]) # get \n",
    "        \n",
    "        if order == \"l2r\":\n",
    "            masked = sorted(masked)\n",
    "        elif order == \"r2l\":\n",
    "            masked = sorted(masked, reverse=True)\n",
    "        elif order == \"random\":\n",
    "            np.random.shuffle(masked) # shuffle the order of random tokens\n",
    "\n",
    "        while len(masked) != 0:\n",
    "            i = masked.pop(0)\n",
    "            # print(\"prediction index:\",i)\n",
    "            logits = model(\n",
    "                torch.IntTensor(input_seq), \n",
    "            )\n",
    "            \n",
    "            if top_k:\n",
    "                logits[0,i] = top_k_top_p_filtering(logits[0,i], top_k=top_k, filter_value=-torch.inf)\n",
    "            \n",
    "            elif top_p:\n",
    "                logits[0,i] = top_k_top_p_filtering(logits[0,i], top_p=top_p, filter_value=-torch.inf)\n",
    "                if verbose: print(sum(logits[0,i] > -torch.inf))\n",
    "            \n",
    "            probs = torch.nn.functional.softmax(logits,dim=-1)\n",
    "            \n",
    "            dist = torch.distributions.Categorical(probs[0,i])\n",
    "            sampled = dist.sample()\n",
    "            input_seq[0,i] = sampled\n",
    "            \n",
    "            if verbose:\n",
    "                print( \"\".join([datamodule.itos[t] for t in input_seq[0]]).replace(\"<mask>\",\"_\"))\n",
    "                    \n",
    "        strout = [datamodule.itos[t] for t in input_seq[0] if t != 128]\n",
    "        # if verbose: print(\" \".join(strout).replace(\"<mask>\",\"_\"))\n",
    "        \n",
    "        strout.insert(2,\"\\n\")\n",
    "        strout.insert(4,\"\\n\")\n",
    "        strout.insert(6,\"\\n\")\n",
    "        if from_structure:\n",
    "            strout = \"X:{}\\nT:based on test {}\\n{}\".format(str(idx),str(idx),\"\".join(strout[1:-1]))\n",
    "        else:\n",
    "            strout = strout = \"X:{}\\nT:{}\\n{}\".format('999'+str(idx),\"random\",\"\".join(strout[1:-1]))\n",
    "            \n",
    "        reels.append(strout)\n",
    "            \n",
    "    #save the outputs in a file\n",
    "    if savepath:\n",
    "        file = Path(savepath)\n",
    "        file.parent.mkdir(parents=True, exist_ok=True)\n",
    "        file.write_text(\"\\n\\n\".join(reels))\n",
    "\n",
    "    return reels\n",
    "\n",
    "def generate_jigs(datamodule, n=10, from_structure=True, sample_length=True, order=\"random\", verbose=False, top_k=None, top_p=0.9, savepath=None):\n",
    "    jigs = []\n",
    "    for idx,e in datamodule.test_df[datamodule.test_df[\"M\"] == \"M:6/8\"].sample(n).iterrows():\n",
    "        if from_structure:\n",
    "            input_seq = '<s> ' + e.full_abc + ' </s>' #add SOS and EOS\n",
    "            input_seq = input_seq.split() \n",
    "            input_seq = np.array([ t if (datamodule.stoi[t] < 42 or datamodule.stoi[t] >= 126)  else \"<mask>\" for t in input_seq ]) # encode non-structural tokens\n",
    "        else:\n",
    "            M = ['M:6/8'] #list(tunes_df[\"M\"].sample(1).values)\n",
    "            K = [np.random.choice(['K:Cmaj','K:Cmin','K:Cmix','K:Cdor'])] # list(tunes_df[\"K\"].sample(1).values)\n",
    "            if sample_length:\n",
    "                input_seq = [\"<s>\", \"L:1/8\"] + M + K + [\"<mask>\"]*e.length  + [\"</s>\"]\n",
    "            else:\n",
    "                input_seq = [\"<s>\", \"L:1/8\"] + M + K + [\"<mask>\"]*252\n",
    "        \n",
    "        seq_len = len(input_seq)\n",
    "        # transform sequence to indices\n",
    "        input_seq = np.array([datamodule.stoi[t] for t in input_seq])\n",
    "        input_seq = np.pad(input_seq,(0,256-seq_len),\"constant\", constant_values=datamodule.stoi[\"<pad>\"]).reshape(1,-1) # pad\n",
    "        masked = list(np.where(input_seq[0]==datamodule.stoi[\"<mask>\"])[0]) # get \n",
    "\n",
    "        if order == \"l2r\":\n",
    "            masked = sorted(masked)\n",
    "        elif order == \"r2l\":\n",
    "            masked = sorted(masked, reverse=True)\n",
    "        elif order == \"random\":\n",
    "            np.random.shuffle(masked) # shuffle the order of random tokens\n",
    "            \n",
    "        while len(masked) != 0:\n",
    "            i = masked.pop(0)\n",
    "            # print(\"prediction index:\",i)\n",
    "            logits = model(\n",
    "                torch.IntTensor(input_seq), \n",
    "            )\n",
    "            \n",
    "            if top_k:\n",
    "                logits[0,i] = top_k_top_p_filtering(logits[0,i], top_k=top_k, filter_value=-torch.inf)\n",
    "            \n",
    "            elif top_p:\n",
    "                logits[0,i] = top_k_top_p_filtering(logits[0,i], top_p=top_p, filter_value=-torch.inf)\n",
    "                # if verbose: print(sum(logits[0,i] > -torch.inf))\n",
    "            \n",
    "            probs = torch.nn.functional.softmax(logits,dim=-1)\n",
    "            \n",
    "            dist = torch.distributions.Categorical(probs[0,i])\n",
    "            sampled = dist.sample()\n",
    "            input_seq[0,i] = sampled\n",
    "            \n",
    "            if verbose:\n",
    "                print( \"\".join([datamodule.itos[t] for t in input_seq[0]]).replace(\"<mask>\",\"_\"))\n",
    "                    \n",
    "        strout = [datamodule.itos[t] for t in input_seq[0] if t != 128]\n",
    "        # if verbose: print(\" \".join(strout).replace(\"<mask>\",\"_\"))\n",
    "        \n",
    "        strout.insert(2,\"\\n\")\n",
    "        strout.insert(4,\"\\n\")\n",
    "        strout.insert(6,\"\\n\")\n",
    "        if from_structure:\n",
    "            strout = \"X:{}\\nT:based on test {}\\n{}\".format(str(idx),str(idx),\"\".join(strout[1:-1]))\n",
    "        else:\n",
    "            strout = \"X:{}\\nT:{}\\n{}\".format('999'+str(idx),\"random\",\"\".join(strout[1:-1]))\n",
    "        jigs.append(strout)\n",
    "            \n",
    "    #save the outputs in a file\n",
    "    if savepath:\n",
    "        file = Path(savepath)\n",
    "        file.parent.mkdir(parents=True, exist_ok=True)\n",
    "        file.write_text(\"\\n\\n\".join(jigs))\n",
    "\n",
    "    return jigs\n",
    "\n",
    "def generate_autoregressive(prompt, datamodule, n=10, order=\"l2r\", verbose=False, temperature=1.0, top_k=None, top_p=None, savepath=None, early_stop=True,):\n",
    "    tunes = []\n",
    "    \n",
    "    for idx in range(n):\n",
    "        if prompt == None:\n",
    "            M = [np.random.choice(['M:4/4','M:6/8'])]\n",
    "            K = [np.random.choice(['K:Cmaj','K:Cmin','K:Cmix','K:Cdor'])] \n",
    "            input_seq = [\"<s>\", \"L:1/8\"] + M + K + [\"<mask>\"]*(256-4)\n",
    "        else:\n",
    "            input_seq = prompt.split(\" \") + [\"<mask>\"]*(256-len(prompt.split(\" \")))\n",
    "        \n",
    "        seq_len = len(input_seq)\n",
    "        assert seq_len == 256\n",
    "        # transform sequence to indices\n",
    "        input_seq = np.array([datamodule.stoi[t] for t in input_seq])\n",
    "        input_seq = np.pad(input_seq,(0,256-seq_len),\"constant\", constant_values=datamodule.stoi[\"<pad>\"]).reshape(1,-1) # pad\n",
    "        masked = list(np.where(input_seq[0]==datamodule.stoi[\"<mask>\"])[0]) # get \n",
    "        masked = sorted(masked)\n",
    "\n",
    "        if order == \"l2r\":\n",
    "            masked = sorted(masked)\n",
    "        elif order == \"r2l\":\n",
    "            masked = sorted(masked, reverse=True)\n",
    "        elif order == \"random\":\n",
    "            np.random.shuffle(masked) # shuffle the order of random tokens\n",
    "        \n",
    "        while len(masked) != 0:\n",
    "            i = masked.pop(0)\n",
    "            # print(\"prediction index:\",i)\n",
    "            logits = model(torch.IntTensor(input_seq))\n",
    "            \n",
    "            logits = logits[0,i] / temperature\n",
    "            \n",
    "            if top_k:\n",
    "                logits = top_k_top_p_filtering(logits, top_k=top_k, filter_value=-torch.inf)\n",
    "            \n",
    "            elif top_p:\n",
    "                logits = top_k_top_p_filtering(logits, top_p=top_p, filter_value=-torch.inf)\n",
    "                # if verbose: print(sum(logits[0,i] > -torch.inf))\n",
    "            \n",
    "            probs = torch.nn.functional.softmax(logits,dim=-1)\n",
    "            dist = torch.distributions.Categorical(probs)\n",
    "            sampled = dist.sample()\n",
    "            input_seq[0,i] = sampled\n",
    "            if verbose:\n",
    "                print( \"\".join([datamodule.itos[t] for t in input_seq[0]]).replace(\"<mask>\",\"_\").replace(\"<pad>\",\"~\"))\n",
    "            # once we get the end token we exit\n",
    "            if sampled == datamodule.stoi[\"</s>\"] and early_stop:\n",
    "                if verbose: print(\"eos i:\", i)\n",
    "                break\n",
    "\n",
    "        strout = [datamodule.itos[t] for t in input_seq[0] if t != datamodule.stoi[\"<pad>\"]]\n",
    "        strout_len = len(strout)\n",
    "        # if verbose: print(\"length:\", strout_len)\n",
    "        # if verbose: print(\" \".join(strout).replace(\"<mask>\",\"_\"))\n",
    "        \n",
    "        strout.insert(2,\"\\n\")\n",
    "        strout.insert(4,\"\\n\")\n",
    "        strout.insert(6,\"\\n\")\n",
    "        \n",
    "        strout = \"X:{}\\nT:{}\\nN:tokens={}\\n{}\".format('999'+str(idx),\"random autoregressive\",strout_len,\"\".join(strout[:]))\n",
    "        tunes.append(strout)\n",
    "            \n",
    "    #save the outputs in a file\n",
    "    if savepath:\n",
    "        file = Path(savepath)\n",
    "        file.parent.mkdir(parents=True, exist_ok=True)\n",
    "        file.write_text(\"\\n\\n\".join(tunes))\n",
    "\n",
    "    return tunes\n",
    "\n",
    "def fill_masked(prompt, datamodule, title=\"fillmask\", verbose=False, temperature=1.0, top_k=None, top_p=None, savepath=None):\n",
    " \n",
    "    input_seq = prompt.split(\" \")\n",
    "    seq_len = len(input_seq)\n",
    "    \n",
    "    # transform sequence to indices\n",
    "    input_seq = np.array([datamodule.stoi[t] for t in input_seq])\n",
    "    input_seq = np.pad(input_seq,(0,256-seq_len),\"constant\", constant_values=datamodule.stoi[\"<pad>\"]).reshape(1,-1) # pad\n",
    "    \n",
    "    masked = list(np.where(input_seq[0]==datamodule.stoi[\"<mask>\"])[0]) # get \n",
    "    masked = sorted(masked)\n",
    "\n",
    "    while len(masked) != 0:\n",
    "        i = masked.pop(0)\n",
    "        # print(\"prediction index:\",i)\n",
    "        logits = model(torch.IntTensor(input_seq))\n",
    "\n",
    "        logits = logits[0,i] / temperature\n",
    "\n",
    "        if top_k:\n",
    "            logits = top_k_top_p_filtering(logits, top_k=top_k, filter_value=-torch.inf)\n",
    "\n",
    "        elif top_p:\n",
    "            logits = top_k_top_p_filtering(logits, top_p=top_p, filter_value=-torch.inf)\n",
    "            # if verbose: print(sum(logits[0,i] > -torch.inf))\n",
    "\n",
    "        probs = torch.nn.functional.softmax(logits,dim=-1)\n",
    "        dist = torch.distributions.Categorical(probs)\n",
    "        sampled = dist.sample()\n",
    "        input_seq[0,i] = sampled\n",
    "        if verbose:\n",
    "            print( \"\".join([datamodule.itos[t] for t in input_seq[0] if t != datamodule.stoi[\"<pad>\"] ]).replace(\"<mask>\",\"_\"))\n",
    "        # once we get the end token we exit\n",
    "        if sampled == datamodule.stoi[\"</s>\"]:\n",
    "            if verbose: print(\"eos i:\", i)\n",
    "            break\n",
    "\n",
    "    strout = [datamodule.itos[t] if t != datamodule.stoi[\"<pad>\"] else \"%\" for t in input_seq[0]]\n",
    "\n",
    "    strout.insert(2,\"\\n\")\n",
    "    strout.insert(4,\"\\n\")\n",
    "    strout.insert(6,\"\\n\")\n",
    "\n",
    "    strout = \"X:{}\\nT:{}\\n{}\".format(0,title,\"\".join(strout[1:-1]))\n",
    "           \n",
    "    #save the outputs in a file\n",
    "    if savepath:\n",
    "        file = Path(savepath)\n",
    "        file.parent.mkdir(parents=True, exist_ok=True)\n",
    "        file.write_text(\"\\n\\n\".join(tunes))\n",
    "\n",
    "    return strout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Resume a Run from WadnB and continue training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n",
      " FOLK MLM training script\n",
      "--------------------------------------------------\n",
      "14 Sep 2023 - 14:47:20\n",
      "--------------------------------------------------\n",
      "-------------------------------------------------- \n",
      "Creating Dataset\n",
      "setting up train\n",
      "Loading dataset: ./datasets/df_v4.pickle\n",
      "Loading tokens: ./datasets/TOKENS_V4_arranged.pickle\n",
      "-------------------------------------------------- \n",
      "Resuming Model\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b36d9f993ad24c9b80280bbb145edc97",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.01666974600714942, max=1.0)…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Problem at: /tmp/ipykernel_1979062/3918897933.py 17 <cell line: 17>\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m/workspace/abcMLM/folk_mlm_test.ipynb Cell 16\u001b[0m line \u001b[0;36m<cell line: 17>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=13'>14</a>\u001b[0m run_id \u001b[39m=\u001b[39m artifcat_ref\u001b[39m.\u001b[39msplit(\u001b[39m\"\u001b[39m\u001b[39m/\u001b[39m\u001b[39m\"\u001b[39m)[\u001b[39m2\u001b[39m][\u001b[39m6\u001b[39m:\u001b[39m-\u001b[39m\u001b[39m3\u001b[39m]\n\u001b[1;32m     <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=15'>16</a>\u001b[0m \u001b[39m# download checkpoint locally (if not already cached)\u001b[39;00m\n\u001b[0;32m---> <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=16'>17</a>\u001b[0m run \u001b[39m=\u001b[39m wandb\u001b[39m.\u001b[39;49minit(\n\u001b[1;32m     <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=17'>18</a>\u001b[0m     project \u001b[39m=\u001b[39;49m \u001b[39m\"\u001b[39;49m\u001b[39mabcMLM\u001b[39;49m\u001b[39m\"\u001b[39;49m,\n\u001b[1;32m     <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=18'>19</a>\u001b[0m     resume \u001b[39m=\u001b[39;49m \u001b[39m\"\u001b[39;49m\u001b[39mmust\u001b[39;49m\u001b[39m\"\u001b[39;49m, \n\u001b[1;32m     <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=19'>20</a>\u001b[0m     \u001b[39mid\u001b[39;49m \u001b[39m=\u001b[39;49m run_id\n\u001b[1;32m     <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=20'>21</a>\u001b[0m     )\n\u001b[1;32m     <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=21'>22</a>\u001b[0m \u001b[39m# reference can be retrieved in artifacts panel\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=22'>23</a>\u001b[0m \u001b[39m# \"VERSION\" can be a version (ex: \"v2\") or an alias (\"latest or \"best\")\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=23'>24</a>\u001b[0m \u001b[39m# checkpoint_reference = \"luca-casini/musaic/9dxy5qm5\"\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://attached-container%2B7b22636f6e7461696e65724e616d65223a222f636173696e695f6c696768746e696e67222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6770755f736572766572227d7d/workspace/abcMLM/folk_mlm_test.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=25'>26</a>\u001b[0m artifact \u001b[39m=\u001b[39m run\u001b[39m.\u001b[39muse_artifact(artifcat_ref, \u001b[39mtype\u001b[39m\u001b[39m=\u001b[39m\u001b[39m'\u001b[39m\u001b[39mmodel\u001b[39m\u001b[39m'\u001b[39m)\n",
      "File \u001b[0;32m/usr/local/lib/python3.9/dist-packages/wandb/sdk/wandb_init.py:1177\u001b[0m, in \u001b[0;36minit\u001b[0;34m(job_type, dir, config, project, entity, reinit, tags, group, name, notes, magic, config_exclude_keys, config_include_keys, anonymous, mode, allow_val_change, resume, force, tensorboard, sync_tensorboard, monitor_gym, save_code, id, settings)\u001b[0m\n\u001b[1;32m   1175\u001b[0m     \u001b[39massert\u001b[39;00m logger\n\u001b[1;32m   1176\u001b[0m     logger\u001b[39m.\u001b[39mwarning(\u001b[39m\"\u001b[39m\u001b[39minterrupted\u001b[39m\u001b[39m\"\u001b[39m, exc_info\u001b[39m=\u001b[39me)\n\u001b[0;32m-> 1177\u001b[0m     \u001b[39mraise\u001b[39;00m e\n\u001b[1;32m   1178\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mException\u001b[39;00m \u001b[39mas\u001b[39;00m e:\n\u001b[1;32m   1179\u001b[0m     error_seen \u001b[39m=\u001b[39m e\n",
      "File \u001b[0;32m/usr/local/lib/python3.9/dist-packages/wandb/sdk/wandb_init.py:1154\u001b[0m, in \u001b[0;36minit\u001b[0;34m(job_type, dir, config, project, entity, reinit, tags, group, name, notes, magic, config_exclude_keys, config_include_keys, anonymous, mode, allow_val_change, resume, force, tensorboard, sync_tensorboard, monitor_gym, save_code, id, settings)\u001b[0m\n\u001b[1;32m   1152\u001b[0m except_exit \u001b[39m=\u001b[39m wi\u001b[39m.\u001b[39msettings\u001b[39m.\u001b[39m_except_exit\n\u001b[1;32m   1153\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m-> 1154\u001b[0m     run \u001b[39m=\u001b[39m wi\u001b[39m.\u001b[39;49minit()\n\u001b[1;32m   1155\u001b[0m     except_exit \u001b[39m=\u001b[39m wi\u001b[39m.\u001b[39msettings\u001b[39m.\u001b[39m_except_exit\n\u001b[1;32m   1156\u001b[0m \u001b[39mexcept\u001b[39;00m (\u001b[39mKeyboardInterrupt\u001b[39;00m, \u001b[39mException\u001b[39;00m) \u001b[39mas\u001b[39;00m e:\n",
      "File \u001b[0;32m/usr/local/lib/python3.9/dist-packages/wandb/sdk/wandb_init.py:741\u001b[0m, in \u001b[0;36m_WandbInit.init\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    738\u001b[0m logger\u001b[39m.\u001b[39minfo(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mcommunicating run to backend with \u001b[39m\u001b[39m{\u001b[39;00mtimeout\u001b[39m}\u001b[39;00m\u001b[39m second timeout\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m    740\u001b[0m run_init_handle \u001b[39m=\u001b[39m backend\u001b[39m.\u001b[39minterface\u001b[39m.\u001b[39mdeliver_run(run)\n\u001b[0;32m--> 741\u001b[0m result \u001b[39m=\u001b[39m run_init_handle\u001b[39m.\u001b[39;49mwait(\n\u001b[1;32m    742\u001b[0m     timeout\u001b[39m=\u001b[39;49mtimeout,\n\u001b[1;32m    743\u001b[0m     on_progress\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_on_progress_init,\n\u001b[1;32m    744\u001b[0m     cancel\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m,\n\u001b[1;32m    745\u001b[0m )\n\u001b[1;32m    746\u001b[0m \u001b[39mif\u001b[39;00m result:\n\u001b[1;32m    747\u001b[0m     run_result \u001b[39m=\u001b[39m result\u001b[39m.\u001b[39mrun_result\n",
      "File \u001b[0;32m/usr/local/lib/python3.9/dist-packages/wandb/sdk/lib/mailbox.py:283\u001b[0m, in \u001b[0;36mMailboxHandle.wait\u001b[0;34m(self, timeout, on_probe, on_progress, release, cancel)\u001b[0m\n\u001b[1;32m    280\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_interface\u001b[39m.\u001b[39m_transport_keepalive_failed():\n\u001b[1;32m    281\u001b[0m         \u001b[39mraise\u001b[39;00m MailboxError(\u001b[39m\"\u001b[39m\u001b[39mtransport failed\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m--> 283\u001b[0m found, abandoned \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_slot\u001b[39m.\u001b[39;49m_get_and_clear(timeout\u001b[39m=\u001b[39;49mwait_timeout)\n\u001b[1;32m    284\u001b[0m \u001b[39mif\u001b[39;00m found:\n\u001b[1;32m    285\u001b[0m     \u001b[39m# Always update progress to 100% when done\u001b[39;00m\n\u001b[1;32m    286\u001b[0m     \u001b[39mif\u001b[39;00m on_progress \u001b[39mand\u001b[39;00m progress_handle \u001b[39mand\u001b[39;00m progress_sent:\n",
      "File \u001b[0;32m/usr/local/lib/python3.9/dist-packages/wandb/sdk/lib/mailbox.py:130\u001b[0m, in \u001b[0;36m_MailboxSlot._get_and_clear\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_get_and_clear\u001b[39m(\u001b[39mself\u001b[39m, timeout: \u001b[39mfloat\u001b[39m) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m Tuple[Optional[pb\u001b[39m.\u001b[39mResult], \u001b[39mbool\u001b[39m]:\n\u001b[1;32m    129\u001b[0m     found \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n\u001b[0;32m--> 130\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_wait(timeout\u001b[39m=\u001b[39;49mtimeout):\n\u001b[1;32m    131\u001b[0m         \u001b[39mwith\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_lock:\n\u001b[1;32m    132\u001b[0m             found \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_result\n",
      "File \u001b[0;32m/usr/local/lib/python3.9/dist-packages/wandb/sdk/lib/mailbox.py:126\u001b[0m, in \u001b[0;36m_MailboxSlot._wait\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_wait\u001b[39m(\u001b[39mself\u001b[39m, timeout: \u001b[39mfloat\u001b[39m) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mbool\u001b[39m:\n\u001b[0;32m--> 126\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_event\u001b[39m.\u001b[39;49mwait(timeout\u001b[39m=\u001b[39;49mtimeout)\n",
      "File \u001b[0;32m/usr/lib/python3.9/threading.py:581\u001b[0m, in \u001b[0;36mEvent.wait\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    579\u001b[0m signaled \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_flag\n\u001b[1;32m    580\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m signaled:\n\u001b[0;32m--> 581\u001b[0m     signaled \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_cond\u001b[39m.\u001b[39;49mwait(timeout)\n\u001b[1;32m    582\u001b[0m \u001b[39mreturn\u001b[39;00m signaled\n",
      "File \u001b[0;32m/usr/lib/python3.9/threading.py:316\u001b[0m, in \u001b[0;36mCondition.wait\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    314\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    315\u001b[0m     \u001b[39mif\u001b[39;00m timeout \u001b[39m>\u001b[39m \u001b[39m0\u001b[39m:\n\u001b[0;32m--> 316\u001b[0m         gotit \u001b[39m=\u001b[39m waiter\u001b[39m.\u001b[39;49macquire(\u001b[39mTrue\u001b[39;49;00m, timeout)\n\u001b[1;32m    317\u001b[0m     \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    318\u001b[0m         gotit \u001b[39m=\u001b[39m waiter\u001b[39m.\u001b[39macquire(\u001b[39mFalse\u001b[39;00m)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "print(\"%s\\n FOLK MLM training script\\n%s\" % (\"-\"*50,\"-\"*50 ))\n",
    "print(datetime.now().strftime(\"%d %h %Y - %H:%M:%S\"))\n",
    "print(\"-\"*50)\n",
    "\n",
    "print(\"-\"*50,'\\nCreating Dataset')\n",
    "\n",
    "datamodule = MLMDataModule()\n",
    "datamodule.setup(stage=\"train\")\n",
    "# datamodule.test_df.full_abc\n",
    "\n",
    "\n",
    "print(\"-\"*50,'\\nResuming Model')\n",
    "artifcat_ref = \"musaic/abcMLM/model-92fr0wj9:v0\"\n",
    "run_id = artifcat_ref.split(\"/\")[2][6:-3]\n",
    "\n",
    "# download checkpoint locally (if not already cached)\n",
    "run = wandb.init(\n",
    "    project = \"abcMLM\",\n",
    "    resume = \"must\", \n",
    "    id = run_id\n",
    "    )\n",
    "# reference can be retrieved in artifacts panel\n",
    "# \"VERSION\" can be a version (ex: \"v2\") or an alias (\"latest or \"best\")\n",
    "# checkpoint_reference = \"luca-casini/musaic/9dxy5qm5\"\n",
    "\n",
    "artifact = run.use_artifact(artifcat_ref, type='model')\n",
    "artifact_dir = artifact.download()\n",
    "\n",
    "# load checkpoint\n",
    "model = MaskedLM.load_from_checkpoint(Path(artifact_dir) / \"model.ckpt\")\n",
    "\n",
    "\n",
    "wandb_logger = WandbLogger(\n",
    "    name = run.name, \n",
    "    project = 'abcMLM', \n",
    "    log_model = True, #'all',\n",
    "    resume=\"must\",\n",
    "    )\n",
    "\n",
    "wandb_logger.watch(model, log_graph=False)\n",
    "\n",
    "trainer = Trainer(\n",
    "    devices = \"auto\",\n",
    "    accelerator = \"auto\",\n",
    "    strategy = \"ddp_notebook\",#ddp.DDPStrategy(find_unused_parameters=False),\n",
    "    max_epochs = 100,\n",
    "    gradient_clip_val=1.0,\n",
    "    accumulate_grad_batches=1,\n",
    "    log_every_n_steps=1,\n",
    "    check_val_every_n_epoch=1,\n",
    "    logger=wandb_logger,\n",
    "    callbacks=[\n",
    "        RichProgressBar(), \n",
    "        #RichModelSummary(),\n",
    "        StochasticWeightAveraging(swa_lrs=1e-3),\n",
    "        ],\n",
    "    precision=16,\n",
    "    \n",
    ")\n",
    "\n",
    "trainer.fit(\n",
    "    model,\n",
    "    datamodule\n",
    "    )\n",
    "\n",
    "print(\"Loggin generation examples!\")\n",
    "# log generation examples\n",
    "\n",
    "print(\"generating strctured reels\")\n",
    "wandb_logger.log_text(key=\"reels_struct\", columns=[\"generated tune\"], data=np.asarray([generate_reels(datamodule, n=10,from_structure=True, savepath=\"./mlm_outputs/reels_struct_test.abc\")]).T)\n",
    "print(\"generating random reels\")\n",
    "wandb_logger.log_text(key=\"reels_rand\", columns=[\"generated tune\"], data=np.asarray([generate_reels(datamodule, n=10,from_structure=False, savepath=\"./mlm_outputs/reels_rand_test.abc\")]).T)\n",
    "print(\"generating strctured jigs\")\n",
    "wandb_logger.log_text(key=\"jigs_struct\", columns=[\"generated tune\"], data=np.asarray([generate_jigs(datamodule, n=10,from_structure=True, savepath=\"./mlm_outputs/jigs_struct_test.abc\")]).T)\n",
    "print(\"generating random jigs\")    \n",
    "wandb_logger.log_text(key=\"jigs_rand\", columns=[\"generated tune\"], data=np.asarray([generate_jigs(datamodule, n=10,from_structure=False, savepath=\"./mlm_outputs/jigs_rand_test.abc\")]).T)\n",
    "\n",
    "# is this necessary?\n",
    "wandb_logger.experiment.unwatch(model)\n",
    "\n",
    "run.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Train a new model from scratch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n",
      " FOLK MLM training script\n",
      "--------------------------------------------------\n",
      "21 Mar 2023 - 18:23:31\n",
      "--------------------------------------------------\n",
      "max_len: 256\n",
      "batch_size: 256\n",
      "epochs: 100\n",
      "weight_decay: 0.01\n",
      "lr: 0.0005\n",
      "d_model: 256\n",
      "d_hid: 1024\n",
      "nhead: 8\n",
      "nlayers: 4\n",
      "dropout: 0.0\n",
      "lr_decay: False\n",
      "---\n",
      "-------------------------------------------------- \n",
      "Creating Dataset\n",
      "setting up train\n",
      "Loading dataset: ../Tradformer/v4/datasets/df_v4.pickle\n",
      "Loading tokens: ../Tradformer/v4/datasets/TOKENS_V4_arranged.pickle\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using 16bit Automatic Mixed Precision (AMP)\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------------------------- \n",
      "Creating Model\n",
      "-------------------------------------------------- \n",
      "Creating Trainer\n",
      "setting up TrainerFn.FITTING\n",
      "Loading dataset: ../Tradformer/v4/datasets/df_v4.pickle\n",
      "Loading tokens: ../Tradformer/v4/datasets/TOKENS_V4_arranged.pickle\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using fused AdamW: cuda:0 and True !\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━┓\n",
       "┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">   </span>┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Name                </span>┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Type       </span>┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Params </span>┃\n",
       "┡━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━┩\n",
       "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 0 </span>│ embedding           │ Embedding  │ 33.0 K │\n",
       "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 1 </span>│ learned_pos         │ Embedding  │ 65.5 K │\n",
       "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 2 </span>│ transformer_encoder │ ModuleList │  3.2 M │\n",
       "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 3 </span>│ ln_f                │ LayerNorm  │    512 │\n",
       "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 4 </span>│ decoder             │ Linear     │ 33.0 K │\n",
       "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 5 </span>│ dropout             │ Dropout    │      0 │\n",
       "└───┴─────────────────────┴────────────┴────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━┓\n",
       "┃\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mName               \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mType      \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mParams\u001b[0m\u001b[1;35m \u001b[0m┃\n",
       "┡━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━┩\n",
       "│\u001b[2m \u001b[0m\u001b[2m0\u001b[0m\u001b[2m \u001b[0m│ embedding           │ Embedding  │ 33.0 K │\n",
       "│\u001b[2m \u001b[0m\u001b[2m1\u001b[0m\u001b[2m \u001b[0m│ learned_pos         │ Embedding  │ 65.5 K │\n",
       "│\u001b[2m \u001b[0m\u001b[2m2\u001b[0m\u001b[2m \u001b[0m│ transformer_encoder │ ModuleList │  3.2 M │\n",
       "│\u001b[2m \u001b[0m\u001b[2m3\u001b[0m\u001b[2m \u001b[0m│ ln_f                │ LayerNorm  │    512 │\n",
       "│\u001b[2m \u001b[0m\u001b[2m4\u001b[0m\u001b[2m \u001b[0m│ decoder             │ Linear     │ 33.0 K │\n",
       "│\u001b[2m \u001b[0m\u001b[2m5\u001b[0m\u001b[2m \u001b[0m│ dropout             │ Dropout    │      0 │\n",
       "└───┴─────────────────────┴────────────┴────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Trainable params</span>: 3.3 M                                                                                            \n",
       "<span style=\"font-weight: bold\">Non-trainable params</span>: 0                                                                                            \n",
       "<span style=\"font-weight: bold\">Total params</span>: 3.3 M                                                                                                \n",
       "<span style=\"font-weight: bold\">Total estimated model params size (MB)</span>: 13                                                                         \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mTrainable params\u001b[0m: 3.3 M                                                                                            \n",
       "\u001b[1mNon-trainable params\u001b[0m: 0                                                                                            \n",
       "\u001b[1mTotal params\u001b[0m: 3.3 M                                                                                                \n",
       "\u001b[1mTotal estimated model params size (MB)\u001b[0m: 13                                                                         \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ed9c86d5d2b4df98240b44a2da9fcf9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "TypeError",
     "evalue": "_forward_unimplemented() got an unexpected keyword argument 'src_key_padding_mask'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[12], line 90\u001b[0m\n\u001b[1;32m     69\u001b[0m wandb_logger\u001b[39m.\u001b[39mwatch(model, log_graph\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m)\n\u001b[1;32m     71\u001b[0m trainer \u001b[39m=\u001b[39m Trainer(\n\u001b[1;32m     72\u001b[0m     devices \u001b[39m=\u001b[39m \u001b[39m1\u001b[39m,\n\u001b[1;32m     73\u001b[0m     accelerator \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mauto\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     87\u001b[0m     \n\u001b[1;32m     88\u001b[0m )\n\u001b[0;32m---> 90\u001b[0m trainer\u001b[39m.\u001b[39;49mfit(\n\u001b[1;32m     91\u001b[0m     model,\n\u001b[1;32m     92\u001b[0m     datamodule\n\u001b[1;32m     93\u001b[0m     )\n\u001b[1;32m     95\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpathlib\u001b[39;00m \u001b[39mimport\u001b[39;00m Path\n\u001b[1;32m     97\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mLoggin generation examples!\u001b[39m\u001b[39m\"\u001b[39m)\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py:520\u001b[0m, in \u001b[0;36mTrainer.fit\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[1;32m    518\u001b[0m model \u001b[39m=\u001b[39m _maybe_unwrap_optimized(model)\n\u001b[1;32m    519\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstrategy\u001b[39m.\u001b[39m_lightning_module \u001b[39m=\u001b[39m model\n\u001b[0;32m--> 520\u001b[0m call\u001b[39m.\u001b[39;49m_call_and_handle_interrupt(\n\u001b[1;32m    521\u001b[0m     \u001b[39mself\u001b[39;49m, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path\n\u001b[1;32m    522\u001b[0m )\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py:44\u001b[0m, in \u001b[0;36m_call_and_handle_interrupt\u001b[0;34m(trainer, trainer_fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m     42\u001b[0m         \u001b[39mreturn\u001b[39;00m trainer\u001b[39m.\u001b[39mstrategy\u001b[39m.\u001b[39mlauncher\u001b[39m.\u001b[39mlaunch(trainer_fn, \u001b[39m*\u001b[39margs, trainer\u001b[39m=\u001b[39mtrainer, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[1;32m     43\u001b[0m     \u001b[39melse\u001b[39;00m:\n\u001b[0;32m---> 44\u001b[0m         \u001b[39mreturn\u001b[39;00m trainer_fn(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m     46\u001b[0m \u001b[39mexcept\u001b[39;00m _TunerExitException:\n\u001b[1;32m     47\u001b[0m     _call_teardown_hook(trainer)\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py:559\u001b[0m, in \u001b[0;36mTrainer._fit_impl\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[1;32m    549\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_data_connector\u001b[39m.\u001b[39mattach_data(\n\u001b[1;32m    550\u001b[0m     model, train_dataloaders\u001b[39m=\u001b[39mtrain_dataloaders, val_dataloaders\u001b[39m=\u001b[39mval_dataloaders, datamodule\u001b[39m=\u001b[39mdatamodule\n\u001b[1;32m    551\u001b[0m )\n\u001b[1;32m    553\u001b[0m ckpt_path \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_checkpoint_connector\u001b[39m.\u001b[39m_select_ckpt_path(\n\u001b[1;32m    554\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstate\u001b[39m.\u001b[39mfn,\n\u001b[1;32m    555\u001b[0m     ckpt_path,\n\u001b[1;32m    556\u001b[0m     model_provided\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m,\n\u001b[1;32m    557\u001b[0m     model_connected\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlightning_module \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m,\n\u001b[1;32m    558\u001b[0m )\n\u001b[0;32m--> 559\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_run(model, ckpt_path\u001b[39m=\u001b[39;49mckpt_path)\n\u001b[1;32m    561\u001b[0m \u001b[39massert\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstate\u001b[39m.\u001b[39mstopped\n\u001b[1;32m    562\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtraining \u001b[39m=\u001b[39m \u001b[39mFalse\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py:935\u001b[0m, in \u001b[0;36mTrainer._run\u001b[0;34m(self, model, ckpt_path)\u001b[0m\n\u001b[1;32m    930\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_signal_connector\u001b[39m.\u001b[39mregister_signal_handlers()\n\u001b[1;32m    932\u001b[0m \u001b[39m# ----------------------------\u001b[39;00m\n\u001b[1;32m    933\u001b[0m \u001b[39m# RUN THE TRAINER\u001b[39;00m\n\u001b[1;32m    934\u001b[0m \u001b[39m# ----------------------------\u001b[39;00m\n\u001b[0;32m--> 935\u001b[0m results \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_run_stage()\n\u001b[1;32m    937\u001b[0m \u001b[39m# ----------------------------\u001b[39;00m\n\u001b[1;32m    938\u001b[0m \u001b[39m# POST-Training CLEAN UP\u001b[39;00m\n\u001b[1;32m    939\u001b[0m \u001b[39m# ----------------------------\u001b[39;00m\n\u001b[1;32m    940\u001b[0m log\u001b[39m.\u001b[39mdebug(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__class__\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39m: trainer tearing down\u001b[39m\u001b[39m\"\u001b[39m)\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py:976\u001b[0m, in \u001b[0;36mTrainer._run_stage\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    974\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtraining:\n\u001b[1;32m    975\u001b[0m     \u001b[39mwith\u001b[39;00m isolate_rng():\n\u001b[0;32m--> 976\u001b[0m         \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_run_sanity_check()\n\u001b[1;32m    977\u001b[0m     \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39mautograd\u001b[39m.\u001b[39mset_detect_anomaly(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_detect_anomaly):\n\u001b[1;32m    978\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfit_loop\u001b[39m.\u001b[39mrun()\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py:1005\u001b[0m, in \u001b[0;36mTrainer._run_sanity_check\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1002\u001b[0m call\u001b[39m.\u001b[39m_call_callback_hooks(\u001b[39mself\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mon_sanity_check_start\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m   1004\u001b[0m \u001b[39m# run eval step\u001b[39;00m\n\u001b[0;32m-> 1005\u001b[0m val_loop\u001b[39m.\u001b[39;49mrun()\n\u001b[1;32m   1007\u001b[0m call\u001b[39m.\u001b[39m_call_callback_hooks(\u001b[39mself\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mon_sanity_check_end\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m   1009\u001b[0m \u001b[39m# reset logger connector\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/loops/utilities.py:174\u001b[0m, in \u001b[0;36m_no_grad_context.<locals>._decorator\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    172\u001b[0m     context_manager \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mno_grad\n\u001b[1;32m    173\u001b[0m \u001b[39mwith\u001b[39;00m context_manager():\n\u001b[0;32m--> 174\u001b[0m     \u001b[39mreturn\u001b[39;00m loop_run(\u001b[39mself\u001b[39;49m, \u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/loops/evaluation_loop.py:115\u001b[0m, in \u001b[0;36m_EvaluationLoop.run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    113\u001b[0m     previous_dataloader_idx \u001b[39m=\u001b[39m dataloader_idx\n\u001b[1;32m    114\u001b[0m     \u001b[39m# run step hooks\u001b[39;00m\n\u001b[0;32m--> 115\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_evaluation_step(batch, batch_idx, dataloader_idx)\n\u001b[1;32m    116\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mStopIteration\u001b[39;00m:\n\u001b[1;32m    117\u001b[0m     \u001b[39m# this needs to wrap the `*_step` call too (not just `next`) for `dataloader_iter` support\u001b[39;00m\n\u001b[1;32m    118\u001b[0m     \u001b[39mbreak\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/loops/evaluation_loop.py:375\u001b[0m, in \u001b[0;36m_EvaluationLoop._evaluation_step\u001b[0;34m(self, batch, batch_idx, dataloader_idx)\u001b[0m\n\u001b[1;32m    372\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mbatch_progress\u001b[39m.\u001b[39mincrement_started()\n\u001b[1;32m    374\u001b[0m hook_name \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mtest_step\u001b[39m\u001b[39m\"\u001b[39m \u001b[39mif\u001b[39;00m trainer\u001b[39m.\u001b[39mtesting \u001b[39melse\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39mvalidation_step\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[0;32m--> 375\u001b[0m output \u001b[39m=\u001b[39m call\u001b[39m.\u001b[39;49m_call_strategy_hook(trainer, hook_name, \u001b[39m*\u001b[39;49mstep_kwargs\u001b[39m.\u001b[39;49mvalues())\n\u001b[1;32m    377\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mbatch_progress\u001b[39m.\u001b[39mincrement_processed()\n\u001b[1;32m    379\u001b[0m hook_name \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mon_test_batch_end\u001b[39m\u001b[39m\"\u001b[39m \u001b[39mif\u001b[39;00m trainer\u001b[39m.\u001b[39mtesting \u001b[39melse\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39mon_validation_batch_end\u001b[39m\u001b[39m\"\u001b[39m\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py:288\u001b[0m, in \u001b[0;36m_call_strategy_hook\u001b[0;34m(trainer, hook_name, *args, **kwargs)\u001b[0m\n\u001b[1;32m    285\u001b[0m     \u001b[39mreturn\u001b[39;00m\n\u001b[1;32m    287\u001b[0m \u001b[39mwith\u001b[39;00m trainer\u001b[39m.\u001b[39mprofiler\u001b[39m.\u001b[39mprofile(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m[Strategy]\u001b[39m\u001b[39m{\u001b[39;00mtrainer\u001b[39m.\u001b[39mstrategy\u001b[39m.\u001b[39m\u001b[39m__class__\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39m.\u001b[39m\u001b[39m{\u001b[39;00mhook_name\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m):\n\u001b[0;32m--> 288\u001b[0m     output \u001b[39m=\u001b[39m fn(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m    290\u001b[0m \u001b[39m# restore current_fx when nested context\u001b[39;00m\n\u001b[1;32m    291\u001b[0m pl_module\u001b[39m.\u001b[39m_current_fx_name \u001b[39m=\u001b[39m prev_fx_name\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/lightning/pytorch/strategies/strategy.py:378\u001b[0m, in \u001b[0;36mStrategy.validation_step\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    376\u001b[0m \u001b[39mwith\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprecision_plugin\u001b[39m.\u001b[39mval_step_context():\n\u001b[1;32m    377\u001b[0m     \u001b[39massert\u001b[39;00m \u001b[39misinstance\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel, ValidationStep)\n\u001b[0;32m--> 378\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mmodel\u001b[39m.\u001b[39;49mvalidation_step(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "Cell \u001b[0;32mIn[11], line 287\u001b[0m, in \u001b[0;36mMaskedLM.validation_step\u001b[0;34m(self, val_batch, batch_idx)\u001b[0m\n\u001b[1;32m    285\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mvalidation_step\u001b[39m(\u001b[39mself\u001b[39m, val_batch, batch_idx):\n\u001b[1;32m    286\u001b[0m     x, y \u001b[39m=\u001b[39m val_batch\n\u001b[0;32m--> 287\u001b[0m     o \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mforward(x)\n\u001b[1;32m    288\u001b[0m     \u001b[39m# y = torch.where(y==self.PAD_TOKEN,self.IGNORE_TOKEN,y)\u001b[39;00m\n\u001b[1;32m    289\u001b[0m     loss \u001b[39m=\u001b[39m F\u001b[39m.\u001b[39mcross_entropy(\n\u001b[1;32m    290\u001b[0m         o\u001b[39m.\u001b[39mpermute(\u001b[39m0\u001b[39m,\u001b[39m2\u001b[39m,\u001b[39m1\u001b[39m), \n\u001b[1;32m    291\u001b[0m         y, \n\u001b[1;32m    292\u001b[0m         ignore_index\u001b[39m=\u001b[39m\u001b[39m-\u001b[39m\u001b[39m100\u001b[39m, \n\u001b[1;32m    293\u001b[0m         \u001b[39m# weight=TOKENS_WEIGHTS.to(self.device)\u001b[39;00m\n\u001b[1;32m    294\u001b[0m         )\n",
      "Cell \u001b[0;32mIn[11], line 237\u001b[0m, in \u001b[0;36mMaskedLM.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    234\u001b[0m position_embeddings \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlearned_pos(pos)\n\u001b[1;32m    235\u001b[0m x \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdropout(token_embeddings \u001b[39m+\u001b[39m position_embeddings)\n\u001b[0;32m--> 237\u001b[0m x \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mtransformer_encoder(x, src_key_padding_mask \u001b[39m=\u001b[39;49m attn_mask)\n\u001b[1;32m    238\u001b[0m x \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdecoder(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mln_f(x))\n\u001b[1;32m    240\u001b[0m \u001b[39mreturn\u001b[39;00m x\n",
      "File \u001b[0;32m~/anaconda3/envs/musaic/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1502\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "\u001b[0;31mTypeError\u001b[0m: _forward_unimplemented() got an unexpected keyword argument 'src_key_padding_mask'"
     ]
    }
   ],
   "source": [
    "print(\"%s\\n FOLK MLM training script\\n%s\" % (\"-\"*50,\"-\"*50 ))\n",
    "print(datetime.now().strftime(\"%d %h %Y - %H:%M:%S\"))\n",
    "print(\"-\"*50)\n",
    "\n",
    "class Config():\n",
    "    def __init__(self) -> None:\n",
    "        self.max_len = 256\n",
    "        self.batch_size = 256\n",
    "        self.epochs = 100\n",
    "        self.weight_decay = 0.01\n",
    "        self.lr = 5e-4\n",
    "        self.d_model = 256\n",
    "        self.d_hid = 4 * self.d_model\n",
    "        self.nhead = 8 \n",
    "        self.nlayers = 4\n",
    "        self.dropout = 0.0\n",
    "        self.lr_decay = False\n",
    "    \n",
    "    def __repr__(self) -> str:\n",
    "        for e in self.__dict__:\n",
    "            print(str(e)+\": \"+str(self.__dict__[e]))\n",
    "        return \"---\"\n",
    "    \n",
    "    def to_dict(self):\n",
    "        return {e : self.__dict__[e] for e in self.__dict__}\n",
    "\n",
    "        \n",
    "    \n",
    "args = Config()\n",
    "print(args)\n",
    "\n",
    "print(\"-\"*50,'\\nCreating Dataset')\n",
    "\n",
    "datamodule = MLMDataModule()\n",
    "datamodule.setup(stage=\"train\")\n",
    "# datamodule.test_df.full_abc\n",
    "\n",
    "\n",
    "# model\n",
    "print(\"-\"*50,\"\\nCreating Model\")\n",
    "\n",
    "model = MaskedLM(\n",
    "    PAD_TOKEN=datamodule.stoi[\"<pad>\"], IGNORE_TOKEN=-100,\n",
    "    ntoken=datamodule.vocab_size, \n",
    "    d_model=args.d_model, d_hid=args.d_hid, nhead=args.nhead, nlayers=args.nlayers,\n",
    "    dropout=args.dropout, lr=args.lr, lr_sched=args.lr_decay,\n",
    "    weight_decay=args.weight_decay, custom_block=True\n",
    "    )\n",
    "\n",
    "# model = torch.compile(model)\n",
    "\n",
    "# training\n",
    "print(\"-\"*50,\"\\nCreating Trainer\")\n",
    "\n",
    "from lightning.pytorch import Trainer\n",
    "from lightning.pytorch.loggers import WandbLogger\n",
    "from lightning.pytorch.callbacks import StochasticWeightAveraging, RichProgressBar, RichModelSummary\n",
    "from lightning.pytorch.strategies import ddp\n",
    "\n",
    "wandb_logger = WandbLogger(\n",
    "    name = \"onnx_exp_jupyter\", \n",
    "    project = 'onnx', \n",
    "    log_model = True, #'all',\n",
    "    # resume=\"must\",\n",
    "    )\n",
    "\n",
    "wandb_logger.experiment.config.update(args.to_dict())\n",
    "\n",
    "wandb_logger.watch(model, log_graph=False)\n",
    "\n",
    "trainer = Trainer(\n",
    "    devices = 1,\n",
    "    accelerator = \"auto\",\n",
    "    # strategy = \"ddp\",#ddp.DDPStrategy(find_unused_parameters=False),\n",
    "    max_epochs = args.epochs,\n",
    "    gradient_clip_val=1.0,\n",
    "    accumulate_grad_batches=1,\n",
    "    log_every_n_steps=1,\n",
    "    check_val_every_n_epoch=1,\n",
    "    logger=wandb_logger,\n",
    "    callbacks=[\n",
    "        RichProgressBar(), \n",
    "        #RichModelSummary(),\n",
    "        StochasticWeightAveraging(swa_lrs=1e-3),\n",
    "        ],\n",
    "    precision=16,\n",
    "    \n",
    ")\n",
    "\n",
    "trainer.fit(\n",
    "    model,\n",
    "    datamodule\n",
    "    )\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "print(\"Loggin generation examples!\")\n",
    "# log generation examples\n",
    "\n",
    "print(\"generating strctured reels\")\n",
    "wandb_logger.log_text(key=\"reels_struct\", columns=[\"generated tune\"], data=np.asarray([generate_reels(datamodule, n=10,from_structure=True, savepath=\"./mlm_outputs/reels_struct_test.abc\")]).T)\n",
    "print(\"generating random reels\")\n",
    "wandb_logger.log_text(key=\"reels_rand\", columns=[\"generated tune\"], data=np.asarray([generate_reels(datamodule, n=10,from_structure=False, savepath=\"./mlm_outputs/reels_rand_test.abc\")]).T)\n",
    "print(\"generating strctured jigs\")\n",
    "wandb_logger.log_text(key=\"jigs_struct\", columns=[\"generated tune\"], data=np.asarray([generate_jigs(datamodule, n=10,from_structure=True, savepath=\"./mlm_outputs/jigs_struct_test.abc\")]).T)\n",
    "print(\"generating random jigs\")    \n",
    "wandb_logger.log_text(key=\"jigs_rand\", columns=[\"generated tune\"], data=np.asarray([generate_jigs(datamodule, n=10,from_structure=False, savepath=\"./mlm_outputs/jigs_rand_test.abc\")]).T)\n",
    "\n",
    "# is this necessary?\n",
    "wandb_logger.experiment.unwatch(model)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Load a model from WandB and generate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "setting up train\n",
      "Loading dataset: ./datasets/df_v4.pickle\n",
      "Loading tokens: ./datasets/TOKENS_V4_arranged.pickle\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mluca-casini\u001b[0m (\u001b[33mmusaic\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.15.10 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.5"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/workspace/abcMLM/wandb/run-20230914_145458-92fr0wj9</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Resuming run <strong><a href='https://wandb.ai/musaic/abcMLM/runs/92fr0wj9' target=\"_blank\">folkMLM_13Sep2023-144853</a></strong> to <a href='https://wandb.ai/musaic/abcMLM' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/musaic/abcMLM' target=\"_blank\">https://wandb.ai/musaic/abcMLM</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/musaic/abcMLM/runs/92fr0wj9' target=\"_blank\">https://wandb.ai/musaic/abcMLM/runs/92fr0wj9</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Downloading large artifact model-92fr0wj9:v0, 194.63MB. 1 files... \n",
      "\u001b[34m\u001b[1mwandb\u001b[0m:   1 of 1 files downloaded.  \n",
      "Done. 0:0:0.5\n"
     ]
    }
   ],
   "source": [
    "# we need the datamodule for itos and stoi\n",
    "datamodule = MLMDataModule()\n",
    "datamodule.setup(stage=\"train\")\n",
    "\n",
    "artifcat_ref = \"musaic/abcMLM/model-92fr0wj9:v0\"\n",
    "run_id = artifcat_ref.split(\"/\")[2][6:-3]\n",
    "\n",
    "# download checkpoint locally (if not already cached)\n",
    "run = wandb.init(\n",
    "    project = \"abcMLM\",\n",
    "    resume = \"must\", \n",
    "    id = run_id\n",
    "    )\n",
    "\n",
    "artifact = run.use_artifact(artifcat_ref, type='model')\n",
    "artifact_dir = artifact.download()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = MaskedLM(\n",
    "#     PAD_TOKEN=datamodule.stoi[\"<pad>\"], IGNORE_TOKEN=-100,\n",
    "#     ntoken=datamodule.vocab_size, \n",
    "#     d_model=256, d_hid=1024, nhead=8, nlayers=4,\n",
    "#     dropout=0.0, lr=1e-3, lr_sched=False,\n",
    "#     weight_decay=True, custom_block=True\n",
    "#     )\n",
    "\n",
    "# model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Lightning automatically upgraded your loaded checkpoint from v1.9.3 to v2.0.0. To apply the upgrade to your files permanently, run `python -m pytorch_lightning.utilities.upgrade_checkpoint --file artifacts/model-jglxq9af:v0/model.ckpt`\n"
     ]
    }
   ],
   "source": [
    "# load checkpoint\n",
    "model = MaskedLM.load_from_checkpoint(Path(artifact_dir) / \"model.ckpt\")\n",
    "# model.eval()\n",
    "# model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# onnx save\n",
    "dummy_input = datamodule.test_set.__getitem__(0)[0].reshape(1,-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(\n",
    "#     model, dummy_input, opset_version=14\n",
    "# )\n",
    "\n",
    "# print(set(unconvertible_ops))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.onnx.export(model,               # model being run\n",
    "                  dummy_input,                         # model input (or a tuple for multiple inputs)\n",
    "                  \"../onnx_experiment/test.onnx\",   # where to save the model (can be a file or file-like object)\n",
    "                  export_params=True,        # store the trained parameter weights inside the model file\n",
    "                  opset_version=15,          # the ONNX version to export the model to\n",
    "                #   do_constant_folding=True,  # whether to execute constant folding for optimization\n",
    "                input_names = ['mlmInput'],   # the model's input names\n",
    "                output_names = ['mlmOutput'], # the model's output names\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.to_onnx(\n",
    "#     \"../onnx_experiment/test.onnx\", \n",
    "#     input_sample=dummy_input, \n",
    "#     # opset_version=14, \n",
    "#     export_params=True, \n",
    "#     do_constant_folding=True, \n",
    "#     export_modules_as_functions=False, \n",
    "#     operator_export_type=torch.onnx.OperatorExportTypes.ONNX)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bobs Tests\n",
    "\n",
    "- left-to-right generation \n",
    "- mask last two tokes of connaughtman rambles\n",
    "- give as input 5 measures with 6 masked tokens in each"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### load a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "setting up train\n",
      "Loading dataset: ./datasets/df_v4.pickle\n",
      "Loading tokens: ./datasets/TOKENS_V4_arranged.pickle\n",
      "-------------------------------------------------- \n",
      "Resuming Model\n"
     ]
    }
   ],
   "source": [
    "datamodule = MLMDataModule()\n",
    "datamodule.setup(stage=\"train\")\n",
    "# datamodule.test_df.full_abc\n",
    "\n",
    "print(\"-\"*50,'\\nResuming Model')\n",
    "artifact_ref = \"musaic/musaic/model-92fr0wj9:v0\"\n",
    "\n",
    "try:\n",
    "    # load checkpoint\n",
    "    artifact_dir = './artifacts/'+artifact_ref.split('/')[-1]\n",
    "    model = MaskedLM.load_from_checkpoint(Path(artifact_dir) / \"model.ckpt\").to('cpu')\n",
    "except:\n",
    "    api = wandb.Api()\n",
    "    artifact = api.artifact(artifact_ref)\n",
    "    artifact_dir = artifact.download()\n",
    "    model = MaskedLM.load_from_checkpoint(Path(artifact_dir) / \"model.ckpt\").to('cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### parallel decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mCanceled future for execute_request message before replies were done"
     ]
    }
   ],
   "source": [
    "M = [np.random.choice(['M:4/4','M:6/8'])]\n",
    "K = [np.random.choice(['K:Cmaj','K:Cmin','K:Cmix','K:Cdor'])] \n",
    "input_seq = [\"<s>\", \"L:1/8\"] + M + K + [\"<mask>\"]*(256-4)\n",
    "\n",
    "seq_len = len(input_seq)\n",
    "# transform sequence to indices\n",
    "input_seq = np.array([datamodule.stoi[t] for t in input_seq])\n",
    "input_seq = np.pad(input_seq,(0,256-seq_len),\"constant\", constant_values=datamodule.stoi[\"<pad>\"]).reshape(1,-1) # pad\n",
    "masked = list(np.where(input_seq[0]==datamodule.stoi[\"<mask>\"])[0]) # get \n",
    "        \n",
    "# print(\"prediction index:\",i)\n",
    "logits = model(torch.IntTensor(input_seq)) \n",
    "logits = logits / 1.0\n",
    "            \n",
    "# if top_k:\n",
    "#     logits = top_k_top_p_filtering(logits, top_k=top_k, filter_value=-torch.inf)\n",
    "\n",
    "# elif top_p:\n",
    "#     logits = top_k_top_p_filtering(logits, top_p=top_p, filter_value=-torch.inf)\n",
    "#     # if verbose: print(sum(logits[0,i] > -torch.inf))\n",
    "\n",
    "print(logits.shape)\n",
    "\n",
    "argmax = logits.argmax(dim=-1)\n",
    "print(argmax.shape)\n",
    "\n",
    "strout = [datamodule.itos[t.item()] for t in argmax[0]]\n",
    "strout.insert(2,\"\\n\")\n",
    "strout.insert(4,\"\\n\")\n",
    "strout.insert(6,\"\\n\")\n",
    "# strout = \"X:{}\\nT:{}\\nN:tokens={}\\n{}\".format('999'+str(idx),\"random autoregressive\",strout_len,\"\".join(strout[:]))\n",
    "\n",
    "print(\"\".join(strout))\n",
    "\n",
    "# probs = torch.nn.functional.softmax(logits,dim=-1)\n",
    "# print(sampled.shape)\n",
    "# dist = torch.distributions.Categorical(probs)\n",
    "# sampled = dist.sample()\n",
    "# print(sampled.shape)\n",
    "\n",
    "# input_seq[0,i] = sampled\n",
    "# if verbose:\n",
    "#     print( \"\".join([datamodule.itos[t] for t in input_seq[0]]).replace(\"<mask>\",\"_\").replace(\"<pad>\",\"~\"))\n",
    "# # once we get the end token we exit\n",
    "# if sampled == datamodule.stoi[\"</s>\"] and early_stop:\n",
    "#     if verbose: print(\"eos i:\", i)\n",
    "#     break\n",
    "\n",
    "#         strout = [datamodule.itos[t] for t in input_seq[0] if t != datamodule.stoi[\"<pad>\"]]\n",
    "#         strout_len = len(strout)\n",
    "#         # if verbose: print(\"length:\", strout_len)\n",
    "#         # if verbose: print(\" \".join(strout).replace(\"<mask>\",\"_\"))\n",
    "        \n",
    "#         strout.insert(2,\"\\n\")\n",
    "#         strout.insert(4,\"\\n\")\n",
    "#         strout.insert(6,\"\\n\")\n",
    "        \n",
    "#         strout = \"X:{}\\nT:{}\\nN:tokens={}\\n{}\".format('999'+str(idx),\"random autoregressive\",strout_len,\"\".join(strout[:]))\n",
    "#         tunes.append(strout)\n",
    "            \n",
    "#     #save the outputs in a file\n",
    "#     if savepath:\n",
    "#         file = Path(savepath)\n",
    "#         file.parent.mkdir(parents=True, exist_ok=True)\n",
    "#         file.write_text(\"\\n\\n\".join(tunes))\n",
    "\n",
    "#     return tunes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### one token generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X:9990\n",
      "T:random autoregressive\n",
      "N:tokens=118\n",
      "L:1/8\n",
      "M:6/8\n",
      "K:Cdor\n",
      "|:GCCCDF|GcccBG|BB,B,CB,|FGB,CDF|GCCCDF|GcccBG|FBBDEF|GCCCDF:|B|cc'c'gbg|dfdcBG|GcaB3|fdfbgf|gc'c'gbg|dfdcBF|GBBDEF|GCCCDF:|</s>\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"\\n\\n\".join(generate_autoregressive(None, datamodule, n=1, order='random', temperature=.95, top_p=0.95, verbose=False, early_stop=False)).replace(\"<s>\",\"\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X:9990\n",
      "T:random autoregressive\n",
      "N:tokens=256\n",
      "L:1/8\n",
      "M:6/8\n",
      "K:Cmin\n",
      "ECCDEG|cBcGEC|E2GGFE|DCDFED|ECCDEG|cBcGEC|E2GGFE|1ECB,C2D:||2ECB,C2F|:GcccBc|dcAABA|GBBFBB|dcdedc|GcccBc|dcAABA|GBBGFE|1ECB,CEF:||2ECB,C2d|:ecAGEC|ededcB|GBBBB|FBBfdB|eccGEC|ecAGEC|ededcB|1ECB,C2:||2ECB,C2d|ecAGEC|ededcB|GBBdcB|dcBedB|edcdcB|c/2d/2edcBG|1ECB,C2d:||2|2ECB,C2G|</s>\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"\\n\\n\".join(generate_autoregressive(None, datamodule, n=1, order='l2r', temperature=.95, top_p=0.95, verbose=False, early_stop=False)).replace(\"<s>\",\"\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X:9990\n",
      "T:random autoregressive\n",
      "N:tokens=256\n",
      "L:1/8\n",
      "M:6/8\n",
      "K:Cmin\n",
      "DBBAGA|B2dfdB|cBcFAc|GcBAcA|DBBAFA|B2dfdB|g2efed|1c3CEG:||2c3ccA|GccFA2|FAcAcA|GBBFDF|BcedcB|GccFAc|e3ecG|e3efd|gfedcB|1c3cc'2:||2c3c2d|:e2ggef|c'2ggec|b2ffdf|b2dfdG|e2efdG|e2edef|gfedcB|1c3c2=B:||2c3c2B|:edcdBG|e2gfdf|b2bfdB|e2efdB|efgc'3|efg=ab3|gfed2B|1c3c2F:||2c3c3|</s>\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"\\n\\n\".join(generate_autoregressive(None, datamodule, n=1, order='r2l', temperature=.95, top_p=0.95, verbose=False, early_stop=False)).replace(\"<s>\",\"\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### jigs and reels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X:16349\n",
      "T:based on test 16349\n",
      "L:1/8\n",
      "M:6/8\n",
      "K:Cmaj\n",
      "c_BGFDC|DEFG2C|DEFEDC|FAFGAB|c_BGFDC|DEFG2C|E2FE2G|FDB,C3:|ccAG^FG|efed2B|cecdcB|cdec2G|cecdBG|GABG2F|E2FG2G|FDB,C3:| \n",
      "\n",
      "X:10248\n",
      "T:based on test 10248\n",
      "L:1/8\n",
      "M:4/4\n",
      "K:Cmaj\n",
      "G2GEDFED|C2CCE2F/2G/2A|G2GEC2CC|E/2F/2G3GG2cG|A3GE2C2|F2FFG2E2|G2F2D2C2|G,2CCD3E|G2E2E2C2|D2D2D2D2|F2E2D2C2|G,2CCC2D/2E/2F|G2E2E2C2|G2FFG2E2|F2F2D2C2|G,2CCD2C2| \n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"\\n\\n\".join(generate_jigs(datamodule, n=1, order=\"random\", from_structure=True, top_p=.9, verbose=False)),'\\n'\n",
    ")\n",
    "\n",
    "print(\n",
    "    \"\\n\\n\".join(generate_reels(datamodule, n=1, order=\"random\", from_structure=True, top_p=.9, verbose=False)),'\\n'\n",
    ")\n",
    "\n",
    "\n",
    "# print(\n",
    "#     \"\\n\\n\".join(generate_jigs(datamodule, n=1, order=\"l2r\", from_structure=True, top_k=5, verbose=False)),'\\n'\n",
    "# )\n",
    "\n",
    "# print(\n",
    "#     \"\\n\\n\".join(generate_jigs(datamodule, n=1, order=\"r2l\", from_structure=True, top_k=5, verbose=False)),'\\n'\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(\n",
    "#     \"\\n\\n\".join(generate_jigs(datamodule, n=1, order=\"random\", from_structure=False, top_k=5, verbose=False)),'\\n'\n",
    "# )\n",
    "\n",
    "print(\n",
    "    \"\\n\\n\".join(generate_reels(datamodule, n=10, order=\"random\", sample_length=False, from_structure=False, top_p=.9, verbose=False)),'\\n'\n",
    ")\n",
    "\n",
    "\n",
    "# print(\n",
    "#     \"\\n\\n\".join(generate_reels(datamodule, n=1, order=\"l2r\", from_structure=True, top_k=5, verbose=False)),'\\n'\n",
    "# )\n",
    "\n",
    "# print(\n",
    "#     \"\\n\\n\".join(generate_reels(datamodule, n=1, order=\"r2l\", from_structure=True, top_k=5, verbose=False)),'\\n'\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X:9990\n",
      "T:random autoregressive\n",
      "N:tokens=256\n",
      "<s>L:1/8\n",
      "M:6/8\n",
      "K:Cmaj\n",
      "G|:GcccBc|AcAAG^F|GEGCEG|FDDDAA|GcccBc|AcAAG^F|GEGCEG|1FDCC2G:||2FDCC2E|:GccAcc|AcAA2^F|GEGCEG|DDDFED|CccAcc|AcAA2^F|GEGCEG|1FDCC2G,:||2FDCC3|CEGCEG|CEGCF^F|GEGCEG|DEFDB,G,|CEGCEG|CEGCF^F|GEGCEG|FEDC2G,|:CEGCEG|CFACF^F|GEGCEG|DDDFED|CEGCEG|CFACFA|CFAAF^F|GEGCEG|FDCCC||</s></s></s></s>\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"\\n\\n\".join(\n",
    "        generate_autoregressive('<s> L:1/8 M:6/8', datamodule, n=1, order=\"random\", verbose=False, temperature=1.0, top_k=None, top_p=0.9, savepath=None, early_stop=False,)\n",
    "        )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Mask the endings of The Connaughtman's Rambles  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "```\n",
    "X:1  \n",
    "T:The Connaughtman's Rambles  \n",
    "R:jig  \n",
    "M:6/8  \n",
    "L:1/8  \n",
    "K:Cmaj  \n",
    "|:EGG cGG|AGG cGF|EGG ced|cAA AGF|  \n",
    "EGG cGG|AGG cde|fed ced|1 cAA AGF:|2 cAA A3||  \n",
    "|:eaa ege|edc dcd|eaa ege|edc d3|  \n",
    "eaa ege|edc cde|fed ced|1 cAA A3:|2 cAA AGF||  \n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "connaughtman_ending = \"<s> M:6/8 L:1/8 K:Cmaj |: E G G c G G | A G G c G F | E G G c e d | c A A A G F | E G G c G G | A G G c d e | f e d c e d |1 <mask> <mask> <mask> <mask> <mask> <mask> :| |2 <mask> <mask> <mask> <mask> 3 | \"+\\\n",
    "\"|: e a a e g e | e d c d c d | e a a e g e | e d c d 3 | e a a e g e | e d c c d e | f e d c e d |1 <mask> <mask> <mask> <mask> 3 :| |2 <mask> <mask> <mask> <mask> <mask> <mask> | </s>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGced|cAAAGF|EGGcGG|AGGcde|fedced|1cedc2c:||2cedA3|\n",
      "|:eaaege|edcdcd|eaaege|edcd3|eaaege|edccde|fedced|1Bcdc3:||2cedcAc| \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGced|cAAAGF|EGGcGG|AGGcde|fedced|1ecccBA:||2eccc3|\n",
      "|:eaaege|edcdcd|eaaege|edcd3|eaaege|edccde|fedced|1eccc3:||2eccGcA| \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGced|cAAAGF|EGGcGG|AGGcde|fedced|1ecccAF:||2eccc3|\n",
      "|:eaaege|edcdcd|eaaege|edcd3|eaaege|edccde|fedced|1eccc3:||2eccccA| \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGced|cAAAGF|EGGcGG|AGGcde|fedced|1edccGF:||2edcB3|\n",
      "|:eaaege|edcdcd|eaaege|edcd3|eaaege|edccde|fedced|1edcc3:||2edccc3| \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGced|cAAAGF|EGGcGG|AGGcde|fedced|1cAGAGF:||2cAGc3|\n",
      "|:eaaege|edcdcd|eaaege|edcd3|eaaege|edccde|fedced|1cAGG3:||2cAGAGF| \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGced|cAAAGF|EGGcGG|AGGcde|fedced|1eccc3F:||2eccc3|\n",
      "|:eaaege|edcdcd|eaaege|edcd3|eaaege|edccde|fedced|1eccc3:||2eccc2c| \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGced|cAAAGF|EGGcGG|AGGcde|fedced|1cfedge:||2cffg3|\n",
      "|:eaaege|edcdcd|eaaege|edcd3|eaaege|edccde|fedced|1cfed3:||2cfedcd| \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGced|cAAAGF|EGGcGG|AGGcde|fedced|1ceccBA:||2cecc3|\n",
      "|:eaaege|edcdcd|eaaege|edcd3|eaaege|edccde|fedced|1cecc3:||2ceccc3| \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGced|cAAAGF|EGGcGG|AGGcde|fedced|1cAAAGF:||2cAAA3|\n",
      "|:eaaege|edcdcd|eaaege|edcd3|eaaege|edccde|fedced|1cAAA3:||2cAAAGF| \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGced|cAAAGF|EGGcGG|AGGcde|fedced|1cAGABc:||2cdcc3|\n",
      "|:eaaege|edcdcd|eaaege|edcd3|eaaege|edccde|fedced|1cccc3:||2cggecc| \n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    print(\n",
    "        fill_masked(connaughtman_ending, datamodule, title=\"The Connaughtman's Fillmask\", verbose=False).replace(\"||:\",\"|\\n|:\"),\"\\n\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGEGG|AGGFED|EGGcGG|AGGcde|fedced|1dBBBAG:||2dBBB3|\n",
      "|:eaaege|edcdcd|egggeg|aaaa3|eaaege|edccde|fedced|1dBBB3:||2dBBBAG|</s>%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGAGG|AGGcAG|EGGcGG|AGGcde|fedced|1dBBBAG:||2dBBB3|\n",
      "|:eaaege|edcdcd|ecgaeg|gecd3|eaaege|edccde|fedced|1dBBB3:||2dBBBAG|</s>%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGAGG|GFFFGF|EGGcGG|AGGcde|fedced|1dBBBAG:||2dBBB3|\n",
      "|:eaaege|edcdcd|edcdcc|dBBd3|eaaege|edccde|fedced|1dBBB3:||2dBBBAG|</s>%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EFFFFF|DDCDEF|EGGcGG|AGGcde|fedced|1dBBBAG:||2dBBB3|\n",
      "|:eaaege|edcdcd|Bfeede|cAGc3|eaaege|edccde|fedced|1dBBB3:||2dBBBAG|</s>%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGcGG|AGAAGF|EGGcGG|AGGcde|fedced|1dBBBAG:||2dBBB3|\n",
      "|:eaaege|edcdcd|cccdef|eddd3|eaaege|edccde|fedced|1dBBB3:||2dBBBAG|</s>%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGcGG|FEFDED|EGGcGG|AGGcde|fedced|1dBBBAG:||2dBBB3|\n",
      "|:eaaege|edcdcd|eaaegg|edcd3|eaaege|edccde|fedced|1dBBB3:||2dBBBAG|</s>%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EGGcGG|efedcF|EGGcGG|AGGcde|fedced|1dBBBAG:||2dBBB3|\n",
      "|:eaaege|edcdcd|eAAcBc|de^fg3|eaaege|edccde|fedced|1dBBB3:||2dBBBAG|</s>%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% \n",
      "\n",
      "X:0\n",
      "T:The Connaughtman's Fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "|:EGGcGG|AGGcGF|EEEcGE|DDDDCD|EGGcGG|AGGcde|fedced|1dBBBAG:||2dBBB3|\n",
      "|:eaaege|edcdcd|eedcde|dddd3|eaaege|edccde|fedced|1dBBB3:||2dBBBAG|</s>%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% \n",
      "\n"
     ]
    }
   ],
   "source": [
    "connaughtman_34 = \"<s> M:6/8 L:1/8 K:Cmaj |: E G G c G G | A G G c G F | <mask> <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> <mask> | E G G c G G | A G G c d e | f e d c e d |1 d B B B A G :| |2 d B B B 3 | \" + \\\n",
    "\"|: e a a e g e | e d c d c d | <mask> <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> 3 | e a a e g e | e d c c d e | f e d c e d |1 d B B B 3 :| |2 d B B B A G | </s>\"\n",
    "\n",
    "# connaughtman_34\n",
    "\n",
    "for i in range(8):\n",
    "    print(\n",
    "        fill_masked(connaughtman_34, datamodule, title=\"The Connaughtman's Fillmask\", verbose=False).replace(\"||:\",\"|\\n|:\"),\"\\n\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### generate 5 bars with 6 tokens each"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmin\n",
      "cB|G|G|B3/2||ee|e|d|cB|z|:|G|D||A|AcF|E4|c3| \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "Bddcdf|g2fede|fdGece|fegecB|GABcfe|g3|</s><mask><mask> \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "AGGGGG|G2||G2|G2|</s><mask><mask>|<mask><mask><mask><mask><mask><mask>|<mask><mask><mask><mask><mask><mask>|<mask><mask><mask><mask><mask><mask> \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmin\n",
      "CD|E4E|CD|B,<D|C2|DCD|E|E2ED|C3|CDE|C4|C|F, \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cdor\n",
      "CE|EFG|C8|CDA,|E|E|CC|E3|C|G|E|AFAF|G6||EF \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmix\n",
      "ECCDFD|FDD2C2|ECCGFD|BBc=ABc|GFECDE|DEDCD2 \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmin\n",
      "C|C|CC||B,/2B,|A,|C|</s><mask><mask><mask>|<mask><mask><mask><mask><mask><mask>|<mask><mask><mask><mask><mask><mask>|<mask><mask><mask><mask><mask><mask> \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmin\n",
      "GE|B,2C|C3|B,2C|C2|</s><mask><mask>|<mask><mask><mask><mask><mask><mask>|<mask><mask><mask><mask><mask><mask>|<mask><mask><mask><mask><mask><mask> \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmix\n",
      "c'3/2c'3/2g/2|agfecc|eccczd|c'bc'd'c'a|gafece|dfdbc'b \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "CEGAFD|AF2FDA|AC2B,DC|DA2G2D|D2DG,2B,|C2ECDE \n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt = [\"<s>\"] + [\"M:6/8\",\"L:1/8\",\"<mask>\"] + [\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\"|\"]*6 \n",
    "prompt = \" \".join(prompt) \n",
    "# print(prompt)\n",
    "\n",
    "for i in range(10):\n",
    "    print(\n",
    "        fill_masked(prompt, datamodule, verbose=False, temperature=0.99, top_p=0.99),\"\\n\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## generate seven bars sections with 5 tokens each"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> M:6/8 L:1/8 <mask> <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> :| <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> | <mask> <mask> <mask> <mask> <mask> :| </s>\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "c>dec|AFG>F|EGC>D|EDD>G|c>dec|AFG>F|EGC>z:||:cGEG|cGE>G|AGF>G|AGE>G|cGE>G|AFG>F|EGC>z:| \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "G\n",
      "c2dcB|A>GAB|c2B>G|G>GGF|E>FGE|D>EFD|C>CC2:|D>EFG|A>GAB|c>BAG|G>GGF|E>FGE|D>EFG|C>CC2:| \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "A\n",
      "|:CCDE|DGD>E|EGc>d|cGc>B|AAG>F|EGA>B|cdBc2:||:cBcd|ecd>c|BGG>d|cBc>d|cBc>d|ecd>c|BGG>A:| \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmin\n",
      "|:EGAG|c>BAG|c>BAG|c>BAG|EGA>B|c>BAG|D>EFD:||:EGcG|EGc>d|edc>d|ecBAG|EGc>d|edc>d|edc>A:| \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "|\n",
      "G,CCDC|FDB,>G,|G,CCDC|DCD>G,|G,CCDC|FDB,>G,|C>CC2:||:G2GE|G>EDC|A,>FED|C>A,G,2|C>A,G,2|A,>FED|C>CC2:| \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "e\n",
      "c>dec|G2G>A|B>cde|dcBAG|c>dec|G2G>A|B>AGE:|D2D>G|G>dBG|G>dec|d>cBG|G>dBd|G2G>A|G>AGE:| \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "K:Cmaj\n",
      "GEGEG|AFD>F|EGc>c|AGG>F|EGE>G|FDB,>D|DCC>D:||:EGcG|EGc>G|EGc>G|AGG>F|EGc>G|AGG>F|ECC>G:| \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "C\n",
      "GEFDE|CEGc2|BAG>F|DEF>A|GEF>E|EFG>F|D>CC2:|cBAGF|EGc>c|BAG>F|DEF>A|GEF>E|E>FGE|D>CC2:| \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "A\n",
      "G2FGB|GFF>C|F2FGB|G2G>F|G2F>B|AGF>F|G>CCB,:|C2C>D|E>FGA|G>GGA|F>G>A|G>AGF|D>CCD|F>EED:| \n",
      "\n",
      "X:0\n",
      "T:fillmask\n",
      "M:6/8\n",
      "L:1/8\n",
      "C\n",
      "c>BAG|F>EDC|F>GAG|FDC>B|c>BAG|F>EDC|D>EDC:||:G,CC2|GCC2C|G>AAG|FDC>B,|G,CC>C|E>DEC|D>EDC:| \n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt = [\"<s>\"] + [\"M:6/8\",\"L:1/8\",\"<mask>\"] + [\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\"|\"]*6 + [\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\":|\"] + [\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\"|\"]*6 + [\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\"<mask>\",\":|\", \"</s>\"]\n",
    "prompt = \" \".join(prompt) \n",
    "print(prompt)\n",
    "# print(prompt)\n",
    "\n",
    "for i in range(10):\n",
    "    print(\n",
    "        fill_masked(prompt, datamodule, verbose=False, temperature=0.9, top_p=0.9),\"\\n\"\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "91c06897c2ada423a836e2c65925763ecd1b24085fa893398bfd2024239c4086"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}