# -*- coding: utf-8 -*-
"""SuperPiano5 (Reformer_TPU).ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1MENA73Z5KnMvi-mWPZCSUMnHS-YWx8Dp

# Super Piano 5 (v.0.5)

Uses integer version of the https://github.com/asigalov61/Tegridy-MIDI-Dataset/blob/master/Tegridy-Piano-CC-BY-NC-SA.zip

Most of the colab as found here:

https://github.com/jlilleberg/Philosophical-Text-Generation-using-Reformers
"""

# Commented out IPython magic to ensure Python compatibility.
# Install dependencies

!git clone https://github.com/asigalov61/tegridy-tools
# %cd /content/tegridy-tools/tegridy-tools/
import TMIDI
# %cd /content/

!wget https://github.com/asigalov61/SuperPiano/raw/master/MUSIC_INT_EVAL_DATASET.zip

!pip install --upgrade -q jax
!pip install --upgrade -q jaxlib
!pip install --upgrade -q trax==1.3.6
!pip install --upgrade -q sentencepiece
!pip install --upgrade -q gin 

# Make sure the Colab Runtime is set to Accelerator: TPU.
import requests
import os
if 'TPU_DRIVER_MODE' not in globals():
  url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver0.1-dev20191206'
  resp = requests.post(url)
  TPU_DRIVER_MODE = 1

# The following is required to use TPU Driver as JAX's backend.
from jax.config import config
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://" + os.environ['COLAB_TPU_ADDR']
print(config.FLAGS.jax_backend_target)

!wget https://github.com/asigalov61/SuperPiano/raw/master/MUSIC_INT_EVAL_DATASET.zip
!unzip MUSIC_INT_EVAL_DATASET.zip

import gin
import os
import numpy as np
from scipy.special import softmax

# Zipping and downloading files
from google.colab import files
import shutil

# Trax
import jax
import trax
from trax.data import inputs
import jax.numpy as jnp

# NLP Vocab Generation
import sentencepiece as spm

# TensorFlow
from tensorflow.compat.v1.io.gfile import GFile
import tensorflow as tf

# This is for the future MIDI Processor code/section.
# This is to create custom datasets.

# TXT2, STR2 = TMIDI.Tegridy_TXT_to_INT_Converter(TXT1)
# TXT3 = '\n'.join([str(elem) for elem in TXT2])
# with open('/content/INT_DATASET.TXT', 'a') as t:
#  t.write(TXT3)

# Load the dataset
with GFile('/content/MUSIC_INT_EVAL_DATASET.TXT') as f:
    text = f.read()

# Train a BPE model on the dataset
spm.SentencePieceTrainer.train('--input=/content/MUSIC_INT_EVAL_DATASET.TXT \
                                --model_prefix=cp.320 \
                                --vocab_size=320 \
                                --model_type=bpe')
# Load BPE vocabulary
TOKENIZER = spm.SentencePieceProcessor() 
TOKENIZER.load('cp.320.model')

# Tokenize (set to max for the provided dataset)

IDS = TOKENIZER.EncodeAsIds(text[:1600000])
IDS = np.asarray(IDS, dtype=np.int32)
PAD_AMOUNT = 512 * 1024 - len(IDS)
print("Number of tokens:", IDS.shape[0])

# Set up the data pipeline.
def my_inputs(n_devices):
  while True:
    inputs = []
    mask = []
    pad_amounts = np.random.choice(PAD_AMOUNT, n_devices)
    for i in range(n_devices):
      inputs.append(np.pad(IDS, (pad_amounts[i], PAD_AMOUNT - pad_amounts[i]), # Pad IDS by different amount for each device
                            mode='constant'))
      mask.append(np.pad(np.ones_like(IDS, dtype=np.float32),
                          (pad_amounts[i], PAD_AMOUNT - pad_amounts[i]),
                          mode='constant'))
    inputs = np.stack(inputs)
    mask = np.stack(mask)
    yield (inputs, inputs, mask)

print("(device count, tokens per device) = ",
      next(my_inputs(trax.fastmath.device_count()))[0].shape)

