NanoMaestro Full model weights released

app.py

@@ -0,0 +1,318 @@
+import gradio as gr
+import torch
+import os
+import random
+import glob
+import numpy as np
+import pretty_midi
+import scipy.io.wavfile
+
+# --- Dependencies from your project ---
+# (Make sure these files are in the same directory)
+try:
+    from model.music_transformer import MusicTransformer
+    from processor import encode_midi, decode_midi
+    from dataset.e_piano import process_midi
+    from utilities.constants import *
+    from utilities.device import get_device, use_cuda
+except ImportError as e:
+    print(f"Error: Could not import necessary files.")
+    print(f"Make sure app.py is in the same folder as 'model', 'processor.py', etc.")
+    print(f"Details: {e}")
+    exit()
+
+# --- Your Model's Hyperparameters ---
+# (Pulled from your training logs)
+MODEL_CONFIG = {
+    "n_layers": 6,
+    "num_heads": 8,
+    "d_model": 512,
+    "dim_feedforward": 1024,
+    "max_sequence": 2048,
+    "rpr": True
+}
+# ------------------------------------
+
+# Global variable to hold the loaded model
+model = None
+device = get_device()
+print(f"Using device: {device}")
+
+
+def load_model(model_path):
+    """
+    Loads the trained MusicTransformer model into memory.
+    """
+    global model
+    if model_path is None or not os.path.exists(model_path):
+        return "Error: Model file not found. Please check the path."
+
+    try:
+        print("Loading model...")
+        model = MusicTransformer(
+            n_layers=MODEL_CONFIG["n_layers"],
+            num_heads=MODEL_CONFIG["num_heads"],
+            d_model=MODEL_CONFIG["d_model"],
+            dim_feedforward=MODEL_CONFIG["dim_feedforward"],
+            max_sequence=MODEL_CONFIG["max_sequence"],
+            rpr=MODEL_CONFIG["rpr"]
+        ).to(device)
+
+        # Load the weights, mapping to the correct device
+        model.load_state_dict(
+            torch.load(model_path, map_location=device, weights_only=True)
+        )
+        model.eval()
+        print("Model loaded successfully.")
+        return f"Model '{model_path}' loaded successfully."
+    except Exception as e:
+        return f"Error loading model: {e}"
+
+
+# --- NEW FUNCTION ---
+def midi_to_wav(midi_file_path, wav_file_path):
+    """
+    Synthesizes a MIDI file to a WAV file using pretty_midi's
+    built-in (simple) sine wave synthesizer.
+    """
+    try:
+        pm = pretty_midi.PrettyMIDI(midi_file_path)
+        # Synthesize the audio at a 44.1kHz sample rate
+        audio_data = pm.synthesize(fs=44100)
+        # Write as a 16-bit WAV file
+        scipy.io.wavfile.write(wav_file_path, 44100, audio_data.astype(np.int16))
+        return wav_file_path
+    except Exception as e:
+        print(f"Error during MIDI to WAV conversion: {e}")
+        return None
+
+
+# --- END NEW FUNCTION ---
+
+def generate_music(primer_type, uploaded_midi, upload_start_location, maestro_path, maestro_start_location,
+                   primer_length, generation_length_new, progress=gr.Progress(track_tqdm=True)):
+    """
+    The main function called by the Gradio button.
+    """
+    global model
+    if model is None:
+        # --- MODIFICATION: Return 3 values on error ---
+        yield "Error: Model is not loaded. Please load a model first.", None, None
+
+    try:
+        # --- 1. Prepare the Primer ---
+        primer = None
+        num_primer = 0
+
+        total_target_length = primer_length + generation_length_new
+        if total_target_length > MODEL_CONFIG["max_sequence"]:
+            total_target_length = MODEL_CONFIG["max_sequence"]
+            yield f"Warning: Clamping to {total_target_length} tokens.", None, None
+
+        if primer_type == "Generate from Silence":
+            yield "Generating from silence...", None, None
+            primer = torch.tensor([372], dtype=TORCH_LABEL_TYPE, device=device)
+            num_primer = 1
+
+        elif primer_type == "Random Maestro MIDI":
+            yield "Finding random Maestro file...", None, None
+            if maestro_path is None or not os.path.isdir(maestro_path):
+                yield f"Error: Maestro path '{maestro_path}' is not valid.", None, None
+                return
+
+            midi_files = glob.glob(os.path.join(maestro_path, "**", "*.mid"), recursive=True) + \
+                         glob.glob(os.path.join(maestro_path, "**", "*.midi"), recursive=True)
+
+            if not midi_files:
+                yield f"Error: No .mid/.midi files found in '{maestro_path}'.", None, None
+                return
+
+            random_file = random.choice(midi_files)
+            yield f"Tokenizing random file: {os.path.basename(random_file)}...", None, None
+            raw_mid = encode_midi(random_file)
+
+            is_random_start = (maestro_start_location == "Random Location")
+            primer_tokens, _ = process_midi(raw_mid, primer_length, random_seq=is_random_start)
+
+            primer = torch.tensor(primer_tokens, dtype=TORCH_LABEL_TYPE, device=device)
+            num_primer = primer.shape[0]
+
+        elif primer_type == "Upload My Own MIDI":
+            if uploaded_midi is None:
+                yield "Error: Please upload a MIDI file.", None, None
+                return
+
+            yield f"Tokenizing uploaded MIDI: {os.path.basename(uploaded_midi.name)}...", None, None
+            raw_mid = encode_midi(uploaded_midi.name)
+            if not raw_mid:
+                yield "Error: Could not read MIDI messages.", None, None
+                return
+
+            is_random_start = (upload_start_location == "Random Location")
+            primer_tokens, _ = process_midi(raw_mid, primer_length, random_seq=is_random_start)
+            primer = torch.tensor(primer_tokens, dtype=TORCH_LABEL_TYPE, device=device)
+            num_primer = primer.shape[0]
+
+        if num_primer == 0:
+            yield "Error: Primer processing resulted in 0 tokens.", None, None
+            return
+
+        # --- 2. Run Generation ---
+        yield f"Primed with {num_primer} tokens. Generating {generation_length_new} new tokens...", None, None
+
+        primer_batch = primer.unsqueeze(0)
+
+        model.eval()
+        with torch.set_grad_enabled(False):
+            rand_seq = model.generate(primer_batch, total_target_length, beam=0)
+
+        # --- 3. Process and Save Output ---
+        generated_only_tokens = rand_seq[0][num_primer:]
+
+        if len(generated_only_tokens) == 0:
+            yield "Warning: Generation produced 0 new tokens.", None, None
+            return
+
+        # --- MODIFICATION: Define output paths ---
+        midi_output_filename = "generation_output.mid"
+        wav_output_filename = "generation_output.wav"
+
+        # Save the MIDI file
+        decode_midi(generated_only_tokens.cpu().numpy(), midi_output_filename)
+
+        # --- MODIFICATION: Synthesize MIDI to WAV ---
+        yield "Synthesizing audio...", midi_output_filename, None
+        wav_path = midi_to_wav(midi_output_filename, wav_output_filename)
+
+        if wav_path:
+            yield "Generation Complete!", midi_output_filename, wav_path
+        else:
+            yield "Generation complete (WAV synthesis failed).", midi_output_filename, None
+
+    except Exception as e:
+        yield f"An error occurred: {e}", None, None
+
+
+# --- Build the Gradio UI ---
+with gr.Blocks(theme=gr.themes.Soft()) as app:
+    gr.Markdown("# 🎹 Music Transformer Generation UI")
+    gr.Markdown("Load your trained model and generate music from silence, a random seed, or your own MIDI file.")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 1. Load Model")
+            model_path_input = gr.Textbox(
+                label="Path to your .pickle model file",
+                value="best_acc_weights.pickle"
+            )
+            load_button = gr.Button("Load Model", variant="primary")
+            load_status = gr.Textbox(label="Model Status", interactive=False)
+
+        with gr.Column(scale=2):
+            gr.Markdown("### 2. Configure Generation")
+
+            primer_type_input = gr.Radio(
+                label="Choose Primer Type",
+                choices=["Generate from Silence", "Random Maestro MIDI", "Upload My Own MIDI"],
+                value="Generate from Silence"
+            )
+
+            with gr.Column(visible=False) as maestro_options:
+                maestro_path_input = gr.Textbox(
+                    label="Path to RAW Maestro MIDI Folder (searches all subfolders)",
+                    value="./maestro-v2.0.0"
+                )
+                maestro_start_location_input = gr.Radio(
+                    label="Primer Start Location",
+                    choices=["Start of File", "Random Location"],
+                    value="Random Location",
+                    info="Selects a random chunk from the file, giving more variety."
+                )
+
+            with gr.Column(visible=False) as upload_options:
+                uploaded_midi_input = gr.File(
+                    label="Upload Your MIDI Primer",
+                    file_types=[".mid", ".midi"]
+                )
+                upload_start_location_input = gr.Radio(
+                    label="Primer Start Location",
+                    choices=["Start of File", "Random Location"],
+                    value="Start of File"
+                )
+
+            primer_length_slider = gr.Slider(
+                label="Primer Length (Tokens)",
+                minimum=64,
+                maximum=2000,
+                value=512,
+                step=32,
+                info="How many tokens to use from the primer file. Ignored for 'Silence'."
+            )
+
+            generation_length_slider = gr.Slider(
+                label="New Tokens to Generate",
+                minimum=128,
+                maximum=2048,
+                value=1024,
+                step=32,
+                info="How many new tokens to create after the primer."
+            )
+
+            generate_button = gr.Button("Generate Music", variant="primary")
+
+    with gr.Row():
+        gr.Markdown("### 3. Get Your Music")
+        status_output = gr.Textbox(label="Status", interactive=False)
+        with gr.Row():
+            output_midi_file = gr.File(label="Download Generated MIDI")
+            # --- MODIFICATION: Added Audio player ---
+            output_wav_file = gr.Audio(label="Listen to Generated WAV", type="filepath")
+            # --- END MODIFICATION ---
+
+
+    # --- UI Event Listeners ---
+
+    def update_ui(primer_type):
+        return {
+            maestro_options: gr.Column(visible=(primer_type == "Random Maestro MIDI")),
+            upload_options: gr.Column(visible=(primer_type == "Upload My Own MIDI")),
+            primer_length_slider: gr.Slider(visible=(primer_type != "Generate from Silence"))
+        }
+
+
+    primer_type_input.change(
+        fn=update_ui,
+        inputs=primer_type_input,
+        outputs=[maestro_options, upload_options, primer_length_slider]
+    )
+
+    load_button.click(
+        fn=load_model,
+        inputs=model_path_input,
+        outputs=load_status
+    )
+
+    # --- MODIFICATION: Updated outputs list ---
+    generate_button.click(
+        fn=generate_music,
+        inputs=[
+            primer_type_input,
+            uploaded_midi_input,
+            upload_start_location_input,
+            maestro_path_input,
+            maestro_start_location_input,
+            primer_length_slider,
+            generation_length_slider
+        ],
+        outputs=[status_output, output_midi_file, output_wav_file]  # <-- Added WAV output
+    )
+    # --- END MODIFICATION ---
+
+if __name__ == "__main__":
+    # Check if CUDA is available and set device
+    if (not torch.cuda.is_available()):
+        print("----- WARNING: CUDA devices not detected. This will cause the model to run very slow! -----")
+        use_cuda(False)
+
+    print("Launching Gradio UI...")
+    app.launch()

