code/data/maestro/load.py

from pathlib import Path
from random import shuffle
from miditok import REMI, TokenizerConfig
from miditok.utils import split_files_for_training
from miditok.data_augmentation import augment_dataset
from miditok.pytorch_data import DatasetMIDI

TOKENIZER_PARAMS = {
    "pitch_range": (21, 109),
    "beat_res": {(0, 4): 8, (4, 12): 4},
    "num_velocities": 32,
    "special_tokens": ["PAD", "BOS", "EOS", "MASK"],
    "use_chords": True,
    "use_rests": False,
    "use_tempos": True,
    "use_time_signatures": False,
    "use_programs": False,
    "num_tempos": 32,  # number of tempo bins
    "tempo_range": (40, 250),  # (min, max)
}
config = TokenizerConfig(**TOKENIZER_PARAMS)

def get_maestro_dataset(data_dir: str, block_size: int, augment: bool = False):
    midi_paths = list(Path(data_dir).resolve().glob("20[0-9][0-9]/*.midi"))  
    num_total_files = len(midi_paths)
    num_valid_files = round(num_total_files * 0.05)
    num_test_files = round(num_total_files * 0.05)
    shuffle(midi_paths)    
    midi_paths_valid = midi_paths[:num_valid_files]
    midi_paths_test = midi_paths[num_valid_files:num_valid_files + num_test_files]
    midi_paths_train = midi_paths[num_valid_files + num_test_files:]
    tokenizer = REMI(config)
    ds = []
    # Process each subset
    for files_paths, subset_name in ((midi_paths_train, "train"), 
                                     (midi_paths_valid, "valid"), 
                                     (midi_paths_test, "test")):
        subset_chunks_dir = Path(data_dir, f"dataset_{subset_name}")
        if augment and subset_name == "train":
            subset_chunks_dir = subset_chunks_dir.parent / (subset_chunks_dir.name + "_augmented")
        subset_chunks_dir.mkdir(parents=True, exist_ok=True)
        
        midi_paths_subset = process_subset(
            files_paths=files_paths,
            subset_name=subset_name,
            subset_chunks_dir=subset_chunks_dir,
            tokenizer=tokenizer,
            block_size=block_size,
            augment=augment
        )
        ds.append(DatasetMIDI(files_paths=midi_paths_subset, tokenizer=tokenizer, max_seq_len=block_size))
    train_ds, val_ds, test_ds = ds
    print(f"Number of training files: {len(train_ds)}")
    print(f"Number of validation files: {len(val_ds)}")
    print(f"Number of test files: {len(test_ds)}")
    return train_ds, val_ds, test_ds, tokenizer

def process_subset(files_paths, subset_name, subset_chunks_dir, tokenizer, block_size, augment):
    chunks_exist = any(subset_chunks_dir.iterdir())
    if not chunks_exist:
        create_chunks(files_paths, subset_name, subset_chunks_dir, tokenizer, block_size, augment)
    return list(subset_chunks_dir.resolve().glob("**/*.midi"))

def create_chunks(files_paths, subset_name, subset_chunks_dir, tokenizer, block_size, augment):
    print(f"Splitting files for {subset_name} subset into chunks of {block_size}...")
    split_files_for_training(
        files_paths=files_paths,
        tokenizer=tokenizer,
        save_dir=subset_chunks_dir,
        max_seq_len=block_size,
        num_overlap_bars=2,
    )
    # Apply data augmentation on train split if requested
    if augment and subset_name == "train":
        print("Performing data augmentation on the training set...")
        augment_dataset(
            subset_chunks_dir,
            pitch_offsets=[-3, -2, -1, 0, 1, 2, 3],
            velocity_offsets=[0,],
            duration_offsets=[0,],
            out_path=subset_chunks_dir
        )