Super_Piano_3_Alt.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Super Piano 3 Alt.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/asigalov61/SuperPiano/blob/master/Super_Piano_3_Alt.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9opKSK2RSDRg"
      },
      "source": [
        "# Super Piano 3: Google Music Transformer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cEbn_ZT08Qjm"
      },
      "source": [
        "!pip install pretty_midi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rwhwdGvU5snk"
      },
      "source": [
        "import pickle\n",
        "import os\n",
        "import sys\n",
        "import math\n",
        "import random\n",
        "import pickle\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "print(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5jsD-2R6OHtW"
      },
      "source": [
        "## Hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D3cHKI6gOGB4"
      },
      "source": [
        "d_model = 512\n",
        "nhead = 8\n",
        "dim_feedforward = 1024\n",
        "dropout = 0.2\n",
        "num_layer = 6\n",
        "batch_size = 8\n",
        "sequence_length = 1024\n",
        "warmup_steps = 4000\n",
        "pad_token = 1   \n",
        "#vocabulary_size = 388 # depends on the dataset"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gsOdPRHj4iw-",
        "cellView": "form"
      },
      "source": [
        "#@title Functions for MIDI encoder/decoder routines\n",
        "#Import library for encoder function\n",
        "import pretty_midi\n",
        "\n",
        "RANGE_NOTE_ON = 128\n",
        "RANGE_NOTE_OFF = 128\n",
        "RANGE_VEL = 32\n",
        "RANGE_TIME_SHIFT = 100\n",
        "\n",
        "START_IDX = {\n",
        "    'note_on': 0,\n",
        "    'note_off': RANGE_NOTE_ON,\n",
        "    'time_shift': RANGE_NOTE_ON + RANGE_NOTE_OFF,\n",
        "    'velocity': RANGE_NOTE_ON + RANGE_NOTE_OFF + RANGE_TIME_SHIFT\n",
        "}\n",
        "\n",
        "\n",
        "class SustainAdapter:\n",
        "    def __init__(self, time, type):\n",
        "        self.start =  time\n",
        "        self.type = type\n",
        "\n",
        "\n",
        "class SustainDownManager:\n",
        "    def __init__(self, start, end):\n",
        "        self.start = start\n",
        "        self.end = end\n",
        "        self.managed_notes = []\n",
        "        self._note_dict = {} # key: pitch, value: note.start\n",
        "\n",
        "    def add_managed_note(self, note: pretty_midi.Note):\n",
        "        self.managed_notes.append(note)\n",
        "\n",
        "    def transposition_notes(self):\n",
        "        for note in reversed(self.managed_notes):\n",
        "            try:\n",
        "                note.end = self._note_dict[note.pitch]\n",
        "            except KeyError:\n",
        "                note.end = max(self.end, note.end)\n",
        "            self._note_dict[note.pitch] = note.start\n",
        "\n",
        "\n",
        "# Divided note by note_on, note_off\n",
        "class SplitNote:\n",
        "    def __init__(self, type, time, value, velocity):\n",
        "        ## type: note_on, note_off\n",
        "        self.type = type\n",
        "        self.time = time\n",
        "        self.velocity = velocity\n",
        "        self.value = value\n",
        "\n",
        "    def __repr__(self):\n",
        "        return '<[SNote] time: {} type: {}, value: {}, velocity: {}>'\\\n",
        "            .format(self.time, self.type, self.value, self.velocity)\n",
        "\n",
        "\n",
        "class Event:\n",
        "    def __init__(self, event_type, value):\n",
        "        self.type = event_type\n",
        "        self.value = value\n",
        "\n",
        "    def __repr__(self):\n",
        "        return '<Event type: {}, value: {}>'.format(self.type, self.value)\n",
        "\n",
        "    def to_int(self):\n",
        "        return START_IDX[self.type] + self.value\n",
        "\n",
        "    @staticmethod\n",
        "    def from_int(int_value):\n",
        "        info = Event._type_check(int_value)\n",
        "        return Event(info['type'], info['value'])\n",
        "\n",
        "    @staticmethod\n",
        "    def _type_check(int_value):\n",
        "        range_note_on = range(0, RANGE_NOTE_ON)\n",
        "        range_note_off = range(RANGE_NOTE_ON, RANGE_NOTE_ON+RANGE_NOTE_OFF)\n",
        "        range_time_shift = range(RANGE_NOTE_ON+RANGE_NOTE_OFF,RANGE_NOTE_ON+RANGE_NOTE_OFF+RANGE_TIME_SHIFT)\n",
        "\n",
        "        valid_value = int_value\n",
        "\n",
        "        if int_value in range_note_on:\n",
        "            return {'type': 'note_on', 'value': valid_value}\n",
        "        elif int_value in range_note_off:\n",
        "            valid_value -= RANGE_NOTE_ON\n",
        "            return {'type': 'note_off', 'value': valid_value}\n",
        "        elif int_value in range_time_shift:\n",
        "            valid_value -= (RANGE_NOTE_ON + RANGE_NOTE_OFF)\n",
        "            return {'type': 'time_shift', 'value': valid_value}\n",
        "        else:\n",
        "            valid_value -= (RANGE_NOTE_ON + RANGE_NOTE_OFF + RANGE_TIME_SHIFT)\n",
        "            return {'type': 'velocity', 'value': valid_value}\n",
        "\n",
        "\n",
        "def _divide_note(notes):\n",
        "    result_array = []\n",
        "    notes.sort(key=lambda x: x.start)\n",
        "\n",
        "    for note in notes:\n",
        "        on = SplitNote('note_on', note.start, note.pitch, note.velocity)\n",
        "        off = SplitNote('note_off', note.end, note.pitch, None)\n",
        "        result_array += [on, off]\n",
        "    return result_array\n",
        "\n",
        "\n",
        "def _merge_note(snote_sequence):\n",
        "    note_on_dict = {}\n",
        "    result_array = []\n",
        "\n",
        "    for snote in snote_sequence:\n",
        "        # print(note_on_dict)\n",
        "        if snote.type == 'note_on':\n",
        "            note_on_dict[snote.value] = snote\n",
        "        elif snote.type == 'note_off':\n",
        "            try:\n",
        "                on = note_on_dict[snote.value]\n",
        "                off = snote\n",
        "                if off.time - on.time == 0:\n",
        "                    continue\n",
        "                result = pretty_midi.Note(on.velocity, snote.value, on.time, off.time)\n",
        "                result_array.append(result)\n",
        "            except:\n",
        "                print('info removed pitch: {}'.format(snote.value))\n",
        "    return result_array\n",
        "\n",
        "\n",
        "def _snote2events(snote: SplitNote, prev_vel: int):\n",
        "    result = []\n",
        "    if snote.velocity is not None:\n",
        "        modified_velocity = snote.velocity // 4\n",
        "        if prev_vel != modified_velocity:\n",
        "            result.append(Event(event_type='velocity', value=modified_velocity))\n",
        "    result.append(Event(event_type=snote.type, value=snote.value))\n",
        "    return result\n",
        "\n",
        "\n",
        "def _event_seq2snote_seq(event_sequence):\n",
        "    timeline = 0\n",
        "    velocity = 0\n",
        "    snote_seq = []\n",
        "\n",
        "    for event in event_sequence:\n",
        "        if event.type == 'time_shift':\n",
        "            timeline += ((event.value+1) / 100)\n",
        "        if event.type == 'velocity':\n",
        "            velocity = event.value * 4\n",
        "        else:\n",
        "            snote = SplitNote(event.type, timeline, event.value, velocity)\n",
        "            snote_seq.append(snote)\n",
        "    return snote_seq\n",
        "\n",
        "\n",
        "def _make_time_sift_events(prev_time, post_time):\n",
        "    time_interval = int(round((post_time - prev_time) * 100))\n",
        "    results = []\n",
        "    while time_interval >= RANGE_TIME_SHIFT:\n",
        "        results.append(Event(event_type='time_shift', value=RANGE_TIME_SHIFT-1))\n",
        "        time_interval -= RANGE_TIME_SHIFT\n",
        "    if time_interval == 0:\n",
        "        return results\n",
        "    else:\n",
        "        return results + [Event(event_type='time_shift', value=time_interval-1)]\n",
        "\n",
        "\n",
        "def _control_preprocess(ctrl_changes):\n",
        "    sustains = []\n",
        "\n",
        "    manager = None\n",
        "    for ctrl in ctrl_changes:\n",
        "        if ctrl.value >= 64 and manager is None:\n",
        "            # sustain down\n",
        "            manager = SustainDownManager(start=ctrl.time, end=None)\n",
        "        elif ctrl.value < 64 and manager is not None:\n",
        "            # sustain up\n",
        "            manager.end = ctrl.time\n",
        "            sustains.append(manager)\n",
        "            manager = None\n",
        "        elif ctrl.value < 64 and len(sustains) > 0:\n",
        "            sustains[-1].end = ctrl.time\n",
        "    return sustains\n",
        "\n",
        "\n",
        "def _note_preprocess(susteins, notes):\n",
        "    note_stream = []\n",
        "\n",
        "    if susteins:    # if the midi file has sustain controls\n",
        "        for sustain in susteins:\n",
        "            for note_idx, note in enumerate(notes):\n",
        "                if note.start < sustain.start:\n",
        "                    note_stream.append(note)\n",
        "                elif note.start > sustain.end:\n",
        "                    notes = notes[note_idx:]\n",
        "                    sustain.transposition_notes()\n",
        "                    break\n",
        "                else:\n",
        "                    sustain.add_managed_note(note)\n",
        "\n",
        "        for sustain in susteins:\n",
        "            note_stream += sustain.managed_notes\n",
        "    \n",
        "    else:       # else, just push everything into note stream\n",
        "        for note_idx, note in enumerate(notes):\n",
        "            note_stream.append(note)\n",
        "\n",
        "    note_stream.sort(key= lambda x: x.start)\n",
        "    return note_stream\n",
        "\n",
        "\n",
        "def encode_midi(file_path):\n",
        "    events = []\n",
        "    notes = []\n",
        "    mid = pretty_midi.PrettyMIDI(midi_file=file_path)\n",
        "\n",
        "    for inst in mid.instruments:\n",
        "        inst_notes = inst.notes\n",
        "        # ctrl.number is the number of sustain control. If you want to know abour the number type of control,\n",
        "        # see https://www.midi.org/specifications-old/item/table-3-control-change-messages-data-bytes-2\n",
        "        ctrls = _control_preprocess([ctrl for ctrl in inst.control_changes if ctrl.number == 64])\n",
        "        notes += _note_preprocess(ctrls, inst_notes)\n",
        "\n",
        "    dnotes = _divide_note(notes)\n",
        "\n",
        "    # print(dnotes)\n",
        "    dnotes.sort(key=lambda x: x.time)\n",
        "    # print('sorted:')\n",
        "    # print(dnotes)\n",
        "    cur_time = 0\n",
        "    cur_vel = 0\n",
        "    for snote in dnotes:\n",
        "        events += _make_time_sift_events(prev_time=cur_time, post_time=snote.time)\n",
        "        events += _snote2events(snote=snote, prev_vel=cur_vel)\n",
        "        # events += _make_time_sift_events(prev_time=cur_time, post_time=snote.time)\n",
        "\n",
        "        cur_time = snote.time\n",
        "        cur_vel = snote.velocity\n",
        "\n",
        "    return [e.to_int() for e in events]\n",
        "\n",
        "\n",
        "def decode_midi(idx_array, file_path=None):\n",
        "    event_sequence = [Event.from_int(idx) for idx in idx_array]\n",
        "    # print(event_sequence)\n",
        "    snote_seq = _event_seq2snote_seq(event_sequence)\n",
        "    note_seq = _merge_note(snote_seq)\n",
        "    note_seq.sort(key=lambda x:x.start)\n",
        "\n",
        "    mid = pretty_midi.PrettyMIDI()\n",
        "    # if want to change instument, see https://www.midi.org/specifications/item/gm-level-1-sound-set\n",
        "    instument = pretty_midi.Instrument(0, False, \"Composed by Super Piano Music Transformer AI\")\n",
        "    instument.notes = note_seq\n",
        "\n",
        "    mid.instruments.append(instument)\n",
        "    if file_path is not None:\n",
        "        mid.write(file_path)\n",
        "    return mid"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z3_Fd7tRDOOC"
      },
      "source": [
        "!mkdir '/content/midis/'\n",
        "!mkdir '/content/save/'\n",
        "!mkdir '/content/train/'\n",
        "!mkdir '/content/test/'\n",
        "\n",
        "midi_dir_path = '/content/midis/'\n",
        "save_dir_path = '/content/save/'\n",
        "train_dir_path = '/content/train/'\n",
        "test_dir_path = '/content/test/'\n",
        "extension = ['.mid', '.midi']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YYUFXcjlhBwV",
        "cellView": "form"
      },
      "source": [
        "#@title Download Google Magenta MAESTRO v.2.0.0 Piano MIDI Dataset (~1300 MIDIs)\n",
        "%cd /content/midis\n",
        "!wget 'https://storage.googleapis.com/magentadata/datasets/maestro/v2.0.0/maestro-v2.0.0-midi.zip'\n",
        "!unzip -j maestro-v2.0.0-midi.zip\n",
        "%cd /content"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ngsTYWHspwyq",
        "cellView": "form"
      },
      "source": [
        "#@title Encode MIDI files\n",
        "\n",
        "def encode_midi_files(midi_dir_path, save_dir_path, extension):\n",
        "  #create directory for saving files\n",
        "  os.makedirs(save_dir_path, exist_ok=True)\n",
        "  #get all midi files from midi_directory\n",
        "  for file in os.listdir(midi_dir_path):\n",
        "    if file.endswith(tuple(extension)):\n",
        "        print(os.path.join(midi_dir_path, file))\n",
        "        print(file + ' is being processed', flush=True)\n",
        "        try:\n",
        "          encoded_file = encode_midi(midi_dir_path+file)\n",
        "        except KeyboardInterrupt:\n",
        "            print(' Stopped by keyboard')\n",
        "            return\n",
        "        except EOFError:\n",
        "            print('EOF Error')\n",
        "            return\n",
        "        with open(save_dir_path+file+'.encoded', 'wb') as f:\n",
        "            pickle.dump(encoded_file, f)\n",
        "\n",
        "encode_midi_files(midi_dir_path, save_dir_path, extension) # only use once"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lMJCA75_fu9M",
        "cellView": "form"
      },
      "source": [
        "#@title Create and Split MIDI DataSet\n",
        "# for more convenience we split training data and test data in two different folders\n",
        "import shutil\n",
        "\n",
        "def create_dataset(save_dir_path, split_ratio=0.9):\n",
        "    dataset = [file for file in os.listdir(save_dir_path)]\n",
        "    np.random.shuffle(dataset)\n",
        "\n",
        "    train_set = dataset[:int(len(dataset) * split_ratio)]    \n",
        "    test_set = dataset[int(len(dataset) * split_ratio):]\n",
        "    \n",
        "    shutil.rmtree(train_dir_path)\n",
        "    shutil.rmtree(test_dir_path)\n",
        "\n",
        "    os.makedirs(train_dir_path, exist_ok=True)\n",
        "    os.makedirs(test_dir_path, exist_ok=True)\n",
        "\n",
        "    for file in os.listdir(save_dir_path):\n",
        "        if os.stat(save_dir_path+file).st_size != 0:\n",
        "            if file in test_set:\n",
        "                shutil.copyfile(save_dir_path+file, test_dir_path+file)\n",
        "            else:\n",
        "                shutil.copyfile(save_dir_path+file, train_dir_path+file)\n",
        "\n",
        "create_dataset(save_dir_path, split_ratio=0.9) # only use once"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "waBqGI-AP7qo",
        "cellView": "form"
      },
      "source": [
        "#@title Setup all functions and procedures\n",
        "\n",
        "def load_dataset(train_dir_path, test_dir_path):\n",
        "    #load all encoded file\n",
        "    train_set = [file for file in os.listdir(train_dir_path)]\n",
        "    test_set = [file for file in os.listdir(test_dir_path)]\n",
        "\n",
        "    return train_set, test_set\n",
        "\n",
        "def generate_batch(dataset, dir_path, sequence_length=1024, batch_size=4):\n",
        "    sequence_length += 1\n",
        "    batch_midi = []\n",
        "    while len(batch_midi) < batch_size:\n",
        "        file = random.choice(dataset)\n",
        "        #if the midi contains more sequence that the sequence length\n",
        "        try:\n",
        "            with open(dir_path+file, 'rb') as f:\n",
        "                data = pickle.load(f)\n",
        "        except:\n",
        "            print(dir_path+file + \" file not found.\")\n",
        "        if sequence_length <= len(data):\n",
        "            begin_index = random.randrange(0, len(data) - sequence_length)\n",
        "            data = data[begin_index:begin_index + sequence_length]\n",
        "            batch_midi.append(data)\n",
        "    batch_midi = torch.Tensor(batch_midi)\n",
        "    inputs = batch_midi[:, :-1]\n",
        "    labels = batch_midi[:, 1:]\n",
        "    return inputs, labels\n",
        "\n",
        "\n",
        "class MusicMultiheadAttention(torch.nn.MultiheadAttention):\n",
        "    def __init__(self, embed_dim, nhead, dropout=0.1, bias=True, add_bias_kv=False, \n",
        "                 add_zero_attn=False, kdim=None, vdim=None):\n",
        "        \n",
        "        torch.nn.MultiheadAttention.__init__(self, embed_dim, nhead, dropout=0.1, \n",
        "                                             bias=True, add_bias_kv=False, \n",
        "                                             add_zero_attn=False, kdim=None, vdim=None)\n",
        "        \n",
        "        self.embed_dim = embed_dim\n",
        "        self.weights_q = torch.nn.Linear(self.embed_dim, self.embed_dim)\n",
        "        self.weights_k = torch.nn.Linear(self.embed_dim, self.embed_dim)\n",
        "        self.weights_v = torch.nn.Linear(self.embed_dim, self.embed_dim)\n",
        "        self.weights_o = torch.nn.Linear(self.embed_dim, self.embed_dim)\n",
        "\n",
        "    def forward(self, query, key, value, key_padding_mask=None,\n",
        "                need_weights=True, attn_mask=None):\n",
        "        Q, K, V = self.transform_input(query, key, value)\n",
        "        # Reshaping the matrices \n",
        "        # Each L × D query, key, and value matrix is then split into H L × D \n",
        "        # h_D parts or attention heads, indexed by h, and with dimension D_h = D/H\n",
        "        Q = self.matrix_to_heads(Q)\n",
        "        K = self.matrix_to_heads(K)\n",
        "        V = self.matrix_to_heads(V)\n",
        "        \n",
        "        # learning a separate relative position embedding Er of shape (H, L, Dh)\n",
        "        Er = torch.randn([self.num_heads, query.size(1), self.head_dim], requires_grad=False).to(device)\n",
        "\n",
        "        # we transpose the two last dimensions of Er to realize Q*Er^T \n",
        "        QEr = torch.matmul(Q, torch.transpose(Er,1,2))\n",
        "        # QEr of shape (B, H, L, L)     \n",
        "        # QEr = torch.einsum('bhld,ld->bhll', [Q, Er])\n",
        "\n",
        "        # 1. Pad a dummy column vector of length L before the leftmost column.\n",
        "        QEr = torch.nn.functional.pad(QEr, (1,0), mode=\"constant\", value=0)\n",
        "\n",
        "        # 2. Reshape the matrix to have shape (L+1, L). \n",
        "        QEr = torch.reshape(QEr, [QEr.size(0), QEr.size(1), QEr.size(3), QEr.size(2)])\n",
        "        \n",
        "        # 3. Slice that matrix to retain only the last l rows and all the columns, \n",
        "        # resulting in a (L, L) matrix again, but now absolute-by-absolute indexed, \n",
        "        # which is the S rel that we need.\n",
        "        S_rel = QEr[:,:,1:,:]\n",
        "\n",
        "        z_attention = self.attention(Q, K, V, S_rel, attn_mask)\n",
        "        z_attention = self.weights_o(z_attention)\n",
        "        # Masking can be added and Dropout ?\n",
        "\n",
        "        return z_attention\n",
        "\n",
        "    def attention(self, Q, K, V, S, mask):\n",
        "        # Dh = self.head_dim // self.num_heads\n",
        "        logits = torch.add(torch.matmul(Q, torch.transpose(K, 2, 3)), S) / math.sqrt((self.head_dim // self.num_heads))\n",
        "        # print(\"logits : \", logits.size())\n",
        "        # print(\"mask : \", mask.size())\n",
        "        if mask is not None:\n",
        "        #    mask = mask.unsqueeze(1) #shape of mask must be broadcastable with shape of underlying tensor\n",
        "            logits = logits.masked_fill(mask == 0, -1e9) #masked_fill fills elements of scores with -1e9 where mask == 0\n",
        "        #if mask is not None:\n",
        "        #    logits += (mask.to(torch.int64) * -1e9).to(logits.dtype)        \n",
        "            \n",
        "        activation = F.softmax(logits, -1)\n",
        "        attention = torch.matmul(activation, V)\n",
        "        attention = torch.reshape(attention, (attention.size(0), -1, self.embed_dim))\n",
        "        return attention\n",
        "    \n",
        "    def matrix_to_heads(self, qkv):\n",
        "        '''\n",
        "            Takes a query/key/value (qkv) matrix and reshapes it to  B * H * L * D_h heads \n",
        "            with dimension D_h = D/H\n",
        "        '''\n",
        "        batch_size_q = qkv.size(0)\n",
        "        #qkv = torch.reshape(qkv, (batch_size_q, qkv.size(0), self.num_heads, self.head_dim))\n",
        "        qkv = torch.reshape(qkv, (batch_size_q, self.num_heads, qkv.size(1), self.head_dim))\n",
        "        return qkv\n",
        "\n",
        "    def transform_input(self, query, key, value):\n",
        "        '''\n",
        "            Transforming the input vector, X, of LxD dimension \n",
        "            into \n",
        "                queries: Q = XW^Q \n",
        "                keys:    K = XW^K\n",
        "            and values:  V = XW^V\n",
        "            which are all DxD square matrices.\n",
        "        '''\n",
        "        return self.weights_q(query), self.weights_k(key), self.weights_v(value)\n",
        "    \n",
        "\n",
        "class MusicTransformerEncoderLayer(torch.nn.TransformerEncoderLayer):\n",
        "    def __init__(self, d_model, nhead, dim_feedforward=2048,\n",
        "                 dropout=0.1, activation=\"relu\"):\n",
        "        torch.nn.TransformerEncoderLayer.__init__(self, d_model, nhead)\n",
        "        self.d_model = d_model\n",
        "        # OverRide\n",
        "        self.self_attn = MusicMultiheadAttention(d_model, nhead)\n",
        "\n",
        "class MusicTransformerDecoderLayer(torch.nn.TransformerDecoderLayer):\n",
        "    def __init__(self, d_model, nhead, dim_feedforward=2048,\n",
        "                 dropout=0.1, activation=\"relu\"):\n",
        "        torch.nn.TransformerDecoderLayer.__init__(self, d_model,nhead)\n",
        "        self.d_model = d_model\n",
        "        # OverRide\n",
        "        self.self_attn = MusicMultiheadAttention(d_model, nhead)\n",
        "\n",
        "class MusicTransformerEncoder(torch.nn.TransformerEncoder):\n",
        "    def __init__(self, encoder_layer, vocabulary_size=390, num_encoder_layers=6, normalization=None):\n",
        "        super().__init__(encoder_layer, num_encoder_layers, normalization)\n",
        "        self.d_model = encoder_layer.d_model\n",
        "        self.vocabulary_size = vocabulary_size\n",
        "        self.dropout = torch.nn.Dropout(encoder_layer.dropout.p)\n",
        "        self.embedding = torch.nn.Embedding(num_embeddings=self.vocabulary_size, embedding_dim=self.d_model)\n",
        "    \n",
        "    def forward(self, src, mask=None, src_key_padding_mask=None):\n",
        "        \n",
        "        pos_encoding = DynamicPositionEmbedding(self.d_model, src.size(1))\n",
        "        \n",
        "        src = math.sqrt(self.d_model) * self.embedding(src.to(torch.long).to(device))\n",
        "        src = pos_encoding(src)\n",
        "        src = self.dropout(src)\n",
        "\n",
        "        return super().forward(src, mask=mask, src_key_padding_mask=src_key_padding_mask)\n",
        "\n",
        "\n",
        "class MusicTransformerDecoder(torch.nn.TransformerDecoder):\n",
        "    def __init__(self, decoder_layer, vocabulary_size=390, num_decoder_layers=6, normalization=None):\n",
        "        super().__init__(decoder_layer, num_decoder_layers, normalization)\n",
        "        self.d_model = decoder_layer.d_model\n",
        "        self.vocabulary_size = vocabulary_size\n",
        "        self.dropout = torch.nn.Dropout(decoder_layer.dropout.p)\n",
        "        self.embedding = torch.nn.Embedding(num_embeddings=self.vocabulary_size, embedding_dim=self.d_model)\n",
        "\n",
        "    def forward(self, tgt, memory, tgt_mask=None, \n",
        "                memory_mask=None, tgt_key_padding_mask=None,\n",
        "                memory_key_padding_mask=None):        \n",
        "                \n",
        "        pos_encoding = DynamicPositionEmbedding(self.d_model, tgt.size(1))\n",
        "\n",
        "        tgt = pos_encoding(math.sqrt(self.d_model) * self.embedding(tgt.to(torch.long).to(device)))\n",
        "        tgt = self.dropout(tgt)\n",
        "\n",
        "        return super().forward(tgt, memory, tgt_mask=tgt_mask, \n",
        "                memory_mask=memory_mask, tgt_key_padding_mask=tgt_key_padding_mask,\n",
        "                memory_key_padding_mask=memory_key_padding_mask)\n",
        "\n",
        "\n",
        "class MusicTransformer(torch.nn.modules.Transformer):\n",
        "    def __init__(self, d_model=512, nhead=8, vocabulary_size=388,\n",
        "                 num_encoder_layers=6, num_decoder_layers=6, \n",
        "                 dim_feedforward=2048, dropout=0.1, activation=\"relu\", \n",
        "                 custom_encoder=None, custom_decoder=None):\n",
        "        \n",
        "        super().__init__(d_model=d_model, nhead=nhead, \n",
        "                         num_encoder_layers=num_encoder_layers,\n",
        "                         num_decoder_layers=num_decoder_layers, \n",
        "                         dim_feedforward=dim_feedforward, dropout=dropout, activation=activation, \n",
        "                         custom_encoder=custom_encoder, custom_decoder=custom_decoder)\n",
        "        \n",
        "        self.vocabulary_size = vocabulary_size\n",
        "        ###        \n",
        "        self.fc = torch.nn.Linear(self.d_model, self.vocabulary_size)\n",
        "        self._reset_parameters()\n",
        "\n",
        "    def forward(self, src, tgt, src_mask=None, tgt_mask=None,\n",
        "                memory_mask=None, src_key_padding_mask=None,\n",
        "                tgt_key_padding_mask=None, memory_key_padding_mask=None):\n",
        "\n",
        "        #if src.size(1) != tgt.size(1):\n",
        "        #    raise RuntimeError(\"the batch number of src and tgt must be equal\")\n",
        "\n",
        "        #if src.size(2) != self.d_model or tgt.size(2) != self.d_model:\n",
        "        #    raise RuntimeError(\"the feature number of src and tgt must be equal to d_model\")\n",
        "            \n",
        "        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)\n",
        "        output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,\n",
        "                              tgt_key_padding_mask=tgt_key_padding_mask,\n",
        "                              memory_key_padding_mask=memory_key_padding_mask)\n",
        "\n",
        "        output = self.fc(output)        \n",
        "        \n",
        "        return output\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "class DynamicPositionEmbedding(torch.nn.Module):\n",
        "    def __init__(self, embedding_dim, max_seq=1024):\n",
        "        super().__init__()\n",
        "        embed_sinusoid_list = np.array([[\n",
        "            [\n",
        "                math.sin(\n",
        "                    pos * math.exp(-math.log(10000) * i/embedding_dim) *\n",
        "                    math.exp(math.log(10000)/embedding_dim * (i % 2)) + 0.5 * math.pi * (i % 2)\n",
        "                )\n",
        "                for i in range(embedding_dim)\n",
        "            ]\n",
        "            for pos in range(max_seq)\n",
        "        ]])\n",
        "        self.positional_embedding = embed_sinusoid_list\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + torch.from_numpy(self.positional_embedding[:, :x.size(1), :]).to(x.device, dtype=x.dtype)\n",
        "        return x\n",
        "\n",
        "\n",
        "class PositionalEncoder(torch.nn.Module):\n",
        "    def __init__(self, d_model, max_seq_len=1024):\n",
        "        super().__init__()\n",
        "        self.d_model = d_model\n",
        "        pe = torch.zeros(max_seq_len, d_model)\n",
        "        for pos in range(max_seq_len):\n",
        "            for i in range(0, d_model, 2):\n",
        "                pe[pos, i] = \\\n",
        "                    math.sin(pos / (10000 ** ((2 * i) / d_model)))\n",
        "                pe[pos, i + 1] = \\\n",
        "                    math.cos(pos / (10000 ** ((2 * (i + 1)) / d_model)))\n",
        "        pe = pe.unsqueeze(0)\n",
        "        self.register_buffer('pe', pe)\n",
        "\n",
        "    def forward(self, x):\n",
        "        with torch.no_grad():\n",
        "            x = x * math.sqrt(self.d_model)\n",
        "            seq_len = x.size(1)\n",
        "            self.pe.to(device)\n",
        "            print(\"self.pe : \", self.pe.device.type)\n",
        "            pe = self.pe[:, :seq_len]\n",
        "            print(\"pe : \", pe.device.type)\n",
        "            x = x + pe\n",
        "            return x   "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YfU5uhNwKx6u",
        "cellView": "form"
      },
      "source": [
        "#@title More functions\n",
        "from torch.utils.data.dataset import Dataset\n",
        "\n",
        "def tensorFromSequence(sequence):\n",
        "    \"\"\"\n",
        "    Generate tensors from the sequence in numpy.\n",
        "    \"\"\"\n",
        "    output = torch.tensor(sequence).long()\n",
        "\n",
        "    return output\n",
        "\n",
        "\n",
        "def PrepareData(npz_file, split='train', L=1024):\n",
        "    \"\"\"\n",
        "    Function to prepare the data into pairs (input, target).\n",
        "    Adds [PAD], [SOS] and [EOS] tokens into the data,\n",
        "    where [PAD]=1, [SOS]=2, [EOS]=3.\n",
        "    Limits the sequence to length of L.\n",
        "    \"\"\"\n",
        "    print(\"Preparing data for\",split,\"split...\")\n",
        "    # Load in the data\n",
        "    full_data = np.load(npz_file, fix_imports=True, encoding=\"latin1\", allow_pickle=True)\n",
        "    data = full_data[split]\n",
        "\n",
        "    # Extract the vocab from file\n",
        "    vocab = GenerateVocab(npz_file)\n",
        "    # Generate new vocab to map to later\n",
        "    new_vocab = np.arange(len(vocab))\n",
        "\n",
        "    # Initialize the tokens\n",
        "    pad_token = np.array([[1]])\n",
        "\n",
        "    # Repeat for all samples in data\n",
        "    pairs = []\n",
        "    for samples in data:\n",
        "        # Serialise the dataset so that the resulting sequence is\n",
        "        # S_1 A_1 T_1 B_1, S_2 A_2 T_2 B_2, ...\n",
        "\n",
        "        # Generate input\n",
        "        input_seq = samples.flatten()\n",
        "\n",
        "        # Cut off the samples so that it has length of 1024\n",
        "        if(len(input_seq) >= L):\n",
        "            # input_seq = input_seq[:L-1]\n",
        "            input_seq = input_seq[:L]\n",
        "\n",
        "        # Set the NaN values to 0 and reshape accordingly\n",
        "        input_seq = np.nan_to_num(input_seq.reshape(1,input_seq.size))\n",
        "\n",
        "        # Generate target\n",
        "        output_seq = input_seq[:,1:]\n",
        "\n",
        "        # For both sequences, pad to sequence length L\n",
        "        pad_array = pad_token * np.ones((1,L-input_seq.shape[1]))\n",
        "        input_seq = np.append(input_seq, pad_array,axis=1)\n",
        "        pad_array = pad_token * np.ones((1,L-output_seq.shape[1]))\n",
        "        output_seq = np.append(output_seq, pad_array,axis=1)\n",
        "\n",
        "        # Map the pitch value to int values below vocab size\n",
        "        for i, val in enumerate(vocab):\n",
        "            input_seq[input_seq==val] = new_vocab[i]\n",
        "            output_seq[output_seq==val] = new_vocab[i]\n",
        "\n",
        "        # Make it into a pair\n",
        "        pair = [input_seq, output_seq]\n",
        "\n",
        "        # Combine all pairs into one big list of pairs\n",
        "        pairs.append(pair)\n",
        "\n",
        "    print(\"Generated data pairs.\")\n",
        "    return np.array(pairs)\n",
        "\n",
        "def GenerateVocab(npz_file):\n",
        "    \"\"\"\n",
        "    Generate vocabulary for the dataset including the custom tokens.\n",
        "    \"\"\"\n",
        "    full_data = np.load(npz_file, fix_imports=True, encoding=\"latin1\", allow_pickle=True)\n",
        "    train_data = full_data['train']\n",
        "    validation_data = full_data['valid']\n",
        "    test_data = full_data['test']\n",
        "\n",
        "    combined_data = np.concatenate((train_data, validation_data, test_data))\n",
        "\n",
        "    vocab = np.nan\n",
        "    for sequences in combined_data:\n",
        "        vocab = np.append(vocab,np.unique(sequences))\n",
        "\n",
        "    vocab = np.unique(vocab)\n",
        "    vocab = vocab[~np.isnan(vocab)]\n",
        "    vocab = np.append([0,1],vocab)\n",
        "    return vocab \n",
        " \n",
        "\n",
        "def batched_learning(train,batch_size):\n",
        "    for i in range(0, len(train), batch_size):\n",
        "        train1 = train[i:i + batch_size]\n",
        "        yield train1[:,0],train1[:,1]\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NbnfRnF2P0d3"
      },
      "source": [
        "##3 - Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CtKgQYKuP4-V",
        "cellView": "form"
      },
      "source": [
        "#@title Definitions and Functions\n",
        " \n",
        "# vocabulary_size depends on the midi encodding\n",
        "# ~> 388(+2) for encoded_midi / epiano compt \n",
        "# ~> 46(+2) for encoded_midi / epiano compt                 \n",
        "\n",
        "# DEFINING THE MODEL\n",
        "\n",
        "vocabulary_size = 390\n",
        "\n",
        "normalization = torch.nn.LayerNorm(d_model)\n",
        "\n",
        "custom_encoder_layer = MusicTransformerEncoderLayer(d_model=d_model, nhead=nhead, \n",
        "                                               dim_feedforward=dim_feedforward, \n",
        "                                               dropout=dropout, activation=\"relu\")\n",
        "\n",
        "custom_decoder_layer = MusicTransformerDecoderLayer(d_model=d_model, nhead=nhead, \n",
        "                                               dim_feedforward=dim_feedforward, \n",
        "                                               dropout=dropout, activation=\"relu\")\n",
        "\n",
        "custom_encoder = MusicTransformerEncoder(custom_encoder_layer, vocabulary_size, num_layer, normalization)\n",
        "custom_decoder = MusicTransformerDecoder(custom_decoder_layer, vocabulary_size, num_layer, normalization)\n",
        "\n",
        "model = MusicTransformer(d_model=d_model, nhead=nhead, \n",
        "                         vocabulary_size=vocabulary_size, \n",
        "                         num_encoder_layers=num_layer, \n",
        "                         num_decoder_layers=num_layer, \n",
        "                         dim_feedforward=dim_feedforward, \n",
        "                         dropout=dropout, activation=\"relu\", \n",
        "                         custom_encoder=custom_encoder, \n",
        "                         custom_decoder=custom_decoder)\n",
        "# Give model to the current device (hopefully cuda)\n",
        "model.to(device)\n",
        "\n",
        "# Optimizer\n",
        "# Adam optimizer [20] with β 1 = 0.9, β 2 = 0.98 and \u000f = 10 −9\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.98), eps=1e-09, weight_decay=1e-4)\n",
        "\n",
        "# Define a scheduler to vary the learning rate\n",
        "\n",
        "class Scheduler:\n",
        "    def __init__(self, optimizer, d_model=d_model, warmup_steps=4000):\n",
        "        self.optimizer = optimizer\n",
        "        self.d_model = d_model\n",
        "        self.step_num = 1\n",
        "        self.l_rate = 0\n",
        "        self.warmup_steps = warmup_steps\n",
        "\n",
        "    def step(self):        \n",
        "        # increment step\n",
        "        self.step_num += 1\n",
        "\n",
        "        # compute new learning rate        \n",
        "        self.l_rate = self.d_model**(-.5) * min(self.step_num**(-.5), self.step_num * self.warmup_steps**(-1.5))\n",
        "\n",
        "        # update optimizer learning rate\n",
        "        for p in optimizer.param_groups:\n",
        "            p['lr'] = self.l_rate\n",
        "\n",
        "        # update the weights in the network\n",
        "        self.optimizer.step()\n",
        "\n",
        "\n",
        "# See if it is possible to do it using lr_scheduler.LambdaLR lr_scheduler.StepLR\n",
        "# scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1)\n",
        "scheduler = Scheduler(optimizer, d_model, warmup_steps)\n",
        "\n",
        "\n",
        "### Mask Generation from https://github.com/COMP6248-Reproducability-Challenge/music-transformer-comp6248/blob/master/MaskGen.py\n",
        "\n",
        "# Filename: MaskGen.py\n",
        "# Date Created: 15-Mar-2019 2:42:12 pm\n",
        "# Description: Functions used to generate masks w.r.t. given inputs.\n",
        "from torch.autograd import Variable\n",
        "\n",
        "def nopeak_mask(size):\n",
        "    np_mask = np.triu(np.ones((1, size, size)), k=1).astype('uint8')\n",
        "    np_mask =  Variable(torch.from_numpy(np_mask) == 0).to(device)\n",
        "    return np_mask\n",
        "\n",
        "\n",
        "def create_masks(src, trg, pad_token):\n",
        "    src_mask = (src != pad_token).unsqueeze(-2).to(device)\n",
        "\n",
        "    if trg is not None:\n",
        "        trg_mask = (trg != pad_token).unsqueeze(-2).to(device)\n",
        "        size = trg.size(1) # get seq_len for matrix\n",
        "        np_mask = nopeak_mask(size)\n",
        "        trg_mask = trg_mask & np_mask\n",
        "    else:\n",
        "        trg_mask = None\n",
        "    return src_mask, trg_mask\n",
        "\n",
        "\n",
        "def count_nonpad_tokens(target):\n",
        "    nonpads = (target != 1).squeeze()\n",
        "    ntokens = torch.sum(nonpads)\n",
        "    return ntokens      "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m_s9IN7xkia9"
      },
      "source": [
        "### Training with data from Midi_Encoded "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JOSMK5LN1_TB",
        "cellView": "form"
      },
      "source": [
        "#@title Load the existing model checkpoint or define a new model\n",
        "train_data, test_data = load_dataset(train_dir_path, test_dir_path)\n",
        "\n",
        "ckpt_dir = '/content'\n",
        "if not os.path.exists(ckpt_dir):\n",
        "  os.makedirs(ckpt_dir)\n",
        "  \n",
        "best_loss = 10.\n",
        "model_name = 'midi_encoded_6-1'\n",
        "ckpt_path = '/content/'+model_name+'.pt'\n",
        "if os.path.exists(ckpt_path):\n",
        "    ckpt = torch.load(ckpt_path)\n",
        "    try:\n",
        "      model.load_state_dict(ckpt['my_model'])\n",
        "      optimizer.load_state_dict(ckpt['optimizer'])\n",
        "      best_acc = ckpt['best_loss']\n",
        "    except RuntimeError as e:\n",
        "        print('wrong checkpoint')\n",
        "    else:    \n",
        "      print('checkpoint is loaded !')\n",
        "      print('current best loss : %.2f' % best_loss)\n",
        "\n",
        "train_writer = SummaryWriter()\n",
        "test_writer = SummaryWriter()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1XYfecvLVYkh",
        "cellView": "form"
      },
      "source": [
        "#@title Train the model\n",
        "# Training\n",
        "max_epochs = 100\n",
        "n_iter = 0\n",
        "\n",
        "# each 50 iterations we are going to compare the losses\n",
        "total_train_loss = []\n",
        "total_valid_loss = []\n",
        "\n",
        "for e in range(max_epochs):\n",
        "    model.train()\n",
        "    train_loss = []\n",
        "    for b_train in range(len(train_data) // batch_size):\n",
        "        \n",
        "        n_iter += 1\n",
        "        # train phase\n",
        "        # feed data into the network and get outputs.\n",
        "        # feed data into the network and get outputs.\n",
        "        inputs, target = generate_batch(train_data, train_dir_path, sequence_length=sequence_length, batch_size=batch_size)\n",
        "              \n",
        "        # Train on GPU\n",
        "        inputs.to(device)\n",
        "        target.to(device)\n",
        "        ys = target.contiguous().view(-1).to(torch.long).to(device)\n",
        "        \n",
        "        # Create mask for both input and target sequences\n",
        "        input_mask, target_mask = create_masks(torch.reshape(inputs, (batch_size, 1, -1)), torch.reshape(target, (batch_size, 1, -1)), pad_token)        \n",
        "        \n",
        "        # feed data into the network and get outputs.\n",
        "        preds_idx = model(inputs, target, input_mask, target_mask)\n",
        "        \n",
        "        # Flush out gradients computed at the previous step before computing gradients at the current step. \n",
        "        #       Otherwise, gradients would accumulate.\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        # calculate loss\n",
        "        loss = F.cross_entropy(preds_idx.contiguous().view(preds_idx.size(-1), -1).transpose(0,1), ys, ignore_index = pad_token, size_average = False) / (count_nonpad_tokens(ys))\n",
        "\n",
        "        # accumulates the gradient and backprogate loss.\n",
        "        loss.backward()\n",
        "\n",
        "        # performs a parameter update based on the current gradient\n",
        "        scheduler.step()    \n",
        "        \n",
        "        print('\\n====================================================')\n",
        "        print('Epoch/Batch: {}/{}'.format(e, b_train))\n",
        "        print('Train >>>> Loss: {:6.6}'.format(loss))\n",
        "\n",
        "        train_loss.append(loss.item())\n",
        "\n",
        "    print('\\n**************************************************')\n",
        "    print(\"\\n*** Test *** \")\n",
        "    \n",
        "    # Validation step\n",
        "    model.eval()\n",
        "    valid_loss = []\n",
        "    with torch.no_grad():\n",
        "        for b_test in range(len(test_data) // batch_size):        \n",
        "            inputs, target = generate_batch(test_data, test_dir_path, sequence_length=sequence_length, batch_size=batch_size)\n",
        "                  \n",
        "            # Train on GPU\n",
        "            inputs.to(device)\n",
        "            target.to(device)\n",
        "            ys = target.contiguous().view(-1).to(torch.long).to(device)\n",
        "            \n",
        "            # Create mask for both input and target sequences\n",
        "            input_mask, target_mask = create_masks(torch.reshape(inputs, (batch_size, 1,-1)), torch.reshape(target, (batch_size, 1, -1)), pad_token)        \n",
        "\n",
        "            # Feed Forward\n",
        "            preds_validate = model(inputs, target, input_mask, target_mask)\n",
        "            loss = F.cross_entropy(preds_validate.contiguous().view(preds_validate.size(-1), -1).transpose(0,1), ys, \\\n",
        "                                    ignore_index = pad_token, size_average = False) / (count_nonpad_tokens(ys))\n",
        "            valid_loss.append(loss.item())\n",
        "\n",
        "    avg_train_loss = np.mean(train_loss)\n",
        "    avg_valid_loss = np.mean(valid_loss)\n",
        "\n",
        "    total_train_loss.append(avg_train_loss)\n",
        "    total_valid_loss.append(avg_valid_loss)\n",
        "\n",
        "    print(\"[Average Train Loss]: {:6.6}\".format(avg_train_loss))\n",
        "    print(\"[Average Testing Loss]: {:6.6}\".format(avg_valid_loss))\n",
        "\n",
        "    # save checkpoint whenever there is improvement in performance\n",
        "    if avg_valid_loss < best_loss:\n",
        "        best_loss = avg_valid_loss\n",
        "        # Note: optimizer also has states ! don't forget to save them as well.\n",
        "        ckpt = {'my_model':model.state_dict(),\n",
        "                'optimizer':optimizer.state_dict(),\n",
        "                'best_loss':best_loss}\n",
        "        torch.save(ckpt, ckpt_path)\n",
        "        print('checkpoint is saved !')\n",
        "\n",
        "    train_writer.add_scalar('loss/train', avg_train_loss, global_step=n_iter)\n",
        "    test_writer.add_scalar('loss/valid', avg_valid_loss, global_step=n_iter)\n",
        "    print('\\n**************************************************')\n",
        "\n",
        "train_writer = SummaryWriter()\n",
        "test_writer = SummaryWriter()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MTqvzYhl05eK",
        "cellView": "form"
      },
      "source": [
        "#@title Plot loss graphs\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "plt.plot(total_train_loss, label='train loss')\n",
        "plt.plot(total_valid_loss, label='validation loss')\n",
        "plt.legend()\n",
        "plt.xlabel(\"epoch\")\n",
        "plt.ylabel(\"loss\")\n",
        "img_path = os.path.join('/gdrive/My Drive/my_data/library/checkpoints/', \"encoded_midi_6-1.png\")\n",
        "plt.savefig(\"encoded_midi_6-1.png\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zR1FaRoBkvSl"
      },
      "source": [
        "### Training with data from JS Bach Chorales "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F9o7LM7gkfZ2",
        "cellView": "form"
      },
      "source": [
        "#@title Load the existing model checkpoint or define a new model\n",
        "src_data = '/gdrive/My Drive/my_data/library/Jsb16thSeparated.npz'\n",
        "\n",
        "# Generate the vocabulary from the data\n",
        "vocab = GenerateVocab(src_data)\n",
        "vocabulary_size = len(vocab)\n",
        "pad_token = 1\n",
        "\n",
        "# since we change de vocab len the model should be create after this\n",
        "# we just show this as an exemple of training of JS Bach Chorales dataset \n",
        "# thanks to propressing functions from another repository.\n",
        "\n",
        "# Setup the dataset for training split and validation split\n",
        "train_data = PrepareData(src_data ,'train', int(sequence_length))\n",
        "valid_data = PrepareData(src_data ,'valid', int(sequence_length))\n",
        "\n",
        "ckpt_dir = os.path.join(gdrive_root, '/my_data/library/checkpoints/')\n",
        "if not os.path.exists(ckpt_dir):\n",
        "  os.makedirs(ckpt_dir)\n",
        "  \n",
        "best_loss = 10.\n",
        "model_name = 'midi_encoded_6-1'\n",
        "ckpt_path = '/gdrive/My Drive/my_data/library/checkpoints/train-nlayer_'+model_name+'.pt'\n",
        "if os.path.exists(ckpt_path):\n",
        "    ckpt = torch.load(ckpt_path)\n",
        "    try:\n",
        "      model.load_state_dict(ckpt['my_model'])\n",
        "      optimizer.load_state_dict(ckpt['optimizer'])\n",
        "      best_acc = ckpt['best_loss']\n",
        "    except RuntimeError as e:\n",
        "        print('wrong checkpoint')\n",
        "    else:    \n",
        "      print('checkpoint is loaded !')\n",
        "      print('current best loss : %.2f' % best_loss)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bjr9R_wDjwHE",
        "cellView": "form"
      },
      "source": [
        "#@title Train the model\n",
        "# Training\n",
        "max_epochs = 100\n",
        "n_iter = 0\n",
        "\n",
        "# each 50 iterations we are going to compare the losses\n",
        "total_train_loss = []\n",
        "total_valid_loss = []\n",
        "\n",
        "for e in range(max_epochs):\n",
        "    model.train()\n",
        "    random.shuffle(train_data)\n",
        "    train_loss = []\n",
        "    for b_train, batch in enumerate(batched_learning(train_data, batch_size=batch_size)):        \n",
        "        \n",
        "        n_iter += 1\n",
        "        # train phase\n",
        "        # feed data into the network and get outputs.\n",
        "        inputs, target = batch\n",
        "        \n",
        "        #print(inputs.shape, target.shape)\n",
        "\n",
        "        # Train on GPU\n",
        "        inputs = (tensorFromSequence(inputs)).to(device)\n",
        "        target = (tensorFromSequence(target)).to(device)\n",
        "\n",
        "        # Create mask for both input and target sequences\n",
        "        input_mask, target_mask = create_masks(inputs, target, pad_token)\n",
        "\n",
        "        inputs = inputs[:,0,:]\n",
        "        target = target[:,0,:] \n",
        "\n",
        "        ys = target.contiguous().view(-1)              \n",
        "\n",
        "        # ys = labels.contiguous().view(-1).to(torch.long).to(device)\n",
        "        \n",
        "        # feed data into the network and get outputs.\n",
        "        preds_idx = model(inputs, target, input_mask, target_mask)\n",
        "        \n",
        "        # Flush out gradients computed at the previous step before computing gradients at the current step. \n",
        "        #       Otherwise, gradients would accumulate.\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        # calculate loss\n",
        "        loss = F.cross_entropy(preds_idx.contiguous().view(preds_idx.size(-1), -1).transpose(0,1), ys, ignore_index = pad_token, size_average = False) / (count_nonpad_tokens(ys))\n",
        "\n",
        "        # accumulates the gradient and backprogate loss.\n",
        "        loss.backward()\n",
        "\n",
        "        # performs a parameter update based on the current gradient\n",
        "        scheduler.step()    \n",
        "        \n",
        "        print('\\n====================================================')\n",
        "        print('Epoch/Batch: {}/{}'.format(e, b_train))\n",
        "        print('Train >>>> Loss: {:6.6}'.format(loss))\n",
        "\n",
        "        train_loss.append(loss.item())\n",
        "\n",
        "    print('\\n**************************************************')\n",
        "    print(\"\\n*** Test *** \")\n",
        "    \n",
        "    # Validation step\n",
        "    model.eval()\n",
        "    valid_loss = []\n",
        "    with torch.no_grad():\n",
        "        pair = valid_data\n",
        "        inputs = tensorFromSequence(pair[0]).to(device)\n",
        "        target = tensorFromSequence(pair[1]).to(device)\n",
        "        \n",
        "        # Create mask for both input and target sequences\n",
        "        input_mask, target_mask = create_masks(inputs, target, pad_token)\n",
        "\n",
        "        inputs = inputs[:,0,:]\n",
        "        target = target[:,0,:] \n",
        "        ys = target.contiguous().view(-1)\n",
        "\n",
        "        preds_validate = model(inputs, target, input_mask, target_mask)\n",
        "        loss = F.cross_entropy(preds_validate.contiguous().view(preds_validate.size(-1), -1).transpose(0,1), ys, \\\n",
        "                                ignore_index = pad_token, size_average = False) / (count_nonpad_tokens(ys))\n",
        "        valid_loss.append(loss.item())\n",
        "\n",
        "    avg_train_loss = np.mean(train_loss)\n",
        "    avg_valid_loss = np.mean(valid_loss)\n",
        "\n",
        "    total_train_loss.append(avg_train_loss)\n",
        "    total_valid_loss.append(avg_valid_loss)\n",
        "\n",
        "    print(\"[Average Train Loss]: {:6.6}\".format(avg_train_loss))\n",
        "    print(\"[Average Testing Loss]: {:6.6}\".format(avg_valid_loss))\n",
        "\n",
        "    # save checkpoint whenever there is improvement in performance\n",
        "    if avg_valid_loss < best_loss:\n",
        "        best_loss = avg_valid_loss\n",
        "        # Note: optimizer also has states ! don't forget to save them as well.\n",
        "        ckpt = {'my_model':model.state_dict(),\n",
        "                'optimizer':optimizer.state_dict(),\n",
        "                'best_loss':best_loss}\n",
        "        torch.save(ckpt, ckpt_path)\n",
        "        print('checkpoint is saved !')\n",
        "\n",
        "    train_writer.add_scalar('loss/train', avg_train_loss, global_step=n_iter)\n",
        "    test_writer.add_scalar('loss/valid', avg_valid_loss, global_step=n_iter)\n",
        "    print('\\n**************************************************')\n",
        "    # torch.cuda.empty_cache()\n",
        "    # torch.save(model.state_dict(), '/gdrive/My Drive/my_data/library/checkpoints/train-{}.pt'.format(e))\n",
        "    ckpt = {'my_model':model.state_dict(),\n",
        "            'optimizer':optimizer.state_dict(),\n",
        "            'best_loss':best_loss}\n",
        "    torch.save(model.state_dict(), ckpt_backup_path)\n",
        "    torch.save(ckpt, ckpt_fullbackup_path)\n",
        "\n",
        "train_writer = SummaryWriter()\n",
        "test_writer = SummaryWriter()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R1PVt1y82Gzy"
      },
      "source": [
        "# Prediction"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jbj3qbuRjare"
      },
      "source": [
        "model.load_state_dict(ckpt['my_model'])\n",
        "optimizer.load_state_dict(ckpt['optimizer'])\n",
        "best_acc = ckpt['best_loss']\n",
        "\n",
        "model.eval()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qK1Z-ijf2Gd6"
      },
      "source": [
        "inputs = torch.randint(0, 1000, (1,10)).to(device)\n",
        "gen_length = 30\n",
        "generated_midi = torch.Tensor()\n",
        "for seq in range(gen_length):\n",
        "    if seq % 20 == 0: \n",
        "        print(seq)\n",
        "    #print(inputs)\n",
        "    logits = F.softmax(model(inputs[:, :-1], inputs[:, 1:]), -1).to(device)\n",
        "\n",
        "    logits = logits[0, :, :]\n",
        "    \n",
        "    one_hot = torch.distributions.OneHotCategorical(probs=logits[:,-1])\n",
        "    res = one_hot.sample().argmax(-1).unsqueeze(-1).to(device)\n",
        "\n",
        "    #res = torch.transpose(res, 0, 1)\n",
        "\n",
        "    inputs = torch.cat((inputs[0], res), dim=-1)\n",
        "    inputs = torch.reshape(inputs, (1, -1))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1qjnWiWn2a4j"
      },
      "source": [
        "from matplotlib import pyplot as pt\n",
        "\n",
        "x = range(0, len(inputs[0]))\n",
        "y = [ints.item() for ints in inputs[0]]\n",
        "\n",
        "p = pt.scatter(x, y, label='generated pitches')\n",
        "\n",
        "pt.xlabel('index')\n",
        "pt.ylabel('midi value (pitch)')\n",
        "pt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vl7nrHJK4tlU"
      },
      "source": [
        "midi_boy = decode_midi(inputs[0])\n",
        "torch.save(midi_boy, gdrive_root + '/my_data/library/test.midi')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CHL3615jlbAb"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}