Super_GPT2_Piano.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.8.5"
    },
    "colab": {
      "name": "Super GPT2 Piano.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/asigalov61/Amazing-GPT2-Piano/blob/master/Super_GPT2_Piano.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QA0W-VK1JVQl",
        "colab_type": "text"
      },
      "source": [
        "# Super GPT2 Piano\n",
        "\n",
        "## All credit for this beautiful colab implementation of char-based Music GPT2 goes out to Andrej Karpathy and Edtky of GitHub on whose repos/code it is based: \n",
        "\n",
        "https://github.com/karpathy/minGPT\n",
        "\n",
        " https://github.com/edtky/mini-musical-neural-net"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "21sOhWtzySmB",
        "colab_type": "text"
      },
      "source": [
        "## A simple implementation of Music GPT2 trained on the classical music. \n",
        "\n",
        "The IOs here are simple text files, which we chop up to individual characters and then train GPT on. So you could say this is a char-transformer instead of a char-rnn. Doesn't quite roll off the tongue well. In this example we will feed it some Music, which we'll get it to predict on a character-level."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HsUtsJGNz6f2",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Clone minGPT repo and install all dependencies (run only once per session)\n",
        "!git clone https://github.com/asigalov61/minGPT\n",
        " \n",
        "!pip install pyknon\n",
        "!pip install pretty_midi\n",
        "!pip install pypianoroll\n",
        "!pip install mir_eval\n",
        "!apt install fluidsynth #Pip does not work for some reason. Only apt works\n",
        "!pip install midi2audio\n",
        "!cp /usr/share/sounds/sf2/FluidR3_GM.sf2 /content/font.sf2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pf8B3p6QySmE",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Import all modules and setup logging\n",
        "%cd /content/minGPT\n",
        "# make deterministic\n",
        "from mingpt.utils import set_seed\n",
        "set_seed(42)\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "#from torch.nn import functional as F\n",
        "\n",
        "# set up logging\n",
        "import logging\n",
        "logging.basicConfig(\n",
        "        format=\"%(asctime)s - %(levelname)s - %(name)s -   %(message)s\",\n",
        "        datefmt=\"%m/%d/%Y %H:%M:%S\",\n",
        "        level=logging.INFO,\n",
        ")\n",
        "\n",
        "from torch import optim\n",
        "import torch.nn.functional as F\n",
        "\n",
        "import keras\n",
        "from keras.utils import to_categorical\n",
        "\n",
        "import time\n",
        "\n",
        "import pretty_midi\n",
        "from midi2audio import FluidSynth\n",
        "from google.colab import output\n",
        "from IPython.display import display, Javascript, HTML, Audio\n",
        "\n",
        "dtype = torch.float\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "# Assume that we are on a CUDA machine, then this should print a CUDA device:\n",
        "print('Available Device:', device)\n",
        "\n",
        "!mkdir /content/midis\n",
        "\n",
        "sample_freq_variable = 12 #@param {type:\"number\"}\n",
        "note_range_variable = 62 #@param {type:\"number\"}\n",
        "note_offset_variable = 33 #@param {type:\"number\"}\n",
        "number_of_instruments = 2 #@param {type:\"number\"}\n",
        "chamber_option = True #@param {type:\"boolean\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7VsJDrWbIcUo",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title (OPTIONAL) Convert your own MIDIs to Notewise TXT DataSet (before running this cell, upload your MIDI DataSet to /content/midis folder)\n",
        "import tqdm.auto\n",
        "import argparse\n",
        "import random\n",
        "import os\n",
        "import numpy as np\n",
        "from math import floor\n",
        "from pyknon.genmidi import Midi\n",
        "from pyknon.music import NoteSeq, Note\n",
        "import music21\n",
        "from music21 import instrument, volume\n",
        "from music21 import midi as midiModule\n",
        "from pathlib import Path\n",
        "import glob, sys\n",
        "from music21 import converter, instrument\n",
        "%cd /content\n",
        "notes=[]\n",
        "InstrumentID=0\n",
        "folder = '/content/midis/*mid'\n",
        "for file in tqdm.auto.tqdm(glob.glob(folder)):\n",
        "    filename = file[-53:]\n",
        "    print(filename)\n",
        "\n",
        "    # fname = \"../midi-files/mozart/sonat-3.mid\"\n",
        "    fname = filename\n",
        "\n",
        "    mf=music21.midi.MidiFile()\n",
        "    mf.open(fname)\n",
        "    mf.read()\n",
        "    mf.close()\n",
        "    midi_stream=music21.midi.translate.midiFileToStream(mf)\n",
        "    midi_stream\n",
        "\n",
        "\n",
        "\n",
        "    sample_freq=sample_freq_variable\n",
        "    note_range=note_range_variable\n",
        "    note_offset=note_offset_variable\n",
        "    chamber=chamber_option\n",
        "    numInstruments=number_of_instruments\n",
        "\n",
        "    s = midi_stream\n",
        "    #print(s.duration.quarterLength)\n",
        "\n",
        "    s[0].elements\n",
        "\n",
        "\n",
        "    maxTimeStep = floor(s.duration.quarterLength * sample_freq)+1\n",
        "    score_arr = np.zeros((maxTimeStep, numInstruments, note_range))\n",
        "\n",
        "    #print(maxTimeStep, \"\\n\", score_arr.shape)\n",
        "\n",
        "    # define two types of filters because notes and chords have different structures for storing their data\n",
        "    # chord have an extra layer because it consist of multiple notes\n",
        "\n",
        "    noteFilter=music21.stream.filters.ClassFilter('Note')\n",
        "    chordFilter=music21.stream.filters.ClassFilter('Chord')\n",
        "      \n",
        "\n",
        "    # pitch.midi-note_offset: pitch is the numerical representation of a note. \n",
        "    #                         note_offset is the the pitch relative to a zero mark. eg. B-=25, C=27, A=24\n",
        "\n",
        "    # n.offset: the timestamps of each note, relative to the start of the score\n",
        "    #           by multiplying with the sample_freq, you make all the timestamps integers\n",
        "\n",
        "    # n.duration.quarterLength: the duration of that note as a float eg. quarter note = 0.25, half note = 0.5\n",
        "    #                           multiply by sample_freq to represent duration in terms of timesteps\n",
        "\n",
        "    notes = []\n",
        "    instrumentID = 0\n",
        "    parts = instrument.partitionByInstrument(s)\n",
        "    for i in range(len(parts.parts)): \n",
        "      instru = parts.parts[i].getInstrument()\n",
        "      \n",
        "\n",
        "    for n in s.recurse().addFilter(noteFilter):\n",
        "        if chamber:\n",
        "          # assign_instrument where 0 means piano-like and 1 means violin-like, and -1 means neither\n",
        "            notes.append((n.pitch.midi-note_offset, floor(n.offset*sample_freq), \n",
        "              floor(n.duration.quarterLength*sample_freq), 1))\n",
        "\n",
        "          #if instru.instrumentName == 'Piano':\n",
        "           #   instrumentID=0\n",
        "          #if instru.instrumentName == 'Violin':\n",
        "         #     instrumentID=1\n",
        "\n",
        "        notes.append((n.pitch.midi-note_offset, floor(n.offset*sample_freq), \n",
        "              floor(n.duration.quarterLength*sample_freq), 0))\n",
        "        \n",
        "    #print(len(notes))\n",
        "    notes[-5:]\n",
        "\n",
        "    # do the same using a chord filter\n",
        "\n",
        "    for c in s.recurse().addFilter(chordFilter):\n",
        "        # unlike the noteFilter, this line of code is necessary as there are multiple notes in each chord\n",
        "        # pitchesInChord is a list of notes at each chord eg. (<music21.pitch.Pitch D5>, <music21.pitch.Pitch F5>)\n",
        "        pitchesInChord=c.pitches\n",
        "        \n",
        "        # do same as noteFilter and append all notes to the notes list\n",
        "        for p in pitchesInChord:\n",
        "            notes.append((p.midi-note_offset, floor(c.offset*sample_freq), \n",
        "                          floor(c.duration.quarterLength*sample_freq), 1))\n",
        "\n",
        "        # do same as noteFilter and append all notes to the notes list\n",
        "        for p in pitchesInChord:\n",
        "            notes.append((p.midi-note_offset, floor(c.offset*sample_freq), \n",
        "                          floor(c.duration.quarterLength*sample_freq), 0))\n",
        "    #print(len(notes))\n",
        "    notes[-5:]\n",
        "\n",
        "    # the variable/list \"notes\" is a collection of all the notes in the song, not ordered in any significant way\n",
        "\n",
        "    for n in notes:\n",
        "        \n",
        "        # pitch is the first variable in n, previously obtained by n.midi-note_offset\n",
        "        pitch=n[0]\n",
        "        \n",
        "        # do some calibration for notes that fall our of note range\n",
        "        # i.e. less than 0 or more than note_range\n",
        "        while pitch<0:\n",
        "            pitch+=12\n",
        "        while pitch>=note_range:\n",
        "            pitch-=12\n",
        "            \n",
        "        # 3rd element refers to instrument type => if instrument is violin, use different pitch calibration\n",
        "        if n[3]==1:      #Violin lowest note is v22\n",
        "            while pitch<22:\n",
        "                pitch+=12\n",
        "\n",
        "        # start building the 3D-tensor of shape: (796, 1, 38)\n",
        "        # score_arr[0] = timestep\n",
        "        # score_arr[1] = type of instrument\n",
        "        # score_arr[2] = pitch/note out of the range of note eg. 38\n",
        "        \n",
        "        # n[0] = pitch\n",
        "        # n[1] = timestep\n",
        "        # n[2] = duration\n",
        "        # n[3] = instrument\n",
        "        #print(n[3])\n",
        "        score_arr[n[1], n[3], pitch]=1                  # Strike note\n",
        "        score_arr[n[1]+1:n[1]+n[2], n[3], pitch]=2      # Continue holding note\n",
        "\n",
        "    #print(score_arr.shape)\n",
        "    # print first 5 timesteps\n",
        "    score_arr[:5,0,]\n",
        "\n",
        "\n",
        "    for timestep in score_arr:\n",
        "        #print(list(reversed(range(len(timestep)))))\n",
        "        break\n",
        "\n",
        "    instr={}\n",
        "    instr[0]=\"p\"\n",
        "    instr[1]=\"v\"\n",
        "\n",
        "    score_string_arr=[]\n",
        "\n",
        "    # loop through all timesteps\n",
        "    for timestep in score_arr:\n",
        "        \n",
        "        # selecting the instruments: i=0 means piano and i=1 means violin\n",
        "        for i in list(reversed(range(len(timestep)))):   # List violin note first, then piano note\n",
        "            \n",
        "            # \n",
        "            score_string_arr.append(instr[i]+''.join([str(int(note)) for note in timestep[i]]))\n",
        "\n",
        "    #print(type(score_string_arr), len(score_string_arr))\n",
        "    score_string_arr[:5]\n",
        "\n",
        "    modulated=[]\n",
        "    # get the note range from the array\n",
        "    note_range=len(score_string_arr[0])-1\n",
        "\n",
        "    for i in range(0,12):\n",
        "        for chord in score_string_arr:\n",
        "            \n",
        "            # minus the instrument letter eg. 'p'\n",
        "            # add 6 zeros on each side of the string\n",
        "            padded='000000'+chord[1:]+'000000'\n",
        "            \n",
        "            # add back the instrument letter eg. 'p'\n",
        "            # append window of len=note_range back into \n",
        "            # eg. if we have \"00012345000\"\n",
        "            # iteratively, we want to get \"p00012\", \"p00123\", \"p01234\", \"p12345\", \"p23450\", \"p34500\", \"p45000\",\n",
        "            modulated.append(chord[0]+padded[i:i+note_range])\n",
        "\n",
        "    # 796 * 12\n",
        "    #print(len(modulated))\n",
        "    modulated[:5]\n",
        "\n",
        "    # input of this function is a modulated string\n",
        "    long_string = modulated\n",
        "\n",
        "    translated_list=[]\n",
        "\n",
        "    # for every timestep of the string\n",
        "    for j in range(len(long_string)):\n",
        "        \n",
        "        # chord at timestep j eg. 'p00000000000000000000000000000000000100'\n",
        "        chord=long_string[j]\n",
        "        next_chord=\"\"\n",
        "        \n",
        "        # range is from next_timestep to max_timestep\n",
        "        for k in range(j+1, len(long_string)):\n",
        "            \n",
        "            # checking if instrument of next chord is same as current chord\n",
        "            if long_string[k][0]==chord[0]:\n",
        "                \n",
        "                # if same, set next chord as next element in modulation\n",
        "                # otherwise, keep going until you find a chord with the same instrument\n",
        "                # when you do, set it as the next chord\n",
        "                next_chord=long_string[k]\n",
        "                break\n",
        "        \n",
        "        # set prefix as the instrument\n",
        "        # set chord and next_chord to be without the instrument prefix\n",
        "        # next_chord is necessary to check when notes end\n",
        "        prefix=chord[0]\n",
        "        chord=chord[1:]\n",
        "        next_chord=next_chord[1:]\n",
        "        \n",
        "        # checking for non-zero notes at one particular timestep\n",
        "        # i is an integer indicating the index of each note the chord\n",
        "        for i in range(len(chord)):\n",
        "            \n",
        "            if chord[i]==\"0\":\n",
        "                continue\n",
        "            \n",
        "            # set note as 2 elements: instrument and index of note\n",
        "            # examples: p22, p16, p4\n",
        "            #p = music21.pitch.Pitch()\n",
        "            #nt = music21.note.Note(p)\n",
        "            #n.volume.velocity = 20\n",
        "            #nt.volume.client == nt\n",
        "            #V = nt.volume.velocity\n",
        "            #print(V)\n",
        "            #note=prefix+str(i)+' V' + str(V)\n",
        "            note=prefix+str(i)                \n",
        "            \n",
        "            # if note in chord is 1, then append the note eg. p22 to the list\n",
        "            if chord[i]==\"1\":\n",
        "                translated_list.append(note)\n",
        "            \n",
        "            # If chord[i]==\"2\" do nothing - we're continuing to hold the note\n",
        "            \n",
        "            # unless next_chord[i] is back to \"0\" and it's time to end the note.\n",
        "            if next_chord==\"\" or next_chord[i]==\"0\":      \n",
        "                translated_list.append(\"end\"+note)\n",
        "\n",
        "        # wait indicates end of every timestep\n",
        "        if prefix==\"p\":\n",
        "            translated_list.append(\"wait\")\n",
        "\n",
        "    #print(len(translated_list))\n",
        "    translated_list[:10]\n",
        "\n",
        "    # this section transforms the list of notes into a string of notes\n",
        "\n",
        "    # initialize i as zero and empty string\n",
        "    i=0\n",
        "    translated_string=\"\"\n",
        "\n",
        "\n",
        "    while i<len(translated_list):\n",
        "        \n",
        "        # stack all the repeated waits together using an integer to indicate the no. of waits\n",
        "        # eg. \"wait wait\" => \"wait2\"\n",
        "        wait_count=1\n",
        "        if translated_list[i]=='wait':\n",
        "            while wait_count<=sample_freq*2 and i+wait_count<len(translated_list) and translated_list[i+wait_count]=='wait':\n",
        "                wait_count+=1\n",
        "            translated_list[i]='wait'+str(wait_count)\n",
        "            \n",
        "        # add next note\n",
        "        translated_string+=translated_list[i]+\" \"\n",
        "        i+=wait_count\n",
        "\n",
        "    translated_string[:100]\n",
        "    len(translated_string)\n",
        "\n",
        "    #print(\"chordwise encoding type and length:\", type(modulated), len(modulated))\n",
        "    #print(\"notewise encoding type and length:\", type(translated_string), len(translated_string))\n",
        "\n",
        "    # default settings: sample_freq=12, note_range=62\n",
        "\n",
        "    chordwise_folder = \"../\"\n",
        "    notewise_folder = \"../\"\n",
        "\n",
        "    # export chordwise encoding\n",
        "#    f=open(chordwise_folder+fname+\"_chordwise\"+\".txt\",\"w+\")\n",
        "#    f.write(\" \".join(modulated))\n",
        "#    f.close()\n",
        "\n",
        "    # export notewise encoding\n",
        "    f=open(notewise_folder+fname+\"_notewise\"+\".txt\",\"w+\")\n",
        "    f.write(translated_string)\n",
        "    f.close()\n",
        "\n",
        "folder = '/content/midis/*notewise.txt'\n",
        "\n",
        "\n",
        "filenames = glob.glob('/content')\n",
        "with open('notewise_custom_dataset.txt', 'w') as outfile:\n",
        "    for fname in glob.glob(folder)[-53:]:\n",
        "        with open(fname) as infile:\n",
        "            for line in infile:\n",
        "                outfile.write(line)\n",
        "\n",
        "#folder = '/content/midis/*chordwise.txt'\n",
        "\n",
        "#filenames = glob.glob('/content')\n",
        "#with open('chordwise_custom_dataset.txt', 'w') as outfile:\n",
        "#    for fname in glob.glob(folder)[-53:]:\n",
        "#        with open(fname) as infile:\n",
        "#            for line in infile:\n",
        "#                outfile.write(line)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z1MQDOVMySmJ",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Setup functions and procedures\n",
        "model_attention_span_in_tokens = 64 #@param {type:\"slider\", min:0, max:512, step:16}\n",
        "import math\n",
        "from torch.utils.data import Dataset\n",
        "\n",
        "class CharDataset(Dataset):\n",
        "\n",
        "    def __init__(self, data, block_size):\n",
        "        chars = sorted(list(set(data)))\n",
        "        data_size, vocab_size = len(data), len(chars)\n",
        "        print('data has %d characters, %d unique.' % (data_size, vocab_size))\n",
        "        \n",
        "        self.stoi = { ch:i for i,ch in enumerate(chars) }\n",
        "        self.itos = { i:ch for i,ch in enumerate(chars) }\n",
        "        self.block_size = block_size\n",
        "        self.vocab_size = vocab_size\n",
        "        self.data = data\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.data) - self.block_size\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        # grab a chunk of (block_size + 1) characters from the data\n",
        "        chunk = self.data[idx:idx + self.block_size + 1]\n",
        "        # encode every character to an integer\n",
        "        dix = [self.stoi[s] for s in chunk]\n",
        "        \"\"\"\n",
        "        arrange data and targets so that the first i elements of x\n",
        "        will be asked to predict the i-th element of y. Notice that\n",
        "        the eventual language model will actually make block_size\n",
        "        individual predictions at the same time based on this data,\n",
        "        so we are being clever and amortizing the cost of the forward\n",
        "        pass of the network. So for example if block_size is 4, then\n",
        "        we could e.g. sample a chunk of text \"hello\", the integers in\n",
        "        x will correspond to \"hell\" and in y will be \"ello\". This will\n",
        "        then actually \"multitask\" 4 separate examples at the same time\n",
        "        in the language model:\n",
        "        - given just \"h\", please predict \"e\" as next\n",
        "        - given \"he\" please predict \"l\" next\n",
        "        - given \"hel\" predict \"l\" next\n",
        "        - given \"hell\" predict \"o\" next\n",
        "        \n",
        "        In addition, because the DataLoader will create batches of examples,\n",
        "        every forward/backward pass during traning will simultaneously train\n",
        "        a LOT of predictions, amortizing a lot of computation. In particular,\n",
        "        for a batched input of integers X (B, T) where B is batch size and\n",
        "        T is block_size and Y (B, T), the network will during training be\n",
        "        simultaneously training to make B*T predictions, all at once! Of course,\n",
        "        at test time we can paralellize across batch B, but unlike during training\n",
        "        we cannot parallelize across the time dimension T - we have to run\n",
        "        a forward pass of the network to recover the next single character of the \n",
        "        sequence along each batch dimension, and repeatedly always feed in a next\n",
        "        character to get the next one.\n",
        "        \n",
        "        So yes there is a big asymmetry between train/test time of autoregressive\n",
        "        models. During training we can go B*T at a time with every forward pass,\n",
        "        but during test time we can only go B at a time, T times, with T forward \n",
        "        passes.\n",
        "        \"\"\"\n",
        "        x = torch.tensor(dix[:-1], dtype=torch.long)\n",
        "        y = torch.tensor(dix[1:], dtype=torch.long)\n",
        "        return x, y\n",
        "\n",
        "        \n",
        "block_size = model_attention_span_in_tokens # spatial extent of the model for its context"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hLnzSr8clDBj",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title (OPTIONAL) Download ready-to-use Piano and Chamber Notewise DataSets\n",
        "%cd /content/\n",
        "!wget 'https://github.com/asigalov61/SuperPiano/raw/master/Super%20Chamber%20Piano%20Violin%20Notewise%20DataSet.zip'\n",
        "!unzip '/content/Super Chamber Piano Violin Notewise DataSet.zip'\n",
        "!rm '/content/Super Chamber Piano Violin Notewise DataSet.zip'\n",
        "\n",
        "!wget 'https://github.com/asigalov61/SuperPiano/raw/master/Super%20Chamber%20Piano%20Only%20Notewise%20DataSet.zip'\n",
        "!unzip '/content/Super Chamber Piano Only Notewise DataSet.zip'\n",
        "!rm '/content/Super Chamber Piano Only Notewise DataSet.zip'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m4QIgbe3ySmN",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Specify input text file with training data (do not worry, any text format is fine)\n",
        "full_path_to_training_text_file = \"/content/notewise_piano.txt\" #@param {type:\"string\"}\n",
        "text = open(full_path_to_training_text_file, 'r').read() # don't worry we won't run out of file handles\n",
        "train_dataset = CharDataset(text, block_size) # one line of poem is roughly 50 characters"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bpNxwzNkySmQ",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Create GPT2 model\n",
        "model_embed_size = 128 #@param {type:\"slider\", min:0, max:1024, step:64}\n",
        "number_of_heads = 4 #@param {type:\"slider\", min:1, max:16, step:1}\n",
        "number_of_layers = 4 #@param {type:\"slider\", min:1, max:16, step:1}\n",
        "from mingpt.model import GPT, GPTConfig\n",
        "mconf = GPTConfig(train_dataset.vocab_size, train_dataset.block_size,\n",
        "                  n_layer=number_of_layers, n_head=number_of_heads, n_embd=model_embed_size)\n",
        "model = GPT(mconf)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BWaOS0sRySmS",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Setup all training parameters\n",
        "number_of_training_epochs = 100 #@param {type:\"slider\", min:1, max:100, step:1}\n",
        "training_batch_size = 3072 #@param {type:\"slider\", min:0, max:3072, step:16}\n",
        "from mingpt.trainer import Trainer, TrainerConfig\n",
        "\n",
        "# initialize a trainer instance and kick off training\n",
        "tconf = TrainerConfig(max_epochs=number_of_training_epochs, batch_size=training_batch_size, learning_rate=6e-4,\n",
        "                      lr_decay=True, warmup_tokens=512*20, final_tokens=2*len(train_dataset)*block_size,\n",
        "                      num_workers=4)\n",
        "trainer = Trainer(model, train_dataset, None, tconf)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cRffqT9WFBHB",
        "colab_type": "code",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "624d2d9a-e9c3-41cc-c361-1ce3e708d4ea"
      },
      "source": [
        "#@title (OPTION 1) Train the model\n",
        "trainer.train()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "epoch 1 iter 898: train loss 0.49843. lr 5.999570e-04:   1%|          | 899/83377 [07:14<11:01:30,  2.08it/s]"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pMkyEMghC-KR",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Save/Re-Save the model from memory\n",
        "torch.save(model, 'trained-model.pth')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CmD7VRZhDcnJ",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title (OPTION 2) Load existing model/checkpoint\n",
        "model = torch.load('trained-model.pth')\n",
        "model.eval()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6ZEqKJ6NySmV",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Generate and Download Music with the Model (Notewise TXT2MIDI)\n",
        "number_of_tokens_to_generate = 4096 #@param {type:\"slider\", min:0, max:8192, step:64}\n",
        "creativity_temperature = 0.7 #@param {type:\"slider\", min:0.05, max:4, step:0.05}\n",
        "top_k_prob = 3 #@param {type:\"slider\", min:0, max:50, step:1}\n",
        "input_promt = \"v\" #@param {type:\"string\"}\n",
        "# alright, let's sample some character-level Shakespeare\n",
        "from mingpt.utils import sample\n",
        "import tqdm\n",
        "\n",
        "context = input_promt\n",
        "x = torch.tensor([train_dataset.stoi[s] for s in context], dtype=torch.long)[None,...].to(trainer.device)\n",
        "y = sample(model, x, number_of_tokens_to_generate, temperature=creativity_temperature, sample=True, top_k=top_k_prob)[0]\n",
        "completion = ''.join([train_dataset.itos[int(i)] for i in y])\n",
        "print('Done! Saving output.txt!')\n",
        "with open(\"/content/output.txt\", \"w\") as text_file:\n",
        "    print(completion, file=text_file)\n",
        "\n",
        "\n",
        "\n",
        "#@title Generate and Download resulting MIDI file\n",
        "time_coefficient = 4 #@param {type:\"integer\"}\n",
        "\n",
        "import os\n",
        "import dill as pickle\n",
        "from pathlib import Path\n",
        "import random\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from math import floor\n",
        "from pyknon.genmidi import Midi\n",
        "from pyknon.music import NoteSeq, Note\n",
        "import music21\n",
        "import random\n",
        "import os, argparse\n",
        "\n",
        "# default settings: sample_freq=12, note_range=62\n",
        "\n",
        "def decoder(filename):\n",
        "    \n",
        "    filedir = '/content/'\n",
        "\n",
        "    notetxt = filedir + filename\n",
        "\n",
        "    with open(notetxt, 'r') as file:\n",
        "        notestring=file.read()\n",
        "\n",
        "    score_note = notestring.split(\" \")\n",
        "\n",
        "    # define some parameters (from encoding script)\n",
        "    sample_freq=sample_freq_variable\n",
        "    note_range=note_range_variable\n",
        "    note_offset=note_offset_variable\n",
        "    chamber=chamber_option\n",
        "    numInstruments=number_of_instruments\n",
        "\n",
        "    # define variables and lists needed for chord decoding\n",
        "    speed=time_coefficient/sample_freq\n",
        "    piano_notes=[]\n",
        "    violin_notes=[]\n",
        "    time_offset=0\n",
        "\n",
        "    # start decoding here\n",
        "    score = score_note\n",
        "\n",
        "    i=0\n",
        "\n",
        "    # for outlier cases, not seen in sonat-1.txt\n",
        "    # not exactly sure what scores would have \"p_octave_\" or \"eoc\" (end of chord?)\n",
        "    # it seems to insert new notes to the score whenever these conditions are met\n",
        "    while i<len(score):\n",
        "        if score[i][:9]==\"p_octave_\":\n",
        "            add_wait=\"\"\n",
        "            if score[i][-3:]==\"eoc\":\n",
        "                add_wait=\"eoc\"\n",
        "                score[i]=score[i][:-3]\n",
        "            this_note=score[i][9:]\n",
        "            score[i]=\"p\"+this_note\n",
        "            score.insert(i+1, \"p\"+str(int(this_note)+12)+add_wait)\n",
        "            i+=1\n",
        "        i+=1\n",
        "\n",
        "\n",
        "    # loop through every event in the score\n",
        "    for i in range(len(score)):\n",
        "\n",
        "        # if the event is a blank, space, \"eos\" or unknown, skip and go to next event\n",
        "        if score[i] in [\"\", \" \", \"<eos>\", \"<unk>\"]:\n",
        "            continue\n",
        "\n",
        "        # if the event starts with 'end' indicating an end of note\n",
        "        elif score[i][:3]==\"end\":\n",
        "\n",
        "            # if the event additionally ends with eoc, increare the time offset by 1\n",
        "            if score[i][-3:]==\"eoc\":\n",
        "                time_offset+=1\n",
        "            continue\n",
        "\n",
        "        # if the event is wait, increase the timestamp by the number after the \"wait\"\n",
        "        elif score[i][:4]==\"wait\":\n",
        "            time_offset+=int(score[i][4:])\n",
        "            continue\n",
        "\n",
        "        # in this block, we are looking for notes   \n",
        "        else:\n",
        "            # Look ahead to see if an end<noteid> was generated\n",
        "            # soon after.  \n",
        "            duration=1\n",
        "            has_end=False\n",
        "            note_string_len = len(score[i])\n",
        "            for j in range(1,200):\n",
        "                if i+j==len(score):\n",
        "                    break\n",
        "                if score[i+j][:4]==\"wait\":\n",
        "                    duration+=int(score[i+j][4:])\n",
        "                if score[i+j][:3+note_string_len]==\"end\"+score[i] or score[i+j][:note_string_len]==score[i]:\n",
        "                    has_end=True\n",
        "                    break\n",
        "                if score[i+j][-3:]==\"eoc\":\n",
        "                    duration+=1\n",
        "\n",
        "            if not has_end:\n",
        "                duration=12\n",
        "\n",
        "            add_wait = 0\n",
        "            if score[i][-3:]==\"eoc\":\n",
        "                score[i]=score[i][:-3]\n",
        "                add_wait = 1\n",
        "\n",
        "            try: \n",
        "                new_note=music21.note.Note(int(score[i][1:])+note_offset)    \n",
        "                new_note.duration = music21.duration.Duration(duration*speed)\n",
        "                new_note.offset=time_offset*speed\n",
        "                if score[i][0]==\"v\":\n",
        "                    violin_notes.append(new_note)\n",
        "                else:\n",
        "                    piano_notes.append(new_note)                \n",
        "            except:\n",
        "                print(\"Unknown note: \" + score[i])\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "            time_offset+=add_wait\n",
        "\n",
        "    # list of all notes for each instrument should be ready at this stage\n",
        "\n",
        "    # creating music21 instrument objects      \n",
        "    \n",
        "    piano=music21.instrument.fromString(\"Piano\")\n",
        "    violin=music21.instrument.fromString(\"Violin\")\n",
        "\n",
        "    # insert instrument object to start (0 index) of notes list\n",
        "    \n",
        "    piano_notes.insert(0, piano)\n",
        "    violin_notes.insert(0, violin)\n",
        "    # create music21 stream object for individual instruments\n",
        "    \n",
        "    piano_stream=music21.stream.Stream(piano_notes)\n",
        "    violin_stream=music21.stream.Stream(violin_notes)\n",
        "    # merge both stream objects into a single stream of 2 instruments\n",
        "    note_stream = music21.stream.Stream([piano_stream, violin_stream])\n",
        "\n",
        "    \n",
        "    note_stream.write('midi', fp=\"/content/\"+filename[:-4]+\".mid\")\n",
        "    print(\"Done! Decoded midi file saved to 'content/'\")\n",
        "\n",
        "    \n",
        "decoder('output.txt')\n",
        "from google.colab import files\n",
        "files.download('/content/output.mid')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xl1d9LYuJ9uZ",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Plot, Graph, and Listen to the Output :)\n",
        "graphs_length_inches = 18 #@param {type:\"slider\", min:0, max:20, step:1}\n",
        "notes_graph_height = 6 #@param {type:\"slider\", min:0, max:20, step:1}\n",
        "highest_displayed_pitch = 92 #@param {type:\"slider\", min:1, max:128, step:1}\n",
        "lowest_displayed_pitch = 24 #@param {type:\"slider\", min:1, max:128, step:1}\n",
        "\n",
        "import librosa\n",
        "import numpy as np\n",
        "import pretty_midi\n",
        "import pypianoroll\n",
        "from pypianoroll import Multitrack, Track\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "#matplotlib.use('SVG')\n",
        "# For plotting\n",
        "import mir_eval.display\n",
        "import librosa.display\n",
        "%matplotlib inline\n",
        "\n",
        "\n",
        "midi_data = pretty_midi.PrettyMIDI('/content/output.mid')\n",
        "\n",
        "def plot_piano_roll(pm, start_pitch, end_pitch, fs=100):\n",
        "    # Use librosa's specshow function for displaying the piano roll\n",
        "    librosa.display.specshow(pm.get_piano_roll(fs)[start_pitch:end_pitch],\n",
        "                             hop_length=1, sr=fs, x_axis='time', y_axis='cqt_note',\n",
        "                             fmin=pretty_midi.note_number_to_hz(start_pitch))\n",
        "\n",
        "\n",
        "\n",
        "roll = np.zeros([int(graphs_length_inches), 128])\n",
        "# Plot the output\n",
        "\n",
        "track = Multitrack('/content/output.mid', name='track')\n",
        "plt.figure(figsize=[graphs_length_inches, notes_graph_height])\n",
        "fig, ax = track.plot()\n",
        "fig.set_size_inches(graphs_length_inches, notes_graph_height)\n",
        "plt.figure(figsize=[graphs_length_inches, notes_graph_height])\n",
        "ax2 = plot_piano_roll(midi_data, int(lowest_displayed_pitch), int(highest_displayed_pitch))\n",
        "plt.show(block=False)\n",
        "\n",
        "\n",
        "FluidSynth(\"/content/font.sf2\", 16000).midi_to_audio('/content/output.mid', '/content/output.wav')\n",
        "Audio('/content/output.wav', rate=16000)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YVWEhUj1cg7N",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Plot Positional Embeddings\n",
        "import torchvision\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "# visualize some of the learned positional embeddings, maybe they contain structure\n",
        "plt.figure(figsize=(18, 1))  \n",
        "ci = model.pos_emb.data[0, :, 0].cpu()\n",
        "zci = torch.cat((torch.tensor([0.0]), ci)) # pre-cat a zero\n",
        "plt.imshow(zci.view(1, block_size+1).numpy())\n",
        "plt.axis('off')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2Snu3fb4N-Nd",
        "colab_type": "text"
      },
      "source": [
        "#Congrats! :) You did it :)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vbtJ6PPpWxlX",
        "colab_type": "text"
      },
      "source": [
        "## Save the model to Google Drive (Standard GD connect code)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3QpUrkARWvrP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}