SuperPiano5_(Reformer_TPU).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "name": "SuperPiano5 (Reformer_TPU).ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "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.7.7"
    }
  },
  "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/SuperPiano5_(Reformer_TPU).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lgn4KUAATtXS"
      },
      "source": [
        "# Super Piano 5 (v.0.5)\r\n",
        "\r\n",
        "Uses integer version of the https://github.com/asigalov61/Tegridy-MIDI-Dataset/blob/master/Tegridy-Piano-CC-BY-NC-SA.zip\r\n",
        "\r\n",
        "Most of the colab as found here:\r\n",
        "\r\n",
        "https://github.com/jlilleberg/Philosophical-Text-Generation-using-Reformers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QINqPZH6FzB3"
      },
      "source": [
        "# Install dependencies\n",
        "\n",
        "!git clone https://github.com/asigalov61/tegridy-tools\n",
        "%cd /content/tegridy-tools/tegridy-tools/\n",
        "import TMIDI\n",
        "%cd /content/\n",
        "\n",
        "!wget https://github.com/asigalov61/SuperPiano/raw/master/MUSIC_INT_EVAL_DATASET.zip\n",
        "\n",
        "!pip install --upgrade -q jax\n",
        "!pip install --upgrade -q jaxlib\n",
        "!pip install --upgrade -q trax==1.3.6\n",
        "!pip install --upgrade -q sentencepiece\n",
        "!pip install --upgrade -q gin \n",
        "\n",
        "# Make sure the Colab Runtime is set to Accelerator: TPU.\n",
        "import requests\n",
        "import os\n",
        "if 'TPU_DRIVER_MODE' not in globals():\n",
        "  url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver0.1-dev20191206'\n",
        "  resp = requests.post(url)\n",
        "  TPU_DRIVER_MODE = 1\n",
        "\n",
        "# The following is required to use TPU Driver as JAX's backend.\n",
        "from jax.config import config\n",
        "config.FLAGS.jax_xla_backend = \"tpu_driver\"\n",
        "config.FLAGS.jax_backend_target = \"grpc://\" + os.environ['COLAB_TPU_ADDR']\n",
        "print(config.FLAGS.jax_backend_target)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s4-aGTZBUxAy"
      },
      "source": [
        "!wget https://github.com/asigalov61/SuperPiano/raw/master/MUSIC_INT_EVAL_DATASET.zip\r\n",
        "!unzip MUSIC_INT_EVAL_DATASET.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pv0PuxC3F-jn"
      },
      "source": [
        "import gin\n",
        "import os\n",
        "import numpy as np\n",
        "from scipy.special import softmax\n",
        "\n",
        "# Zipping and downloading files\n",
        "from google.colab import files\n",
        "import shutil\n",
        "\n",
        "# Trax\n",
        "import jax\n",
        "import trax\n",
        "from trax.data import inputs\n",
        "import jax.numpy as jnp\n",
        "\n",
        "# NLP Vocab Generation\n",
        "import sentencepiece as spm\n",
        "\n",
        "# TensorFlow\n",
        "from tensorflow.compat.v1.io.gfile import GFile\n",
        "import tensorflow as tf"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L7shBqm0T6rz"
      },
      "source": [
        "# This is for the future MIDI Processor code/section.\r\n",
        "# This is to create custom datasets.\r\n",
        "\r\n",
        "# TXT2, STR2 = TMIDI.Tegridy_TXT_to_INT_Converter(TXT1)\r\n",
        "# TXT3 = '\\n'.join([str(elem) for elem in TXT2])\r\n",
        "# with open('/content/INT_DATASET.TXT', 'a') as t:\r\n",
        "#  t.write(TXT3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yRnIoCBkNXyd"
      },
      "source": [
        "# Load the dataset\n",
        "with GFile('/content/MUSIC_INT_EVAL_DATASET.TXT') as f:\n",
        "    text = f.read()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_l9Znw2bS-Pg"
      },
      "source": [
        "# Train a BPE model on the dataset\n",
        "spm.SentencePieceTrainer.train('--input=/content/MUSIC_INT_EVAL_DATASET.TXT \\\n",
        "                                --model_prefix=cp.320 \\\n",
        "                                --vocab_size=320 \\\n",
        "                                --model_type=bpe')\n",
        "# Load BPE vocabulary\n",
        "TOKENIZER = spm.SentencePieceProcessor() \n",
        "TOKENIZER.load('cp.320.model')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OQwhEqJzS-I2"
      },
      "source": [
        "# Tokenize (set to max for the provided dataset)\n",
        "\n",
        "IDS = TOKENIZER.EncodeAsIds(text[:1600000])\n",
        "IDS = np.asarray(IDS, dtype=np.int32)\n",
        "PAD_AMOUNT = 512 * 1024 - len(IDS)\n",
        "print(\"Number of tokens:\", IDS.shape[0])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bdKnx2FyS-AO"
      },
      "source": [
        "# Set up the data pipeline.\n",
        "def my_inputs(n_devices):\n",
        "  while True:\n",
        "    inputs = []\n",
        "    mask = []\n",
        "    pad_amounts = np.random.choice(PAD_AMOUNT, n_devices)\n",
        "    for i in range(n_devices):\n",
        "      inputs.append(np.pad(IDS, (pad_amounts[i], PAD_AMOUNT - pad_amounts[i]), # Pad IDS by different amount for each device\n",
        "                            mode='constant'))\n",
        "      mask.append(np.pad(np.ones_like(IDS, dtype=np.float32),\n",
        "                          (pad_amounts[i], PAD_AMOUNT - pad_amounts[i]),\n",
        "                          mode='constant'))\n",
        "    inputs = np.stack(inputs)\n",
        "    mask = np.stack(mask)\n",
        "    yield (inputs, inputs, mask)\n",
        "\n",
        "print(\"(device count, tokens per device) = \",\n",
        "      next(my_inputs(trax.fastmath.device_count()))[0].shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-NeFGlGGS959"
      },
      "source": [
        "# Configure hyperparameters.\n",
        "gin.parse_config(\"\"\"\n",
        "import trax.layers\n",
        "import trax.models\n",
        "import trax.optimizers\n",
        "import trax.data.inputs\n",
        "import trax.supervised.trainer_lib\n",
        "\n",
        "# Parameters that will vary between experiments:\n",
        "# ==============================================================================\n",
        "train.model = @trax.models.ReformerLM\n",
        "# Model will have 6 layers, alternating between the LSH attention\n",
        "# and local attention within a certain context window.\n",
        "n_layers = 6\n",
        "attn_type = [\n",
        "  @trax.layers.SelfAttention,\n",
        "  @LSHSelfAttention,  \n",
        "  @trax.layers.SelfAttention,\n",
        "  @LSHSelfAttention,\n",
        "  @trax.layers.SelfAttention,\n",
        "  @LSHSelfAttention,\n",
        "  ]\n",
        "share_qk = False  # LSH attention ignores this flag and always shares q & k\n",
        "n_heads = 2\n",
        "attn_kv = 64\n",
        "dropout = 0.05\n",
        "n_tokens = 524288\n",
        "\n",
        "# Parameters for multifactor:\n",
        "# ==============================================================================\n",
        "multifactor.constant = 0.01\n",
        "multifactor.factors = 'constant * linear_warmup * cosine_decay'\n",
        "multifactor.warmup_steps = 100\n",
        "multifactor.steps_per_cycle = 900\n",
        "\n",
        "# Parameters for Adam:\n",
        "# ==============================================================================\n",
        "Adam.weight_decay_rate=0.0\n",
        "Adam.b1 = 0.86\n",
        "Adam.b2 = 0.92\n",
        "Adam.eps = 1e-9\n",
        "\n",
        "# Parameters for SelfAttention:\n",
        "# ==============================================================================\n",
        "trax.layers.SelfAttention.attention_dropout = 0.05\n",
        "trax.layers.SelfAttention.chunk_len = 64\n",
        "trax.layers.SelfAttention.n_chunks_before = 1\n",
        "trax.layers.SelfAttention.n_parallel_heads = 1\n",
        "\n",
        "# Parameters for LSHSelfAttention:\n",
        "# ==============================================================================\n",
        "LSHSelfAttention.attention_dropout = 0.0\n",
        "LSHSelfAttention.chunk_len = 64\n",
        "LSHSelfAttention.n_buckets = [64, 128]\n",
        "LSHSelfAttention.n_chunks_after = 0\n",
        "LSHSelfAttention.n_chunks_before = 1\n",
        "LSHSelfAttention.n_hashes = 1\n",
        "LSHSelfAttention.n_parallel_heads = 1\n",
        "LSHSelfAttention.predict_drop_len = 128\n",
        "LSHSelfAttention.predict_mem_len = 1024\n",
        "\n",
        "# Parameters for ReformerLM:\n",
        "# ==============================================================================\n",
        "ReformerLM.attention_type = %attn_type\n",
        "ReformerLM.d_attention_key = %attn_kv\n",
        "ReformerLM.d_attention_value = %attn_kv\n",
        "ReformerLM.d_model = 256\n",
        "ReformerLM.d_ff = 512\n",
        "ReformerLM.dropout = %dropout\n",
        "ReformerLM.ff_activation = @trax.layers.Relu\n",
        "ReformerLM.max_len = %n_tokens\n",
        "ReformerLM.mode = 'train'\n",
        "ReformerLM.n_heads = %n_heads\n",
        "ReformerLM.n_layers = %n_layers\n",
        "ReformerLM.vocab_size = 320\n",
        "ReformerLM.axial_pos_shape = (512, 1024)\n",
        "ReformerLM.d_axial_pos_embs= (64, 192)\n",
        "\"\"\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X8ay-6VxS9xw"
      },
      "source": [
        "# Trainer.\n",
        "output_dir = os.path.expanduser('model')\n",
        "!rm -f ~/model/model.pkl.gz  # Remove old model\n",
        "\n",
        "trainer = trax.supervised.Trainer(\n",
        "    model=trax.models.ReformerLM,\n",
        "    loss_fn=trax.layers.CrossEntropyLoss(),\n",
        "    optimizer=trax.optimizers.Adam,\n",
        "    lr_schedule=trax.lr.multifactor(),\n",
        "    inputs=trax.data.inputs.Inputs(my_inputs),\n",
        "    output_dir=output_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qf0HITY7Wz8C"
      },
      "source": [
        "# Train Model\n",
        "import tqdm\n",
        "for _ in tqdm.tqdm(range(50)):\n",
        "  trainer.train_epoch(n_steps=100, n_eval_steps=20)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JEikLctyte-m"
      },
      "source": [
        "# Zip directory contents\n",
        "shutil.make_archive(\"project\", \"zip\", \".\")\n",
        "\n",
        "# Download zipped directory\n",
        "files.download('project.zip')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FLebeiAFrItK"
      },
      "source": [
        "# In the Reformer paper, increasing the number of hashing rounds helps with quality. \n",
        "# The number of hashing rounds at can be increased at evaluation time only.\n",
        "gin.parse_config(\"\"\"LSHSelfAttention.n_hashes = 8\"\"\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "thdiRvGDrJhy"
      },
      "source": [
        "# Load the trained Reformer in 'predict' mode\n",
        "model = trax.models.ReformerLM(mode='predict')\n",
        "output_dir = os.path.expanduser('model')\n",
        "model.init_from_file(os.path.join(output_dir,'model.pkl.gz'),\n",
        "                     weights_only=True)\n",
        "\n",
        "# Sample from ReformerLM\n",
        "output_token_ids = trax.supervised.decoding.autoregressive_sample(\n",
        "    model, temperature=0.8, max_length=1024, accelerate=True)\n",
        "\n",
        "# Decode token IDs\n",
        "# Reformer outputed a batch with one item so access it using [0]\n",
        "# tolist() converts from int64 to int, the type SentencePiece expects\n",
        "input = TOKENIZER.DecodeIds(output_token_ids[0].tolist())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ox03iX3KHCKa"
      },
      "source": [
        "# Run the cells below to convert generated output to MIDI.\r\n",
        "# If you getting errors/halts, regenerate the output again.\r\n",
        "# Model must be sufficiently trained. Rec. 0.90+ accuracy for the output to make sense and pass error control.\r\n",
        "\r\n",
        "TXT = TMIDI.Tegridy_INT_String_to_TXT_Converter(input, line_by_line_input=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WggjV2NNHCti"
      },
      "source": [
        "SONG = TMIDI.Tegridy_Reduced_TXT_to_Notes_Converter(TXT, has_MIDI_channels=False, has_velocities=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z0MvzNCOHEXZ"
      },
      "source": [
        "stats = TMIDI.Tegridy_SONG_to_MIDI_Converter(SONG=SONG[0], output_file_name='/content/OUTPUT_MIDI')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}