Fix PretrainingDataset.build bottleneck (#92)

* improve logging

* fix

* fix dataset building bottlekneck

aria/data/datasets.py

@@ -190,6 +190,10 @@ def build_to_file(
     def combine_datasets_from_file(cls, *args: str, output_path: str):
         """Function for concatenating jsonl files, checking for duplicates"""
         logger = setup_logger()
+
+        for input_path in args:
+            assert os.path.isfile(input_path), f"{input_path} doesn't exist"
+
         dupe_cnt = 0
         hashes = {}
         with jsonlines.open(output_path, mode="w") as f_out:
@@ -208,6 +212,10 @@ def combine_datasets_from_file(cls, *args: str, output_path: str):
                             f_out.write(msg_dict)
                             hashes[midi_dict_hash] = True
                 logger.info(f"Finished processing: {input_path}")
+                logger.info(
+                    f"{len(hashes)} unique midi_dicts and {dupe_cnt} duplicates so far"
+                )
+
 
         logger.info(
             f"Found {len(hashes)} unique midi_dicts and {dupe_cnt} duplicates"
@@ -299,25 +307,32 @@ def _get_mididicts_mp(_paths):
         with Pool() as pool:
             results = pool.imap(_get_mididict, _paths)
             seen_hashes = defaultdict(list)
+            dupe_cnt = 0
+            failed_cnt = 0
             for idx, (success, result) in enumerate(results):
                 if idx % 50 == 0 and idx != 0:
                     logger.info(f"Processed MIDI files: {idx}/{num_paths}")
 
                 if not success:
+                    failed_cnt += 1
                     continue
                 else:
                     mid_dict, mid_hash, mid_path = result
 
                 if seen_hashes.get(mid_hash):
                     logger.info(
                         f"MIDI located at '{mid_path}' is a duplicate - already"
-                        f" seen at: {seen_hashes[mid_hash]}"
+                        f" seen at: {seen_hashes[mid_hash][0]}"
                     )
                     seen_hashes[mid_hash].append(str(mid_path))
+                    dupe_cnt += 1
                 else:
                     seen_hashes[mid_hash].append(str(mid_path))
                     yield mid_dict
 
+        print(f"Total duplicates: {dupe_cnt}")
+        print(f"Total processing fails (tests or otherwise): {failed_cnt}")
+
     logger = setup_logger()
     if get_start_method() == "spawn":
         logger.warning(
@@ -523,8 +538,12 @@ def _get_seqs(_entry: MidiDict | dict, _tokenizer: Tokenizer):
 
     if isinstance(_entry, str):
         _midi_dict = MidiDict.from_msg_dict(json.loads(_entry.rstrip()))
-    else:
+    elif isinstance(_entry, dict):
+        _midi_dict = MidiDict.from_msg_dict(_entry)
+    elif isinstance(_entry, MidiDict):
         _midi_dict = _entry
+    else:
+        raise Exception
 
     try:
         _tokenized_seq = _tokenizer.tokenize(_midi_dict)
@@ -584,6 +603,19 @@ def get_seqs(
             if not any(proc.is_alive() for proc in workers):
                 break
 
+def reservoir(_iterable: Iterable, k: int):
+    _reservoir = []
+    for entry in _iterable:
+        if entry is not None:
+            _reservoir.append(entry)
+
+        if len(_reservoir) >= k:
+            random.shuffle(_reservoir)
+            yield from _reservoir
+            _reservoir = []
+
+    if _reservoir != []:
+        yield from _reservoir
 
 class PretrainingDataset(TrainingDataset):
     def __init__(self, dir_path: str, tokenizer: Tokenizer):
@@ -695,13 +727,18 @@ def _build_epoch(_save_path, _midi_dataset):
                 )
 
                 buffer = []
-                # TODO: Profile why mp takes a while to spit up
-                for entry in get_seqs(tokenizer, _midi_dataset):
+                _idx = 0
+                for entry in reservoir(get_seqs(tokenizer, _midi_dataset), 5):
                     if entry is not None:
                         buffer += entry
                     while len(buffer) >= max_seq_len:
                         writer.write(buffer[:max_seq_len])
                         buffer = buffer[max_seq_len:]
+
+                    _idx += 1
+                    if _idx % 250 == 0:
+                        logger.info(f"Finished processing {_idx}")
+
                 buffer += [tokenizer.pad_tok] * (max_seq_len - len(buffer))
 
         logger = setup_logger()
@@ -727,32 +764,27 @@ def _build_epoch(_save_path, _midi_dataset):
         if not os.path.exists(save_dir):
             os.mkdir(save_dir)
 
-        # TODO: This is very slow right now
-        if not midi_dataset:
-            midi_dataset = MidiDataset.load(midi_dataset_path)
-        else:
+        if not midi_dataset and not midi_dataset_path:
             Exception("Must provide either midi_dataset or midi_dataset_path")
+        if midi_dataset and midi_dataset_path:
+            Exception("Can't provide both midi_dataset and midi_dataset_path")
 
         logger.info(
             f"Building PretrainingDataset with config: "
             f"max_seq_len={max_seq_len}, "
             f"tokenizer_name={tokenizer.name}"
         )
-        _num_proc = os.cpu_count()
-        if 2 * _num_proc > len(midi_dataset):
-            logger.warning(
-                "Number of processes is close to the number of MidiDicts "
-                "in the dataset. This can result in shuffling not working "
-                "as intended when building different epochs"
-            )
         for idx in range(num_epochs):
             logger.info(f"Building epoch {idx}/{num_epochs - 1}...")
+
+            # Reload the dataset on each iter
+            if midi_dataset_path:
+                midi_dataset = jsonlines.open(midi_dataset_path, "r")
+
             _build_epoch(
                 _save_path=os.path.join(save_dir, f"epoch{idx}.jsonl"),
                 _midi_dataset=midi_dataset,
             )
-            # TODO: This is very slow for large datasets
-            midi_dataset.shuffle()
 
         logger.info(
             f"Finished building, saved PretrainingDataset to {save_dir}"

aria/data/midi.py

@@ -177,6 +177,7 @@ def calculate_hash(self):
         msg_dict_to_hash = self.get_msg_dict()
         # Remove meta when calculating hash
         msg_dict_to_hash.pop("meta_msgs")
+        msg_dict_to_hash.pop("ticks_per_beat")
         msg_dict_to_hash.pop("metadata")
 
         return hashlib.md5(

aria/run.py

@@ -263,7 +263,7 @@ def _parse_midi_dataset_args():
         help="manually add metadata key-value pair when building dataset",
     )
     argp.add_argument(
-        "--split", type=float, help="create train/val split", required=False
+        "-split", type=float, help="create train/val split", required=False
     )
 
     return argp.parse_args(sys.argv[2:])
@@ -284,7 +284,7 @@ def build_midi_dataset(args):
     )
 
     if args.split:
-        assert 0.0 < args.split < 1.0, "Invalid range given for --split"
+        assert 0.0 < args.split < 1.0, "Invalid range given for -split"
         MidiDataset.split_from_file(
             load_path=args.save_path,
             train_val_ratio=args.split,

config/config.json

@@ -4,7 +4,7 @@
             "max_programs":{
                 "run": true,
                 "args": {
-                    "max": 7
+                    "max": 12
                 }
             },
             "max_instruments":{
@@ -24,7 +24,7 @@
                 "run": true,
                 "args": {
                     "min_per_second": 0.5,
-                    "max_per_second": 15
+                    "max_per_second": 16
                 }
             },
             "min_length":{