# Configure hyperparameters.
gin.parse_config("""
import trax.layers
import trax.models
import trax.optimizers
import trax.data.inputs
import trax.supervised.trainer_lib

# Parameters that will vary between experiments:
# ==============================================================================
train.model = @trax.models.ReformerLM
# Model will have 6 layers, alternating between the LSH attention
# and local attention within a certain context window.
n_layers = 6
attn_type = [
  @trax.layers.SelfAttention,
  @LSHSelfAttention,  
  @trax.layers.SelfAttention,
  @LSHSelfAttention,
  @trax.layers.SelfAttention,
  @LSHSelfAttention,
  ]
share_qk = False  # LSH attention ignores this flag and always shares q & k
n_heads = 2
attn_kv = 64
dropout = 0.05
n_tokens = 524288

# Parameters for multifactor:
# ==============================================================================
multifactor.constant = 0.01
multifactor.factors = 'constant * linear_warmup * cosine_decay'
multifactor.warmup_steps = 100
multifactor.steps_per_cycle = 900

# Parameters for Adam:
# ==============================================================================
Adam.weight_decay_rate=0.0
Adam.b1 = 0.86
Adam.b2 = 0.92
Adam.eps = 1e-9

# Parameters for SelfAttention:
# ==============================================================================
trax.layers.SelfAttention.attention_dropout = 0.05
trax.layers.SelfAttention.chunk_len = 64
trax.layers.SelfAttention.n_chunks_before = 1
trax.layers.SelfAttention.n_parallel_heads = 1

# Parameters for LSHSelfAttention:
# ==============================================================================
LSHSelfAttention.attention_dropout = 0.0
LSHSelfAttention.chunk_len = 64
LSHSelfAttention.n_buckets = [64, 128]
LSHSelfAttention.n_chunks_after = 0
LSHSelfAttention.n_chunks_before = 1
LSHSelfAttention.n_hashes = 1
LSHSelfAttention.n_parallel_heads = 1
LSHSelfAttention.predict_drop_len = 128
LSHSelfAttention.predict_mem_len = 1024

# Parameters for ReformerLM:
# ==============================================================================
ReformerLM.attention_type = %attn_type
ReformerLM.d_attention_key = %attn_kv
ReformerLM.d_attention_value = %attn_kv
ReformerLM.d_model = 256
ReformerLM.d_ff = 512
ReformerLM.dropout = %dropout
ReformerLM.ff_activation = @trax.layers.Relu
ReformerLM.max_len = %n_tokens
ReformerLM.mode = 'train'
ReformerLM.n_heads = %n_heads
ReformerLM.n_layers = %n_layers
ReformerLM.vocab_size = 320
ReformerLM.axial_pos_shape = (512, 1024)
ReformerLM.d_axial_pos_embs= (64, 192)
""")

# Trainer.
output_dir = os.path.expanduser('model')
!rm -f ~/model/model.pkl.gz  # Remove old model

trainer = trax.supervised.Trainer(
    model=trax.models.ReformerLM,
    loss_fn=trax.layers.CrossEntropyLoss(),
    optimizer=trax.optimizers.Adam,
    lr_schedule=trax.lr.multifactor(),
    inputs=trax.data.inputs.Inputs(my_inputs),
    output_dir=output_dir)

# Train Model
import tqdm
for _ in tqdm.tqdm(range(50)):
  trainer.train_epoch(n_steps=100, n_eval_steps=20)

# Zip directory contents
shutil.make_archive("project", "zip", ".")

# Download zipped directory
files.download('project.zip')

# In the Reformer paper, increasing the number of hashing rounds helps with quality. 
# The number of hashing rounds at can be increased at evaluation time only.
gin.parse_config("""LSHSelfAttention.n_hashes = 8""")

# Load the trained Reformer in 'predict' mode
model = trax.models.ReformerLM(mode='predict')
output_dir = os.path.expanduser('model')
model.init_from_file(os.path.join(output_dir,'model.pkl.gz'),
                     weights_only=True)

# Sample from ReformerLM
output_token_ids = trax.supervised.decoding.autoregressive_sample(
    model, temperature=0.8, max_length=1024, accelerate=True)

# Decode token IDs
# Reformer outputed a batch with one item so access it using [0]
# tolist() converts from int64 to int, the type SentencePiece expects
input = TOKENIZER.DecodeIds(output_token_ids[0].tolist())

# Run the cells below to convert generated output to MIDI.
# If you getting errors/halts, regenerate the output again.
# Model must be sufficiently trained. Rec. 0.90+ accuracy for the output to make sense and pass error control.

TXT = TMIDI.Tegridy_INT_String_to_TXT_Converter(input, line_by_line_input=False)

SONG = TMIDI.Tegridy_Reduced_TXT_to_Notes_Converter(TXT, has_MIDI_channels=False, has_velocities=False)

stats = TMIDI.Tegridy_SONG_to_MIDI_Converter(SONG=SONG[0], output_file_name='/content/OUTPUT_MIDI')