fix data augment with unique_track, Program tokens created from confi…

…g.programs

miditok/constants.py

@@ -2,7 +2,7 @@
 
 """
 
-CURRENT_VERSION_PACKAGE = "2.1.0"  # used when saving the config of a tokenizer
+CURRENT_VERSION_PACKAGE = "2.1.1"  # used when saving the config of a tokenizer
 
 MIDI_FILES_EXTENSIONS = [".mid", ".midi", ".MID", ".MIDI"]
 

miditok/data_augmentation/data_augmentation.py

@@ -90,7 +90,10 @@ def data_augmentation_dataset(
                 Tuple[int, int, int], List[Union[int, List[int]]]
             ] = {}
             for track, (_, is_drum) in zip(ids, programs):
-                if is_drum:  # we dont augment drums
+                # we dont augment drums
+                if not tokenizer.unique_track and is_drum:
+                    continue
+                elif tokenizer.unique_track and all(p[1] for p in programs):
                     continue
                 corrected_offsets = deepcopy(offsets)
                 vel_dim = int(128 / len(tokenizer.velocities))
@@ -113,12 +116,13 @@ def data_augmentation_dataset(
                         augmented_tokens[aug_offsets].append(seq)
                     except KeyError:
                         augmented_tokens[aug_offsets] = [seq]
-            for i, (track, (_, is_drum)) in enumerate(
-                zip(ids, programs)
-            ):  # adding drums to all already augmented
-                if is_drum:
-                    for aug_offsets in augmented_tokens:
-                        augmented_tokens[aug_offsets].insert(i, track)
+            if not tokenizer.unique_track:
+                for i, (track, (_, is_drum)) in enumerate(
+                    zip(ids, programs)
+                ):  # adding drums to all already augmented
+                    if is_drum:
+                        for aug_offsets in augmented_tokens:
+                            augmented_tokens[aug_offsets].insert(i, track)
 
             # Save augmented tracks as json
             for aug_offsets, tracks_seq in augmented_tokens.items():
@@ -138,6 +142,8 @@ def data_augmentation_dataset(
                 nb_augmentations += 1
                 nb_tracks_augmented += len(tracks_seq)
             if copy_original_in_new_location and out_path is not None:
+                if tokenizer.unique_track:
+                    ids = ids[0]
                 tokenizer.save_tokens(
                     ids, out_path / f"{file_path.stem}.json", programs
                 )
@@ -448,6 +454,17 @@ def data_augmentation_tokens(
             )
             note_off_tokens = np.array(tokenizer.token_ids_of_type("NoteOff"))
             mask_pitch = np.isin(tokens, pitch_tokens)
+            # If applicable, removes drum notes from the mask
+            if tokenizer.unique_track:
+                for idx, is_note in enumerate(mask_pitch):
+                    if (
+                        is_note
+                        and idx > 0
+                        and tokenizer[tokens[idx - 1]] == "Program_-1"
+                    ):
+                        mask_pitch[idx] = False
+                        if len(note_off_tokens) > 0:
+                            note_off_tokens[idx] = False
         else:
             pitch_tokens = np.array(tokenizer.token_ids_of_type("Pitch", pitch_voc_idx))
             mask_pitch = np.full_like(tokens, 0, dtype=np.bool_)

miditok/midi_tokenizer.py

@@ -1025,9 +1025,7 @@ def learn_bpe(
                     sample["ids"], as_one_str=True
                 )  # list of str (bytes)
                 iterator += (
-                    [[byte_] for byte_ in bytes_]
-                    if not self.unique_track
-                    else [bytes_]
+                    [[byte_] for byte_ in bytes_] if not self.unique_track else [bytes_]
                 )
 
             # This doesn't seem to work, the trainer pre-processes the sequences, but then no word remains

miditok/tokenizations/cp_word.py

@@ -410,7 +410,8 @@ def _create_base_vocabulary(self) -> List[List[str]]:
         # PROGRAM
         if self.config.use_programs:
             vocab += [
-                ["Ignore_None"] + [f"Program_{program}" for program in range(-1, 128)]
+                ["Ignore_None"]
+                + [f"Program_{program}" for program in self.config.programs]
             ]
 
         # CHORD

miditok/tokenizations/midi_like.py

@@ -312,7 +312,7 @@ def _create_base_vocabulary(self) -> List[str]:
 
         # PROGRAM
         if self.config.use_programs:
-            vocab += [f"Program_{program}" for program in range(-1, 128)]
+            vocab += [f"Program_{program}" for program in self.config.programs]
 
         return vocab
 

miditok/tokenizations/octuple.py

@@ -301,7 +301,7 @@ def tokens_to_midi(
         current_time_sig_tick = 0
         current_time_sig_bar = 0
 
-        tracks = dict([(n, []) for n in range(-1, 128)])
+        tracks = dict([(n, []) for n in self.config.programs])
         for time_step in tokens:
             if any(tok.split("_")[1] == "None" for tok in time_step[:6]):
                 continue  # Either padding, mask: error of prediction or end of sequence anyway

miditok/tokenizations/remi.py

@@ -329,7 +329,7 @@ def _create_base_vocabulary(self, sos_eos_tokens: bool = None) -> List[str]:
 
         # PROGRAM
         if self.config.use_programs:
-            vocab += [f"Program_{program}" for program in range(-1, 128)]
+            vocab += [f"Program_{program}" for program in self.config.programs]
 
         return vocab
 

miditok/tokenizations/structured.py

@@ -222,7 +222,7 @@ def _create_base_vocabulary(self, sos_eos_tokens: bool = None) -> List[str]:
 
         # PROGRAM
         if self.config.use_programs:
-            vocab += [f"Program_{program}" for program in range(-1, 128)]
+            vocab += [f"Program_{program}" for program in self.config.programs]
 
         return vocab
 

miditok/tokenizations/tsd.py

@@ -426,7 +426,7 @@ def _create_base_vocabulary(self, sos_eos_tokens: bool = False) -> List[str]:
 
         # PROGRAM
         if self.config.use_programs:
-            vocab += [f"Program_{program}" for program in range(-1, 128)]
+            vocab += [f"Program_{program}" for program in self.config.programs]
 
         # TIME SIGNATURE
         if self.config.use_time_signatures:

setup.py

@@ -8,7 +8,7 @@
     author="Nathan Fradet",
     url="https://github.com/Natooz/MidiTok",
     packages=find_packages(exclude=("tests",)),
-    version="2.1.0",
+    version="2.1.1",
     license="MIT",
     description="A convenient MIDI tokenizer for Deep Learning networks, with multiple encoding strategies",
     long_description=long_description,