code/data/jsb/load.py

import os
import numpy as np
import json
from torch.utils.data import Dataset

class JSBDataset(Dataset):
    """ Preprocessor for JSB Chorales dataset. From Boulanger-Lewandowski (2012): 
    "This will load a dictionary with 'train', 'valid' and 'test' keys, 
    with the corresponding values being a list of sequences. 
    Each sequence is itself a list of time steps, and each time step is a 
    list of the non-zero elements in the piano-roll at this instant 
    (in MIDI note numbers, between 21 and 108 inclusive)". """
    def __init__(self, data_dir: str, split: str, max_len: int, pad_token_id: int = 0):
        """
        Arguments:
            data_dir (string): path to the dataset directory.
            split (string): dataset split (train/val/test). 
        """
        self.data_dir = data_dir
        self.max_len = max_len
        with open(os.path.join(data_dir, 'Jsb16thSeparated.json')) as f:
            d = json.load(f)
        self.data = d[split]
        self.vocab_size = 88 + 1 + 1 # [21, 108] MIDI notes + silence (-1) + pad (0)
        self.pad_token_id = pad_token_id

    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        seq = np.array(self.data[idx], dtype=np.int64)
        full_seq = seq.flatten()
        # padding to max_len if shorter
        if len(full_seq) < self.max_len:
            full_seq = np.pad(full_seq, (0, self.max_len - len(full_seq)), 'constant', constant_values=self.pad_token_id)
        else:
            full_seq = full_seq[:self.max_len]
        # NOTE: nn.Embedding layer expects 0 to vocab_size-1.
        # Shift only the token_id for silence to be +1 from its default value of -1.
        full_seq[full_seq == -1] = 1
        x, y = full_seq[:-1], full_seq[1:]
        return x, y