best_acc_weights.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bba9e1c91f449753383895c224383dd0ab8402ed003d0d6eb83a4d3f9a3de5df
+size 59442741

dataset/e_piano.py

@@ -0,0 +1,169 @@
+import os
+import pickle
+import random
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset
+
+from utilities.constants import *
+from utilities.device import cpu_device
+
+SEQUENCE_START = 0
+
+# EPianoDataset
+class EPianoDataset(Dataset):
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Pytorch Dataset for the Maestro e-piano dataset (https://magenta.tensorflow.org/datasets/maestro).
+    Recommended to use with Dataloader (https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader)
+
+    Uses all files found in the given root directory of pre-processed (preprocess_midi.py)
+    Maestro midi files.
+    ----------
+    """
+
+    def __init__(self, root, max_seq=2048, random_seq=True):
+        self.root       = root
+        self.max_seq    = max_seq
+        self.random_seq = random_seq
+
+        fs = [os.path.join(root, f) for f in os.listdir(self.root)]
+        self.data_files = [f for f in fs if os.path.isfile(f)]
+
+    # __len__
+    def __len__(self):
+        """
+        ----------
+        Author: Damon Gwinn
+        ----------
+        How many data files exist in the given directory
+        ----------
+        """
+
+        return len(self.data_files)
+
+    # __getitem__
+    def __getitem__(self, idx):
+        """
+        ----------
+        Author: Damon Gwinn
+        ----------
+        Gets the indexed midi batch. Gets random sequence or from start depending on random_seq.
+
+        Returns the input and the target.
+        ----------
+        """
+
+        # All data on cpu to allow for the Dataloader to multithread
+        i_stream    = open(self.data_files[idx], "rb")
+        # return pickle.load(i_stream), None
+        raw_mid     = torch.tensor(pickle.load(i_stream), dtype=TORCH_LABEL_TYPE, device=cpu_device())
+        i_stream.close()
+
+        x, tgt = process_midi(raw_mid, self.max_seq, self.random_seq)
+
+        return x, tgt
+
+# process_midi
+def process_midi(raw_mid, max_seq, random_seq):
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Takes in pre-processed raw midi and returns the input and target. Can use a random sequence or
+    go from the start based on random_seq.
+    ----------
+    """
+
+    x   = torch.full((max_seq, ), TOKEN_PAD, dtype=TORCH_LABEL_TYPE, device=cpu_device())
+    tgt = torch.full((max_seq, ), TOKEN_PAD, dtype=TORCH_LABEL_TYPE, device=cpu_device())
+
+    raw_len     = len(raw_mid)
+    full_seq    = max_seq + 1 # Performing seq2seq
+
+    if(raw_len == 0):
+        return x, tgt
+
+    if(raw_len < full_seq):
+        x[:raw_len]         = raw_mid
+        tgt[:raw_len-1]     = raw_mid[1:]
+        tgt[raw_len]        = TOKEN_END
+    else:
+        # Randomly selecting a range
+        if(random_seq):
+            end_range = raw_len - full_seq
+            start = random.randint(SEQUENCE_START, end_range)
+
+        # Always taking from the start to as far as we can
+        else:
+            start = SEQUENCE_START
+
+        end = start + full_seq
+
+        data = raw_mid[start:end]
+
+        x = data[:max_seq]
+        tgt = data[1:full_seq]
+
+
+    # print("x:",x)
+    # print("tgt:",tgt)
+
+    return x, tgt
+
+
+# create_epiano_datasets
+def create_epiano_datasets(dataset_root, max_seq, random_seq=True):
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Creates train, evaluation, and test EPianoDataset objects for a pre-processed (preprocess_midi.py)
+    root containing train, val, and test folders.
+    ----------
+    """
+
+    train_root = os.path.join(dataset_root, "train")
+    val_root = os.path.join(dataset_root, "val")
+    test_root = os.path.join(dataset_root, "test")
+
+    train_dataset = EPianoDataset(train_root, max_seq, random_seq)
+    val_dataset = EPianoDataset(val_root, max_seq, random_seq)
+    test_dataset = EPianoDataset(test_root, max_seq, random_seq)
+
+    return train_dataset, val_dataset, test_dataset
+
+# compute_epiano_accuracy
+def compute_epiano_accuracy(out, tgt):
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Computes the average accuracy for the given input and output batches. Accuracy uses softmax
+    of the output.
+    ----------
+    """
+
+    softmax = nn.Softmax(dim=-1)
+    out = torch.argmax(softmax(out), dim=-1)
+
+    out = out.flatten()
+    tgt = tgt.flatten()
+
+    mask = (tgt != TOKEN_PAD)
+
+    out = out[mask]
+    tgt = tgt[mask]
+
+    # Empty
+    if(len(tgt) == 0):
+        return 1.0
+
+    num_right = (out == tgt)
+    num_right = torch.sum(num_right).type(TORCH_FLOAT)
+
+    acc = num_right / len(tgt)
+
+    return acc

generate.py

@@ -0,0 +1,134 @@
+import torch
+import torch.nn as nn
+import os
+import random
+import pretty_midi
+import processor
+
+from processor import encode_midi, decode_midi
+
+from utilities.argument_funcs import parse_generate_args, print_generate_args
+from model.music_transformer import MusicTransformer
+from dataset.e_piano import create_epiano_datasets, compute_epiano_accuracy, process_midi
+from torch.utils.data import DataLoader
+from torch.optim import Adam
+
+from utilities.constants import *
+from utilities.device import get_device, use_cuda
+
+
+# main
+def main():
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Entry point. Generates music from a model specified by command line arguments
+    ----------
+    """
+
+    args = parse_generate_args()
+    print_generate_args(args)
+
+    if (args.force_cpu):
+        use_cuda(False)
+        print("WARNING: Forced CPU usage, expect model to perform slower")
+        print("")
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # --- MODIFIED LOGIC ---
+    # Can be None, an integer index to dataset, or a file path
+    if (args.primer_file is None):
+        # --- Load dataset ONLY if no primer file is given ---
+        print("No primer file provided, loading dataset to pick a random primer...")
+        _, _, dataset = create_epiano_datasets(args.midi_root, args.num_prime, random_seq=False)
+        # --- --- --- --- --- --- --- --- --- --- --- --- --- ---
+
+        f = str(random.randrange(len(dataset)))
+        idx = int(f)
+        primer, _ = dataset[idx]
+        primer = primer.to(get_device())
+        num_primer = primer.shape[0]
+        print("Using primer index:", idx, "(", dataset.data_files[idx], ")")
+
+    else:
+        # --- Primer file is provided, NO DATASET NEEDED (unless it's an index) ---
+        f = args.primer_file
+
+        # --- NEW: Check for "silence" ---
+        if (f.lower() == "silence"):
+            print("Generating from silence...")
+            # Create a primer with one token: a medium velocity (64 // 4 = 16)
+            # Velocity START_IDX = 356. Token = 356 + 16 = 372
+            primer = torch.tensor([372], dtype=TORCH_LABEL_TYPE, device=get_device())
+            num_primer = primer.shape[0]  # This will be 1
+
+        # This part handles if the primer is an integer index (e.g., "3")
+        elif (f.isdigit()):
+            print("Primer file is an index, loading dataset...")
+            # --- Load dataset ONLY if primer is an index ---
+            _, _, dataset = create_epiano_datasets(args.midi_root, args.num_prime, random_seq=False)
+            # --- --- --- --- --- --- --- --- --- --- --- ---
+            idx = int(f)
+            primer, _ = dataset[idx]
+            primer = primer.to(get_device())
+            num_primer = primer.shape[0]
+            print("Using primer index:", idx, "(", dataset.data_files[idx], ")")
+
+        # This part handles if the primer is a MIDI file path (e.t., "my_primer.mid")
+        else:
+            print("Primer file is a MIDI path. Loading and tokenizing...")
+            raw_mid = encode_midi(f)
+            if (len(raw_mid) == 0):
+                print("Error: No midi messages in primer file:", f)
+                return
+
+            primer, _ = process_midi(raw_mid, args.num_prime, random_seq=False)
+            primer = torch.tensor(primer, dtype=TORCH_LABEL_TYPE, device=get_device())
+            num_primer = primer.shape[0]  # Get the actual primer length
+            print("Using primer file:", f)
+
+    # --- END MODIFIED LOGIC ---
+
+    model = MusicTransformer(n_layers=args.n_layers, num_heads=args.num_heads,
+                             d_model=args.d_model, dim_feedforward=args.dim_feedforward,
+                             max_sequence=args.max_sequence, rpr=args.rpr).to(get_device())
+
+    model.load_state_dict(torch.load(args.model_weights))
+
+    # --- MODIFICATION: Don't save a primer if we started from silence ---
+    if (args.primer_file.lower() != "silence"):
+        f_path = os.path.join(args.output_dir, "primer.mid")
+        decode_midi(primer.cpu().numpy(), file_path=f_path)
+    # --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- --- ---
+
+    # GENERATION
+    model.eval()
+    with torch.set_grad_enabled(False):
+        if (args.beam > 0):
+            print("BEAM:", args.beam)
+
+            beam_seq = model.generate(primer.unsqueeze(0), args.target_seq_length, beam=args.beam)
+
+            # --- MODIFICATION: Slice the primer ---
+            generated_only = beam_seq[0][num_primer:]
+            # ------------------------------------
+
+            f_path = os.path.join(args.output_dir, "beam.mid")
+            decode_midi(generated_only.cpu().numpy(), file_path=f_path)
+        else:
+            print("RAND DIST")
+
+            rand_seq = model.generate(primer.unsqueeze(0), args.target_seq_length, beam=0)
+
+            # --- MODIFICATION: Slice the primer ---
+            generated_only = rand_seq[0][num_primer:]
+            # ------------------------------------
+
+            f_path = os.path.join(args.output_dir, "rand.mid")
+            decode_midi(generated_only.cpu().numpy(), file_path=f_path)
+
+
+if __name__ == "__main__":
+    main()

model/music_transformer.py

@@ -0,0 +1,135 @@
+import torch
+import torch.nn as nn
+from torch.nn.modules.normalization import LayerNorm
+import random
+
+from utilities.constants import *
+from utilities.device import get_device
+
+from .positional_encoding import PositionalEncoding
+from .rpr import TransformerEncoderRPR, TransformerEncoderLayerRPR
+
+
+# MusicTransformer
+class MusicTransformer(nn.Module):
+    def __init__(self, n_layers=6, num_heads=8, d_model=512, dim_feedforward=1024,
+                 dropout=0.1, max_sequence=2048, rpr=False):
+        super(MusicTransformer, self).__init__()
+
+        self.dummy      = DummyDecoder()
+
+        self.nlayers    = n_layers
+        self.nhead      = num_heads
+        self.d_model    = d_model
+        self.d_ff       = dim_feedforward
+        self.dropout    = dropout
+        self.max_seq    = max_sequence
+        self.rpr        = rpr
+
+        # Input embedding
+        self.embedding = nn.Embedding(VOCAB_SIZE, self.d_model)
+
+        # Positional encoding
+        self.positional_encoding = PositionalEncoding(self.d_model, self.dropout, self.max_seq)
+
+        # Base transformer
+        if(not self.rpr):
+            self.transformer = nn.Transformer(
+                d_model=self.d_model, nhead=self.nhead, num_encoder_layers=self.nlayers,
+                num_decoder_layers=0, dropout=self.dropout,
+                dim_feedforward=self.d_ff, custom_decoder=self.dummy
+            )
+        else:
+            encoder_norm = LayerNorm(self.d_model)
+            encoder_layer = TransformerEncoderLayerRPR(self.d_model, self.nhead, self.d_ff, self.dropout, er_len=self.max_seq)
+            encoder = TransformerEncoderRPR(encoder_layer, self.nlayers, encoder_norm)
+            self.transformer = nn.Transformer(
+                d_model=self.d_model, nhead=self.nhead, num_encoder_layers=self.nlayers,
+                num_decoder_layers=0, dropout=self.dropout,
+                dim_feedforward=self.d_ff, custom_decoder=self.dummy, custom_encoder=encoder
+            )
+
+        # Final output is a softmaxed linear layer
+        self.Wout       = nn.Linear(self.d_model, VOCAB_SIZE)
+        self.softmax    = nn.Softmax(dim=-1)
+
+    # forward
+    def forward(self, x, mask=True):
+        # --- FIX: USE DEVICE OF INPUT TENSOR x ---
+        if(mask is True):
+            # Generate mask on the same device as input x
+            mask = self.transformer.generate_square_subsequent_mask(x.shape[1]).to(x.device)
+        else:
+            mask = None
+        # -----------------------------------------
+
+        x = self.embedding(x)
+
+        # Input shape is (max_seq, batch_size, d_model)
+        x = x.permute(1,0,2)
+
+        x = self.positional_encoding(x)
+
+        # Since there are no true decoder layers, the tgt is unused
+        x_out = self.transformer(src=x, tgt=x, src_mask=mask)
+
+        # Back to (batch_size, max_seq, d_model)
+        x_out = x_out.permute(1,0,2)
+
+        y = self.Wout(x_out)
+        return y
+
+    # generate
+    def generate(self, primer=None, target_seq_length=1024, beam=0, beam_chance=1.0):
+        assert (not self.training), "Cannot generate while in training mode"
+
+        print("Generating sequence of max length:", target_seq_length)
+
+        batch_size = primer.shape[0]
+        gen_seq = torch.full((batch_size, target_seq_length), TOKEN_PAD, dtype=TORCH_LABEL_TYPE, device=get_device())
+
+        num_primer = primer.shape[1]
+        gen_seq[..., :num_primer] = primer.type(TORCH_LABEL_TYPE).to(get_device())
+
+        cur_i = num_primer
+        while(cur_i < target_seq_length):
+            y = self.softmax(self.forward(gen_seq[..., :cur_i]))[..., :TOKEN_END]
+            token_probs = y[:, cur_i-1, :]
+
+            if(beam == 0):
+                beam_ran = 2.0
+            else:
+                beam_ran = random.uniform(0,1)
+
+            if(beam_ran <= beam_chance):
+                token_probs = token_probs.flatten()
+                top_res, top_i = torch.topk(token_probs, beam)
+
+                beam_rows = top_i // VOCAB_SIZE
+                beam_cols = top_i % VOCAB_SIZE
+
+                gen_seq = gen_seq[beam_rows, :]
+                gen_seq[..., cur_i] = beam_cols
+
+            else:
+                distrib = torch.distributions.categorical.Categorical(probs=token_probs)
+                next_token = distrib.sample()
+                gen_seq[:, cur_i] = next_token
+
+                if(next_token == TOKEN_END):
+                    print("Model called end of sequence at:", cur_i, "/", target_seq_length)
+                    break
+
+            cur_i += 1
+            if(cur_i % 50 == 0):
+                print(cur_i, "/", target_seq_length)
+
+        return gen_seq[:, :cur_i]
+
+# Used as a dummy to nn.Transformer
+class DummyDecoder(nn.Module):
+    def __init__(self):
+        super(DummyDecoder, self).__init__()
+
+    def forward(self, tgt, memory, tgt_mask, memory_mask,tgt_key_padding_mask,memory_key_padding_mask, **kwargs):
+        return memory

model/positional_encoding.py

@@ -0,0 +1,23 @@
+import torch
+import torch.nn as nn
+import math
+
+# PositionalEncoding
+# Taken from https://pytorch.org/tutorials/beginner/transformer_tutorial.html
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + self.pe[:x.size(0), :]
+        return self.dropout(x)

model/rpr.py

@@ -0,0 +1,171 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.nn.parameter import Parameter
+from torch.nn import Module
+from torch.nn.modules.linear import Linear
+from torch.nn.modules.dropout import Dropout
+from torch.nn.modules.normalization import LayerNorm
+from torch.nn.init import *
+
+class TransformerEncoderRPR(Module):
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super(TransformerEncoderRPR, self).__init__()
+        self.layers = torch.nn.ModuleList([encoder_layer for _ in range(num_layers)]) # Fix for tracing
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, src, mask=None, src_key_padding_mask=None, **kwargs):
+        output = src
+        for layer in self.layers:
+            output = layer(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
+        if self.norm:
+            output = self.norm(output)
+        return output
+
+class TransformerEncoderLayerRPR(Module):
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, er_len=None):
+        super(TransformerEncoderLayerRPR, self).__init__()
+        self.self_attn = MultiheadAttentionRPR(d_model, nhead, dropout=dropout, er_len=er_len)
+        self.linear1 = Linear(d_model, dim_feedforward)
+        self.dropout = Dropout(dropout)
+        self.linear2 = Linear(dim_feedforward, d_model)
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.dropout1 = Dropout(dropout)
+        self.dropout2 = Dropout(dropout)
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        src2 = self.self_attn(src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout(F.relu(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+class MultiheadAttentionRPR(Module):
+    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None, er_len=None):
+        super(MultiheadAttentionRPR, self).__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        
+        self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))
+        if not self._qkv_same_embed_dim:
+            self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
+            self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
+            self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
+
+        if bias:
+            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
+        else:
+            self.register_parameter('in_proj_bias', None)
+        self.out_proj = Linear(embed_dim, embed_dim, bias=bias)
+        self.add_zero_attn = add_zero_attn
+
+        if er_len is not None:
+            self.Er = Parameter(torch.rand((er_len, self.head_dim), dtype=torch.float32))
+        else:
+            self.Er = None
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        if self._qkv_same_embed_dim: xavier_uniform_(self.in_proj_weight)
+        else:
+            xavier_uniform_(self.q_proj_weight)
+            xavier_uniform_(self.k_proj_weight)
+            xavier_uniform_(self.v_proj_weight)
+        if self.in_proj_bias is not None:
+            constant_(self.in_proj_bias, 0.)
+            constant_(self.out_proj.bias, 0.)
+
+    def forward(self, query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None):
+        return multi_head_attention_forward_rpr(
+            query, key, value, self.embed_dim, self.num_heads, self.head_dim,
+            self.in_proj_weight, self.in_proj_bias,
+            None, None, self.add_zero_attn,
+            self.dropout, self.out_proj.weight, self.out_proj.bias,
+            training=self.training,
+            key_padding_mask=key_padding_mask, need_weights=need_weights,
+            attn_mask=attn_mask, use_separate_proj_weight=not self._qkv_same_embed_dim,
+            q_proj_weight=getattr(self, 'q_proj_weight', None), 
+            k_proj_weight=getattr(self, 'k_proj_weight', None),
+            v_proj_weight=getattr(self, 'v_proj_weight', None), 
+            rpr_mat=self.Er)
+
+def multi_head_attention_forward_rpr(query, key, value, embed_dim_to_check, num_heads, head_dim,
+                                     in_proj_weight, in_proj_bias, bias_k, bias_v, add_zero_attn,
+                                     dropout_p, out_proj_weight, out_proj_bias, training=True,
+                                     key_padding_mask=None, need_weights=True, attn_mask=None,
+                                     use_separate_proj_weight=False, q_proj_weight=None,
+                                     k_proj_weight=None, v_proj_weight=None, static_k=None,
+                                     static_v=None, rpr_mat=None):
+    
+    tgt_len, bsz, embed_dim = query.size()
+    scaling = float(head_dim) ** -0.5
+
+    if not use_separate_proj_weight:
+        q, k, v = F.linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1)
+    else:
+        q = F.linear(query, q_proj_weight, in_proj_bias[0:embed_dim])
+        k = F.linear(key, k_proj_weight, in_proj_bias[embed_dim:(embed_dim * 2)])
+        v = F.linear(value, v_proj_weight, in_proj_bias[(embed_dim * 2):])
+    
+    q = q * scaling
+    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+    k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+    v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+
+    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
+
+    if rpr_mat is not None:
+        # Safe Explicit Skew
+        len_q = q.shape[1]
+        start_idx = rpr_mat.shape[0] - len_q
+        rpr_mat_valid = rpr_mat[start_idx:, :]
+        qe = torch.einsum("hld,md->hlm", q, rpr_mat_valid)
+        
+        # Indices logic (Flatten -> Gather -> Reshape)
+        B, L, _ = qe.shape
+        # Mask out upper triangle BEFORE skewing
+        mask_tri = torch.triu(torch.ones((L, L), device=qe.device, dtype=torch.bool)).flip(0)
+        qe = qe.masked_fill(~mask_tri, 0.0) # Fill with 0 before shift
+
+        zeros = torch.zeros((B, L, 1), device=qe.device, dtype=qe.dtype)
+        qe_pad = torch.cat([zeros, qe], dim=2).view(B, -1)
+        
+        offsets = torch.arange(L * L, device=qe.device, dtype=torch.int64) + L
+        offsets = offsets.unsqueeze(0).expand(B, -1)
+        srel = torch.gather(qe_pad, 1, offsets).view(B, L, L)
+        
+        attn_output_weights = attn_output_weights + srel
+
+    # --- MASKING FIX (Boolean Masked Fill) ---
+    if attn_mask is not None:
+        if attn_mask.dim() == 2: attn_mask = attn_mask.unsqueeze(0)
+        # ONNX prefers masked_fill with boolean mask over adding -inf
+        is_causal_mask = (attn_mask == float('-inf')) | (attn_mask < -1e4)
+        attn_output_weights = attn_output_weights.masked_fill(is_causal_mask, float('-inf'))
+
+    if key_padding_mask is not None:
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, tgt_len)
+        attn_output_weights = attn_output_weights.masked_fill(
+            key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf')
+        )
+        attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, tgt_len)
+
+    attn_output_weights = F.softmax(attn_output_weights, dim=-1)
+    attn_output_weights = F.dropout(attn_output_weights, p=dropout_p, training=training)
+
+    attn_output = torch.bmm(attn_output_weights, v)
+    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+    attn_output = F.linear(attn_output, out_proj_weight, out_proj_bias)
+
+    if need_weights:
+        return attn_output, attn_output_weights.view(bsz, num_heads, tgt_len, tgt_len).sum(dim=1) / num_heads
+    return attn_output, None

processor.py

@@ -0,0 +1,267 @@
+import pretty_midi
+
+
+RANGE_NOTE_ON = 128
+RANGE_NOTE_OFF = 128
+RANGE_VEL = 32
+RANGE_TIME_SHIFT = 100
+
+START_IDX = {
+    'note_on': 0,
+    'note_off': RANGE_NOTE_ON,
+    'time_shift': RANGE_NOTE_ON + RANGE_NOTE_OFF,
+    'velocity': RANGE_NOTE_ON + RANGE_NOTE_OFF + RANGE_TIME_SHIFT
+}
+
+
+class SustainAdapter:
+    def __init__(self, time, type):
+        self.start =  time
+        self.type = type
+
+
+class SustainDownManager:
+    def __init__(self, start, end):
+        self.start = start
+        self.end = end
+        self.managed_notes = []
+        self._note_dict = {} # key: pitch, value: note.start
+
+    def add_managed_note(self, note: pretty_midi.Note):
+        self.managed_notes.append(note)
+
+    def transposition_notes(self):
+        for note in reversed(self.managed_notes):
+            try:
+                note.end = self._note_dict[note.pitch]
+            except KeyError:
+                note.end = max(self.end, note.end)
+            self._note_dict[note.pitch] = note.start
+
+
+# Divided note by note_on, note_off
+class SplitNote:
+    def __init__(self, type, time, value, velocity):
+        ## type: note_on, note_off
+        self.type = type
+        self.time = time
+        self.velocity = velocity
+        self.value = value
+
+    def __repr__(self):
+        return '<[SNote] time: {} type: {}, value: {}, velocity: {}>'\
+            .format(self.time, self.type, self.value, self.velocity)
+
+
+class Event:
+    def __init__(self, event_type, value):
+        self.type = event_type
+        self.value = value
+
+    def __repr__(self):
+        return '<Event type: {}, value: {}>'.format(self.type, self.value)
+
+    def to_int(self):
+        return START_IDX[self.type] + self.value
+
+    @staticmethod
+    def from_int(int_value):
+        info = Event._type_check(int_value)
+        return Event(info['type'], info['value'])
+
+    @staticmethod
+    def _type_check(int_value):
+        range_note_on = range(0, RANGE_NOTE_ON)
+        range_note_off = range(RANGE_NOTE_ON, RANGE_NOTE_ON+RANGE_NOTE_OFF)
+        range_time_shift = range(RANGE_NOTE_ON+RANGE_NOTE_OFF,RANGE_NOTE_ON+RANGE_NOTE_OFF+RANGE_TIME_SHIFT)
+
+        valid_value = int_value
+
+        if int_value in range_note_on:
+            return {'type': 'note_on', 'value': valid_value}
+        elif int_value in range_note_off:
+            valid_value -= RANGE_NOTE_ON
+            return {'type': 'note_off', 'value': valid_value}
+        elif int_value in range_time_shift:
+            valid_value -= (RANGE_NOTE_ON + RANGE_NOTE_OFF)
+            return {'type': 'time_shift', 'value': valid_value}
+        else:
+            valid_value -= (RANGE_NOTE_ON + RANGE_NOTE_OFF + RANGE_TIME_SHIFT)
+            return {'type': 'velocity', 'value': valid_value}
+
+
+def _divide_note(notes):
+    result_array = []
+    notes.sort(key=lambda x: x.start)
+
+    for note in notes:
+        on = SplitNote('note_on', note.start, note.pitch, note.velocity)
+        off = SplitNote('note_off', note.end, note.pitch, None)
+        result_array += [on, off]
+    return result_array
+
+
+def _merge_note(snote_sequence):
+    note_on_dict = {}
+    result_array = []
+
+    for snote in snote_sequence:
+        # print(note_on_dict)
+        if snote.type == 'note_on':
+            note_on_dict[snote.value] = snote
+        elif snote.type == 'note_off':
+            try:
+                on = note_on_dict[snote.value]
+                off = snote
+                if off.time - on.time == 0:
+                    continue
+                result = pretty_midi.Note(on.velocity, snote.value, on.time, off.time)
+                result_array.append(result)
+            except:
+                print('info removed pitch: {}'.format(snote.value))
+    return result_array
+
+
+def _snote2events(snote: SplitNote, prev_vel: int):
+    result = []
+    if snote.velocity is not None:
+        modified_velocity = snote.velocity // 4
+        if prev_vel != modified_velocity:
+            result.append(Event(event_type='velocity', value=modified_velocity))
+    result.append(Event(event_type=snote.type, value=snote.value))
+    return result
+
+
+def _event_seq2snote_seq(event_sequence):
+    timeline = 0
+    velocity = 0
+    snote_seq = []
+
+    for event in event_sequence:
+        if event.type == 'time_shift':
+            timeline += ((event.value+1) / 100)
+        if event.type == 'velocity':
+            velocity = event.value * 4
+        else:
+            snote = SplitNote(event.type, timeline, event.value, velocity)
+            snote_seq.append(snote)
+    return snote_seq
+
+
+def _make_time_sift_events(prev_time, post_time):
+    time_interval = int(round((post_time - prev_time) * 100))
+    results = []
+    while time_interval >= RANGE_TIME_SHIFT:
+        results.append(Event(event_type='time_shift', value=RANGE_TIME_SHIFT-1))
+        time_interval -= RANGE_TIME_SHIFT
+    if time_interval == 0:
+        return results
+    else:
+        return results + [Event(event_type='time_shift', value=time_interval-1)]
+
+
+def _control_preprocess(ctrl_changes):
+    sustains = []
+
+    manager = None
+    for ctrl in ctrl_changes:
+        if ctrl.value >= 64 and manager is None:
+            # sustain down
+            manager = SustainDownManager(start=ctrl.time, end=None)
+        elif ctrl.value < 64 and manager is not None:
+            # sustain up
+            manager.end = ctrl.time
+            sustains.append(manager)
+            manager = None
+        elif ctrl.value < 64 and len(sustains) > 0:
+            sustains[-1].end = ctrl.time
+    return sustains
+
+
+def _note_preprocess(susteins, notes):
+    note_stream = []
+
+    if susteins:    # if the midi file has sustain controls
+        for sustain in susteins:
+            for note_idx, note in enumerate(notes):
+                if note.start < sustain.start:
+                    note_stream.append(note)
+                elif note.start > sustain.end:
+                    notes = notes[note_idx:]
+                    sustain.transposition_notes()
+                    break
+                else:
+                    sustain.add_managed_note(note)
+
+        for sustain in susteins:
+            note_stream += sustain.managed_notes
+    
+    else:       # else, just push everything into note stream
+        for note_idx, note in enumerate(notes):
+            note_stream.append(note)
+
+    note_stream.sort(key= lambda x: x.start)
+    return note_stream
+
+
+def encode_midi(file_path):
+    events = []
+    notes = []
+    mid = pretty_midi.PrettyMIDI(midi_file=file_path)
+
+    for inst in mid.instruments:
+        inst_notes = inst.notes
+        # ctrl.number is the number of sustain control. If you want to know abour the number type of control,
+        # see https://www.midi.org/specifications-old/item/table-3-control-change-messages-data-bytes-2
+        ctrls = _control_preprocess([ctrl for ctrl in inst.control_changes if ctrl.number == 64])
+        notes += _note_preprocess(ctrls, inst_notes)
+
+    dnotes = _divide_note(notes)
+
+    # print(dnotes)
+    dnotes.sort(key=lambda x: x.time)
+    # print('sorted:')
+    # print(dnotes)
+    cur_time = 0
+    cur_vel = 0
+    for snote in dnotes:
+        events += _make_time_sift_events(prev_time=cur_time, post_time=snote.time)
+        events += _snote2events(snote=snote, prev_vel=cur_vel)
+        # events += _make_time_sift_events(prev_time=cur_time, post_time=snote.time)
+
+        cur_time = snote.time
+        cur_vel = snote.velocity
+
+    return [e.to_int() for e in events]
+
+
+def decode_midi(idx_array, file_path=None):
+    event_sequence = [Event.from_int(idx) for idx in idx_array]
+    # print(event_sequence)
+    snote_seq = _event_seq2snote_seq(event_sequence)
+    note_seq = _merge_note(snote_seq)
+    note_seq.sort(key=lambda x:x.start)
+
+    mid = pretty_midi.PrettyMIDI()
+    # if want to change instument, see https://www.midi.org/specifications/item/gm-level-1-sound-set
+    instument = pretty_midi.Instrument(0, False, "Composed by Super Piano Music Transformer AI")
+    instument.notes = note_seq
+
+    mid.instruments.append(instument)
+    if file_path is not None:
+        mid.write(file_path)
+    return mid
+
+
+if __name__ == '__main__':
+    encoded = encode_midi('bin/ADIG04.mid')
+    print(encoded)
+    decided = decode_midi(encoded,file_path='bin/test.mid')
+
+    ins = pretty_midi.PrettyMIDI('bin/ADIG04.mid')
+    print(ins)
+    print(ins.instruments[0])
+    for i in ins.instruments:
+        print(i.control_changes)
+        print(i.notes)
+

utilities/argument_funcs.py

@@ -0,0 +1,228 @@
+import argparse
+
+from .constants import SEPERATOR
+
+# parse_train_args
+def parse_train_args():
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Argparse arguments for training a model
+    ----------
+    """
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("-input_dir", type=str, default="./dataset/e_piano", help="Folder of preprocessed and pickled midi files")
+    parser.add_argument("-output_dir", type=str, default="./saved_models", help="Folder to save model weights. Saves one every epoch")
+    parser.add_argument("-weight_modulus", type=int, default=1, help="How often to save epoch weights (ex: value of 10 means save every 10 epochs)")
+    parser.add_argument("-print_modulus", type=int, default=1, help="How often to print train results for a batch (batch loss, learn rate, etc.)")
+
+    parser.add_argument("-n_workers", type=int, default=1, help="Number of threads for the dataloader")
+    parser.add_argument("--force_cpu", action="store_true", help="Forces model to run on a cpu even when gpu is available")
+    parser.add_argument("--no_tensorboard", action="store_true", help="Turns off tensorboard result reporting")
+
+    parser.add_argument("-continue_weights", type=str, default=None, help="Model weights to continue training based on")
+    parser.add_argument("-continue_epoch", type=int, default=None, help="Epoch the continue_weights model was at")
+
+    parser.add_argument("-lr", type=float, default=None, help="Constant learn rate. Leave as None for a custom scheduler.")
+    parser.add_argument("-ce_smoothing", type=float, default=None, help="Smoothing parameter for smoothed cross entropy loss (defaults to no smoothing)")
+    parser.add_argument("-batch_size", type=int, default=2, help="Batch size to use")
+    parser.add_argument("-epochs", type=int, default=100, help="Number of epochs to use")
+
+    parser.add_argument("--rpr", action="store_true", help="Use a modified Transformer for Relative Position Representations")
+    parser.add_argument("-max_sequence", type=int, default=2048, help="Maximum midi sequence to consider")
+    parser.add_argument("-n_layers", type=int, default=6, help="Number of decoder layers to use")
+    parser.add_argument("-num_heads", type=int, default=8, help="Number of heads to use for multi-head attention")
+    parser.add_argument("-d_model", type=int, default=512, help="Dimension of the model (output dim of embedding layers, etc.)")
+
+    parser.add_argument("-dim_feedforward", type=int, default=1024, help="Dimension of the feedforward layer")
+
+    parser.add_argument("-dropout", type=float, default=0.1, help="Dropout rate")
+
+    return parser.parse_args()
+
+# print_train_args
+def print_train_args(args):
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Prints training arguments
+    ----------
+    """
+
+    print(SEPERATOR)
+    print("input_dir:", args.input_dir)
+    print("output_dir:", args.output_dir)
+    print("weight_modulus:", args.weight_modulus)
+    print("print_modulus:", args.print_modulus)
+    print("")
+    print("n_workers:", args.n_workers)
+    print("force_cpu:", args.force_cpu)
+    print("tensorboard:", not args.no_tensorboard)
+    print("")
+    print("continue_weights:", args.continue_weights)
+    print("continue_epoch:", args.continue_epoch)
+    print("")
+    print("lr:", args.lr)
+    print("ce_smoothing:", args.ce_smoothing)
+    print("batch_size:", args.batch_size)
+    print("epochs:", args.epochs)
+    print("")
+    print("rpr:", args.rpr)
+    print("max_sequence:", args.max_sequence)
+    print("n_layers:", args.n_layers)
+    print("num_heads:", args.num_heads)
+    print("d_model:", args.d_model)
+    print("")
+    print("dim_feedforward:", args.dim_feedforward)
+    print("dropout:", args.dropout)
+    print(SEPERATOR)
+    print("")
+
+# parse_eval_args
+def parse_eval_args():
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Argparse arguments for evaluating a model
+    ----------
+    """
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("-dataset_dir", type=str, default="./dataset/e_piano", help="Folder of preprocessed and pickled midi files")
+    parser.add_argument("-model_weights", type=str, default="./saved_models/model.pickle", help="Pickled model weights file saved with torch.save and model.state_dict()")
+    parser.add_argument("-n_workers", type=int, default=1, help="Number of threads for the dataloader")
+    parser.add_argument("--force_cpu", action="store_true", help="Forces model to run on a cpu even when gpu is available")
+
+    parser.add_argument("-batch_size", type=int, default=2, help="Batch size to use")
+
+    parser.add_argument("--rpr", action="store_true", help="Use a modified Transformer for Relative Position Representations")
+    parser.add_argument("-max_sequence", type=int, default=2048, help="Maximum midi sequence to consider in the model")
+    parser.add_argument("-n_layers", type=int, default=6, help="Number of decoder layers to use")
+    parser.add_argument("-num_heads", type=int, default=8, help="Number of heads to use for multi-head attention")
+    parser.add_argument("-d_model", type=int, default=512, help="Dimension of the model (output dim of embedding layers, etc.)")
+
+    parser.add_argument("-dim_feedforward", type=int, default=1024, help="Dimension of the feedforward layer")
+
+    return parser.parse_args()
+
+# print_eval_args
+def print_eval_args(args):
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Prints evaluation arguments
+    ----------
+    """
+
+    print(SEPERATOR)
+    print("dataset_dir:", args.dataset_dir)
+    print("model_weights:", args.model_weights)
+    print("n_workers:", args.n_workers)
+    print("force_cpu:", args.force_cpu)
+    print("")
+    print("batch_size:", args.batch_size)
+    print("")
+    print("rpr:", args.rpr)
+    print("max_sequence:", args.max_sequence)
+    print("n_layers:", args.n_layers)
+    print("num_heads:", args.num_heads)
+    print("d_model:", args.d_model)
+    print("")
+    print("dim_feedforward:", args.dim_feedforward)
+    print(SEPERATOR)
+    print("")
+
+# parse_generate_args
+def parse_generate_args():
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Argparse arguments for generation
+    ----------
+    """
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("-midi_root", type=str, default="./dataset/e_piano/", help="Midi file to prime the generator with")
+    parser.add_argument("-output_dir", type=str, default="./gen", help="Folder to write generated midi to")
+    parser.add_argument("-primer_file", type=str, default=None, help="File path or integer index to the evaluation dataset. Default is to select a random index.")
+    parser.add_argument("--force_cpu", action="store_true", help="Forces model to run on a cpu even when gpu is available")
+
+    parser.add_argument("-target_seq_length", type=int, default=1024, help="Target length you'd like the midi to be")
+    parser.add_argument("-num_prime", type=int, default=256, help="Amount of messages to prime the generator with")
+    parser.add_argument("-model_weights", type=str, default="./saved_models/model.pickle", help="Pickled model weights file saved with torch.save and model.state_dict()")
+    parser.add_argument("-beam", type=int, default=0, help="Beam search k. 0 for random probability sample and 1 for greedy")
+
+    parser.add_argument("--rpr", action="store_true", help="Use a modified Transformer for Relative Position Representations")
+    parser.add_argument("-max_sequence", type=int, default=2048, help="Maximum midi sequence to consider")
+    parser.add_argument("-n_layers", type=int, default=6, help="Number of decoder layers to use")
+    parser.add_argument("-num_heads", type=int, default=8, help="Number of heads to use for multi-head attention")
+    parser.add_argument("-d_model", type=int, default=512, help="Dimension of the model (output dim of embedding layers, etc.)")
+
+    parser.add_argument("-dim_feedforward", type=int, default=1024, help="Dimension of the feedforward layer")
+
+    return parser.parse_args()
+
+# print_generate_args
+def print_generate_args(args):
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Prints generation arguments
+    ----------
+    """
+
+    print(SEPERATOR)
+    print("midi_root:", args.midi_root)
+    print("output_dir:", args.output_dir)
+    print("primer_file:", args.primer_file)
+    print("force_cpu:", args.force_cpu)
+    print("")
+    print("target_seq_length:", args.target_seq_length)
+    print("num_prime:", args.num_prime)
+    print("model_weights:", args.model_weights)
+    print("beam:", args.beam)
+    print("")
+    print("rpr:", args.rpr)
+    print("max_sequence:", args.max_sequence)
+    print("n_layers:", args.n_layers)
+    print("num_heads:", args.num_heads)
+    print("d_model:", args.d_model)
+    print("")
+    print("dim_feedforward:", args.dim_feedforward)
+    print(SEPERATOR)
+    print("")
+
+# write_model_params
+def write_model_params(args, output_file):
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Writes given training parameters to text file
+    ----------
+    """
+
+    o_stream = open(output_file, "w")
+
+    o_stream.write("rpr: " + str(args.rpr) + "\n")
+    o_stream.write("lr: " + str(args.lr) + "\n")
+    o_stream.write("ce_smoothing: " + str(args.ce_smoothing) + "\n")
+    o_stream.write("batch_size: " + str(args.batch_size) + "\n")
+    o_stream.write("max_sequence: " + str(args.max_sequence) + "\n")
+    o_stream.write("n_layers: " + str(args.n_layers) + "\n")
+    o_stream.write("num_heads: " + str(args.num_heads) + "\n")
+    o_stream.write("d_model: " + str(args.d_model) + "\n")
+    o_stream.write("dim_feedforward: " + str(args.dim_feedforward) + "\n")
+    o_stream.write("dropout: " + str(args.dropout) + "\n")
+
+    o_stream.close()

utilities/constants.py

@@ -0,0 +1,28 @@
+import torch
+
+from processor import RANGE_NOTE_ON, RANGE_NOTE_OFF, RANGE_VEL, RANGE_TIME_SHIFT
+
+SEPERATOR               = "========================="
+
+# Taken from the paper
+ADAM_BETA_1             = 0.9
+ADAM_BETA_2             = 0.98
+ADAM_EPSILON            = 10e-9
+
+LR_DEFAULT_START        = 1.0
+SCHEDULER_WARMUP_STEPS  = 4000
+# LABEL_SMOOTHING_E       = 0.1
+
+# DROPOUT_P               = 0.1
+
+TOKEN_END               = RANGE_NOTE_ON + RANGE_NOTE_OFF + RANGE_VEL + RANGE_TIME_SHIFT
+TOKEN_PAD               = TOKEN_END + 1
+
+VOCAB_SIZE              = TOKEN_PAD + 1
+
+TORCH_FLOAT             = torch.float32
+TORCH_INT               = torch.int32
+
+TORCH_LABEL_TYPE        = torch.long
+
+PREPEND_ZEROS_WIDTH     = 4

utilities/device.py

@@ -0,0 +1,67 @@
+# For all things related to devices
+#### ONLY USE PROVIDED FUNCTIONS, DO NOT USE GLOBAL CONSTANTS ####
+
+import torch
+
+TORCH_CPU_DEVICE = torch.device("cpu")
+
+if(torch.cuda.device_count() > 0):
+    TORCH_CUDA_DEVICE = torch.device("cuda")
+else:
+    print("----- WARNING: CUDA devices not detected. This will cause the model to run very slow! -----")
+    print("")
+    TORCH_CUDA_DEVICE = None
+
+USE_CUDA = True
+
+# use_cuda
+def use_cuda(cuda_bool):
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Sets whether to use CUDA (if available), or use the CPU (not recommended)
+    ----------
+    """
+
+    global USE_CUDA
+    USE_CUDA = cuda_bool
+
+# get_device
+def get_device():
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Grabs the default device. Default device is CUDA if available and use_cuda is not False, CPU otherwise.
+    ----------
+    """
+
+    if((not USE_CUDA) or (TORCH_CUDA_DEVICE is None)):
+        return TORCH_CPU_DEVICE
+    else:
+        return TORCH_CUDA_DEVICE
+
+# cuda_device
+def cuda_device():
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Grabs the cuda device (may be None if CUDA is not available)
+    ----------
+    """
+
+    return TORCH_CUDA_DEVICE
+
+# cpu_device
+def cpu_device():
+    """
+    ----------
+    Author: Damon Gwinn
+    ----------
+    Grabs the cpu device
+    ----------
+    """
+
+    return TORCH_CPU_DEVICE