From 8c2349bfb771145c805c8a652392ae8f11ed0756 Mon Sep 17 00:00:00 2001
From: Nathan Fradet <56734983+Natooz@users.noreply.github.com>
Date: Mon, 24 Jul 2023 18:17:45 +0200
Subject: [PATCH] unique_track renamed one_token_stream,
 convert_sequence_to_tokseq method, io_format property

---
 docs/midi_tokenizer.rst                       |  18 +-
 miditok/__init__.py                           |   2 +-
 .../data_augmentation/data_augmentation.py    |  14 +-
 miditok/midi_tokenizer.py                     | 184 ++++++++++++------
 miditok/tokenizations/mumidi.py               |   2 +-
 miditok/tokenizations/octuple.py              |   2 +-
 miditok/tokenizations/remi_plus.py            |   2 +-
 miditok/tokenizations/tsd.py                  |  16 +-
 tests/test_bpe.py                             |   6 +-
 tests/test_io_formats.py                      |  10 +-
 tests/test_methods.py                         |   4 +-
 tests/test_multitrack.py                      |   6 +-
 tests/test_one_track.py                       |   4 +-
 13 files changed, 171 insertions(+), 99 deletions(-)

diff --git a/docs/midi_tokenizer.rst b/docs/midi_tokenizer.rst
index 22e16d98..e8544d2e 100644
--- a/docs/midi_tokenizer.rst
+++ b/docs/midi_tokenizer.rst
@@ -144,16 +144,16 @@ Tokens & TokSequence input / output format
 
 Depending on the tokenizer at use, the **format** of the tokens returned by the ``midi_to_tokens`` method may vary, as well as the expected format for the ``tokens_to_midi`` method. For any tokenizer, the format is the same for both methods.
 
-The format is deduced from the ``is_multi_voc`` and ``unique_track`` tokenizer properties. In short: **unique_track** being True means that the tokenizer will convert a MIDI file into a single stream of tokens for all instrument tracks, otherwise it will convert each track to a distinct token stream; **is_mult_voc** being True means that each token stream is a list of lists of tokens, of shape ``(T,C)`` for T time steps and C subtokens per time step.
+The format is deduced from the ``is_multi_voc`` and ``one_token_stream`` tokenizer properties. In short: **one_token_stream** being True means that the tokenizer will convert a MIDI file into a single stream of tokens for all instrument tracks, otherwise it will convert each track to a distinct token stream; **is_mult_voc** being True means that each token stream is a list of lists of tokens, of shape ``(T,C)`` for T time steps and C subtokens per time step.
 
 This results in four situations, where I is the number of tracks, T is the number of tokens (or time steps) and C the number of subtokens per time step:
 
-* **is_multi_voc** and **unique_track** are both **False**: ``[I,(T)]``
-* **is_multi_voc** is **False** and **unique_track** is **True**: ``(T)``
-* **is_multi_voc** is **True** and **unique_track** is **False**: ``[I,(T,C)]``
-* **is_multi_voc** and **unique_track** are both **True**: ``(T,C)``
+* **is_multi_voc** and **one_token_stream** are both **False**: ``[I,(T)]``
+* **is_multi_voc** is **False** and **one_token_stream** is **True**: ``(T)``
+* **is_multi_voc** is **True** and **one_token_stream** is **False**: ``[I,(T,C)]``
+* **is_multi_voc** and **one_token_stream** are both **True**: ``(T,C)``
 
-**Note that if there is no I dimension in the format, the output of ``midi_to_tokens`` is a ``TokSequence`` object, otherwise it is a list of ``TokSequence`` objects (one per token stream / track).**
+**Note that if there is no I dimension in the format, the output of **``midi_to_tokens``** is a **:class:`miditok.TokSequence`** object, otherwise it is a list of **:class:`miditok.TokSequence`** objects (one per token stream / track).**
 
 Some tokenizer examples to illustrate:
 
@@ -163,6 +163,12 @@ Some tokenizer examples to illustrate:
 * **Octuple** is a multi-voc tokenizer and converts all MIDI track to a single stream of tokens, hence it will convert MIDI files to a ``TokSequence`` object, ``(T,C)`` format.
 
 
+**You can use the **``convert_sequence_to_tokseq``** method to automatically convert a input sequence, of ids (integers) or tokens (string), into a **:class:`miditok.TokSequence`** or list of **:class:`miditok.TokSequence`** objects with the appropriate format of the tokenizer being used.**
+
+.. autofunction:: miditok.convert_sequence_to_tokseq
+    :noindex:
+
+
 Magic methods
 ------------------------
 
diff --git a/miditok/__init__.py b/miditok/__init__.py
index d044f62c..40b5c124 100644
--- a/miditok/__init__.py
+++ b/miditok/__init__.py
@@ -10,7 +10,7 @@
     MuMIDI,
     MMM,
 )
-from .midi_tokenizer import MIDITokenizer
+from .midi_tokenizer import MIDITokenizer, convert_sequence_to_tokseq
 from .classes import Event, TokSequence, TokenizerConfig
 
 from .utils import utils
diff --git a/miditok/data_augmentation/data_augmentation.py b/miditok/data_augmentation/data_augmentation.py
index 74a451b5..3fc50837 100644
--- a/miditok/data_augmentation/data_augmentation.py
+++ b/miditok/data_augmentation/data_augmentation.py
@@ -72,7 +72,7 @@ def data_augmentation_dataset(
                 file = json.load(json_file)
                 ids, programs = file["ids"], file["programs"]
 
-            if tokenizer.unique_track:
+            if tokenizer.one_token_stream:
                 ids = [ids]
 
             # Perform data augmentation for each track
@@ -91,9 +91,9 @@ def data_augmentation_dataset(
             ] = {}
             for track, (_, is_drum) in zip(ids, programs):
                 # we dont augment drums
-                if not tokenizer.unique_track and is_drum:
+                if not tokenizer.one_token_stream and is_drum:
                     continue
-                elif tokenizer.unique_track and all(p[1] for p in programs):
+                elif tokenizer.one_token_stream and all(p[1] for p in programs):
                     continue
                 corrected_offsets = deepcopy(offsets)
                 vel_dim = int(128 / len(tokenizer.velocities))
@@ -109,14 +109,14 @@ def data_augmentation_dataset(
                 if len(aug) == 0:
                     continue
                 for aug_offsets, seq in aug:
-                    if tokenizer.unique_track:
+                    if tokenizer.one_token_stream:
                         augmented_tokens[aug_offsets] = seq
                         continue
                     try:
                         augmented_tokens[aug_offsets].append(seq)
                     except KeyError:
                         augmented_tokens[aug_offsets] = [seq]
-            if not tokenizer.unique_track:
+            if not tokenizer.one_token_stream:
                 for i, (track, (_, is_drum)) in enumerate(
                     zip(ids, programs)
                 ):  # adding drums to all already augmented
@@ -142,7 +142,7 @@ 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:
+                if tokenizer.one_token_stream:
                     ids = ids[0]
                 tokenizer.save_tokens(
                     ids, out_path / f"{file_path.stem}.json", programs
@@ -455,7 +455,7 @@ 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:
+            if tokenizer.one_token_stream:
                 for idx, is_note in enumerate(mask_pitch):
                     if (
                         is_note
diff --git a/miditok/midi_tokenizer.py b/miditok/midi_tokenizer.py
index fc195a1c..48b36292 100644
--- a/miditok/midi_tokenizer.py
+++ b/miditok/midi_tokenizer.py
@@ -32,6 +32,83 @@
 )
 
 
+def convert_sequence_to_tokseq(
+        tokenizer,
+        input_seq,
+        complete_seq: bool = True,
+        decode_bpe: bool = True
+) -> Union[TokSequence, List[TokSequence]]:
+    r"""Converts a sequence into a **:class:`miditok.TokSequence`** or list of **:class:`miditok.TokSequence`**
+    objects with the appropriate format of the tokenizer being used.
+
+    :param tokenizer: tokenizer being used with the sequence.
+    :param input_seq: sequence to convert. It can be a list of ids (integers), tokens (string) or events (Event).
+        It can also be a Pytorch or TensorFlow tensor, or Numpy array representing ids.
+    :param complete_seq: will complete the output sequence(s). (default: True)
+    :param decode_bpe: if the input sequence contains ids, and that they contain BPE tokens, these tokens will
+        be decoded. (default: True)
+    :return:
+    """
+    # Deduce the type of data (ids/tokens/events)
+    try:
+        arg = ("ids", convert_ids_tensors_to_list(input_seq))
+    except (AttributeError, ValueError, TypeError, IndexError):
+        if isinstance(input_seq[0], str) or (
+                isinstance(input_seq[0], list) and isinstance(input_seq[0][0], str)
+        ):
+            arg = ("tokens", input_seq)
+        else:  # list of Event, but unlikely
+            arg = ("events", input_seq)
+
+    # Deduce nb of subscripts / dims
+    nb_io_dims = len(tokenizer.io_format)
+    nb_seq_dims = 1
+    if isinstance(arg[1][0], list):
+        nb_seq_dims += 1
+        if isinstance(arg[1][0][0], list):
+            nb_seq_dims += 1
+
+    # Check the number of dimensions is good
+    # In case of no one_token_stream and one dimension short --> unsqueeze
+    if not tokenizer.one_token_stream and nb_seq_dims == nb_io_dims - 1:
+        print(f"The input sequence has one dimension less than expected ({nb_seq_dims} instead of "
+              f"{nb_io_dims}). It is being unsqueezed to conform with the tokenizer's i/o format "
+              f"({tokenizer.io_format})")
+        arg = (arg[0], [arg[1]])
+
+    elif nb_seq_dims != nb_io_dims:
+        raise ValueError(f"The input sequence does not have the expected dimension "
+                         f"({nb_seq_dims} instead of {nb_io_dims}).")
+
+    # Convert to TokSequence
+    if not tokenizer.one_token_stream and nb_io_dims == nb_seq_dims:
+        seq = []
+        for obj in arg[1]:
+            kwarg = {arg[0]: obj}
+            seq.append(TokSequence(**kwarg))
+            if not tokenizer.is_multi_voc:
+                seq[-1].ids_bpe_encoded = tokenizer._are_ids_bpe_encoded(
+                    seq[-1].ids
+                )
+    else:  # 1 subscript, one_token_stream and no multi-voc
+        kwarg = {arg[0]: arg[1]}
+        seq = TokSequence(**kwarg)
+        if not tokenizer.is_multi_voc:
+            seq.ids_bpe_encoded = tokenizer._are_ids_bpe_encoded(seq.ids)
+
+    # decode BPE and complete the output sequence(s) if requested
+    if tokenizer.has_bpe and decode_bpe:
+        tokenizer.decode_bpe(seq)
+    if complete_seq:
+        if isinstance(seq, TokSequence):
+            tokenizer.complete_sequence(seq)
+        else:
+            for seq_ in seq:
+                tokenizer.complete_sequence(seq_)
+
+    return seq
+
+
 def _in_as_seq(complete: bool = True, decode_bpe: bool = True):
     r"""Decorator creating if necessary and completing a TokSequence object before that the function is called.
     This decorator is made to be used by the :py:meth:`miditok.MIDITokenizer.tokens_to_midi` method.
@@ -47,50 +124,16 @@ def wrapper(*args, **kwargs):
             if not isinstance(seq, TokSequence) and not all(
                 isinstance(seq_, TokSequence) for seq_ in seq
             ):
-                try:
-                    arg = ("ids", convert_ids_tensors_to_list(seq))
-                except (AttributeError, ValueError, TypeError, IndexError):
-                    if isinstance(seq[0], str) or (
-                        isinstance(seq[0], str) and isinstance(seq[0][0], str)
-                    ):
-                        arg = ("tokens", seq)
-                    else:  # list of Event, very unlikely
-                        arg = ("events", seq)
-
-                # Deduce nb of subscript, if tokenizer is multi-voc or unique_track
-                nb_subscripts = nb_real_subscripts = 1
-                if not tokenizer.unique_track:
-                    nb_subscripts += 1
-                if tokenizer.is_multi_voc:
-                    nb_subscripts += 1
-                if isinstance(arg[1][0], list):
-                    nb_real_subscripts += 1
-                    if isinstance(arg[1][0][0], list):
-                        nb_real_subscripts += 1
-
-                if not tokenizer.unique_track and nb_subscripts == nb_real_subscripts:
-                    seq = []
-                    for obj in arg[1]:
-                        kwarg = {arg[0]: obj}
-                        seq.append(TokSequence(**kwarg))
-                        if not tokenizer.is_multi_voc:
-                            seq[-1].ids_bpe_encoded = tokenizer._are_ids_bpe_encoded(
-                                seq[-1].ids
-                            )
-                else:  # 1 subscript, unique_track and no multi-voc
-                    kwarg = {arg[0]: arg[1]}
-                    seq = TokSequence(**kwarg)
-                    if not tokenizer.is_multi_voc:
-                        seq.ids_bpe_encoded = tokenizer._are_ids_bpe_encoded(seq.ids)
-
-            if tokenizer.has_bpe and decode_bpe:
-                tokenizer.decode_bpe(seq)
-            if complete:
-                if isinstance(seq, TokSequence):
-                    tokenizer.complete_sequence(seq)
-                else:
-                    for seq_ in seq:
-                        tokenizer.complete_sequence(seq_)
+                seq = convert_sequence_to_tokseq(tokenizer, seq, complete, decode_bpe)
+            else:
+                if tokenizer.has_bpe and decode_bpe:
+                    tokenizer.decode_bpe(seq)
+                if complete:
+                    if isinstance(seq, TokSequence):
+                        tokenizer.complete_sequence(seq)
+                    else:
+                        for seq_ in seq:
+                            tokenizer.complete_sequence(seq_)
 
             args = list(args)
             args[1] = seq
@@ -117,9 +160,9 @@ class MIDITokenizer(ABC):
     r"""MIDI tokenizer base class, containing common methods and attributes for all tokenizers.
 
     :param tokenizer_config: the tokenizer's configuration, as a :class:`miditok.classes.TokenizerConfig` object.
-    :param unique_track: set to True if the tokenizer works only with a unique track.
-            Tokens will be saved as a single track. This applies to representations that natively handle
-            multiple tracks such as Octuple, resulting in a single "stream" of tokens for all tracks.
+    :param one_token_stream: give True if the tokenizer handle all the tracks of a MIDI as a single sequence of tokens.
+            Tokens will be saved as a single sequence. This applies to representations that natively handle
+            multiple tracks such as Octuple or REMIPlus, resulting in a single "stream" of tokens per MIDI.
             This attribute will be saved in config files of the tokenizer. (default: False)
     :param params: path to a tokenizer config file. This will override other arguments and
             load the tokenizer based on the config file. This is particularly useful if the
@@ -129,7 +172,7 @@ class MIDITokenizer(ABC):
     def __init__(
         self,
         tokenizer_config: TokenizerConfig = None,
-        unique_track: bool = False,
+        one_token_stream: bool = False,
         params: Union[str, Path] = None,
     ):
         # Initialize params
@@ -162,7 +205,7 @@ def __init__(
             assert (
                 0 < self.config.nb_velocities < 128
             ), "You must specify a nb_velocities between 1 and 127 (included)"
-            self.unique_track = unique_track
+            self.one_token_stream = one_token_stream
 
         # Tweak the tokenizer's configuration and / or attributes before creating the vocabulary
         # This method is intended to be overridden by inheriting tokenizer classes
@@ -457,7 +500,7 @@ def midi_to_tokens(
 
         :param midi: the MIDI object to convert.
         :param apply_bpe_if_possible: will apply BPE if the tokenizer's vocabulary was learned with.
-        :return: a :class:`miditok.TokSequence` if `tokenizer.unique_track` is true, else a list of
+        :return: a :class:`miditok.TokSequence` if `tokenizer.one_token_stream` is true, else a list of
                 :class:`miditok.TokSequence` objects.
         """
         # Check if the durations values have been calculated before for this time division
@@ -619,14 +662,14 @@ def tokens_to_midi(
         :param tokens: tokens to convert. Can be either a list of :class:`miditok.TokSequence`,
                 a Tensor (PyTorch and Tensorflow are supported), a numpy array or a Python list of ints.
                 The first dimension represents tracks, unless the tokenizer handle tracks altogether as a
-                single token sequence (e.g. Octuple, MuMIDI): tokenizer.unique_track == True.
+                single token sequence (e.g. Octuple, MuMIDI): tokenizer.one_token_stream == True.
         :param programs: programs of the tracks. If none is given, will default to piano, program 0. (default: None)
         :param output_path: path to save the file. (default: None)
         :param time_division: MIDI time division / resolution, in ticks/beat (of the MIDI to create).
         :return: the midi object (miditoolkit.MidiFile).
         """
         midi = MidiFile(ticks_per_beat=time_division)
-        # if self.unique_track:
+        # if self.one_token_stream:
         #    tokens = [tokens]
         for i, track_tokens in enumerate(tokens):
             if programs is not None:
@@ -1041,7 +1084,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.one_token_stream else [bytes_]
                 )
 
             # This doesn't seem to work, the trainer pre-processes the sequences, but then no word remains
@@ -1178,7 +1221,7 @@ def apply_bpe_to_dataset(
             sample = self.load_tokens(path)
             seq = (
                 TokSequence(ids=sample["ids"])
-                if self.unique_track
+                if self.one_token_stream
                 else [TokSequence(ids=track) for track in sample["ids"]]
             )
             self.apply_bpe(seq)
@@ -1460,7 +1503,7 @@ def save_params(
         }
         params = {
             "config": dict_config,
-            "unique_track": self.unique_track,
+            "one_token_stream": self.one_token_stream,
             "has_bpe": self.has_bpe,
             "tokenization": self.__class__.__name__,
             "miditok_version": CURRENT_VERSION_PACKAGE,
@@ -1532,6 +1575,9 @@ def _load_params(self, config_file_path: Union[str, Path]):
                     key = old_add_tokens_attr[key]
                 setattr(self.config, key, value)
                 continue
+            elif key == "unique_track":
+                # For config files <= v2.1.1 before the attribute is renamed
+                self.one_token_stream = value
 
             setattr(self, key, value)
 
@@ -1544,6 +1590,17 @@ def is_multi_voc(self) -> bool:
         """
         return isinstance(self._vocab_base, list)
 
+    @property
+    def io_format(self) -> Tuple[str]:
+        format_ = []
+        if not self.one_token_stream:
+            format_.append("I")
+        format_.append("T")
+        if self.is_multi_voc:
+            format_.append("C")
+
+        return tuple(d for d in format_)
+
     def __call__(self, obj: Any, *args, **kwargs):
         r"""Calling a tokenizer allows to directly convert a MIDI to tokens or the other way around.
         The method automatically detects MIDI and token objects, as well as paths and can directly load
@@ -1599,13 +1656,22 @@ def len(self) -> Union[int, List[int]]:
         return [len(v) for v in self.vocab] if self.is_multi_voc else len(self)
 
     def __repr__(self):
-        out_str = f"{self.len} tokens"
+        out_str = f"{self.len} tokens with {self.io_format} io format"
+
+        # one_token_stream / multi-voc
+        tmp = []
+        if self.one_token_stream:
+            tmp.append("one token stream")
         if self.is_multi_voc:
-            out_str += " (multi-voc)"
+            tmp.append("multi-voc")
+        if len(tmp) > 0:
+            out_str += f"({', '.join(tmp)})"
+
+        # BPE
         if self.has_bpe:
-            out_str += " with BPE"
+            out_str += ", with BPE"
         else:
-            out_str += " without BPE"
+            out_str += ", without BPE"
         return out_str
 
     def __getitem__(
diff --git a/miditok/tokenizations/mumidi.py b/miditok/tokenizations/mumidi.py
index c3c7f7f1..85d156ce 100644
--- a/miditok/tokenizations/mumidi.py
+++ b/miditok/tokenizations/mumidi.py
@@ -67,7 +67,7 @@ def __init__(
     def _tweak_config_before_creating_voc(self):
         self.config.use_rests = False
         self.config.use_time_signatures = False
-        # self.unique_track = True
+        # self.one_token_stream = True
 
         self.vocab_types_idx = {
             "Pitch": 0,
diff --git a/miditok/tokenizations/octuple.py b/miditok/tokenizations/octuple.py
index 5a40036e..4bd68f7d 100644
--- a/miditok/tokenizations/octuple.py
+++ b/miditok/tokenizations/octuple.py
@@ -44,7 +44,7 @@ class Octuple(MIDITokenizer):
     def _tweak_config_before_creating_voc(self):
         self.config.use_chords = False
         self.config.use_rests = False
-        self.unique_track = True
+        self.one_token_stream = True
 
         # used in place of positional encoding
         # This attribute might increase over tokenizations, if the tokenizer encounter longer MIDIs
diff --git a/miditok/tokenizations/remi_plus.py b/miditok/tokenizations/remi_plus.py
index 6777e0ab..5dc5cde7 100644
--- a/miditok/tokenizations/remi_plus.py
+++ b/miditok/tokenizations/remi_plus.py
@@ -52,7 +52,7 @@ def _tweak_config_before_creating_voc(self):
         self.config.use_programs = True
         # code handling rest decoding is writen, but not for detection (encoding)
         self.config.use_rests = False
-        # self.unique_track = True
+        # self.one_token_stream = True
 
         # In case the tokenizer has been created without specifying any config or params file path
         if "max_bar_embedding" not in self.config.additional_params:
diff --git a/miditok/tokenizations/tsd.py b/miditok/tokenizations/tsd.py
index c5fe0a56..eaab2d73 100644
--- a/miditok/tokenizations/tsd.py
+++ b/miditok/tokenizations/tsd.py
@@ -27,7 +27,7 @@ class TSD(MIDITokenizer):
     def _tweak_config_before_creating_voc(self):
         self.config.use_time_signatures = False
         if self.config.use_programs:
-            self.unique_track = True
+            self.one_token_stream = True
 
     def __notes_to_events(self, track: Instrument) -> List[Event]:
         r"""Converts notes of a track (``miditoolkit.Instrument``) into a sequence of `Event` objects.
@@ -201,7 +201,7 @@ def _midi_to_tokens(
         r"""Converts a preprocessed MIDI object to a sequence of tokens.
 
         :param midi: the MIDI objet to convert.
-        :return: a :class:`miditok.TokSequence` if `tokenizer.unique_track` is true, else a list of
+        :return: a :class:`miditok.TokSequence` if `tokenizer.one_token_stream` is true, else a list of
                 :class:`miditok.TokSequence` objects.
         """
         # Convert each track to tokens
@@ -210,7 +210,7 @@ def _midi_to_tokens(
         # Adds note tokens
         for track in midi.instruments:
             note_events = self.__notes_to_events(track)
-            if self.unique_track:
+            if self.one_token_stream:
                 all_events += note_events
             else:
                 all_events.append(note_events)
@@ -226,14 +226,14 @@ def _midi_to_tokens(
                         desc=tempo_change.tempo,
                     )
                 )
-            if self.unique_track:
+            if self.one_token_stream:
                 all_events += tempo_events
             else:
                 for i in range(len(all_events)):
                     all_events[i] += tempo_events
 
         # Add time events
-        if self.unique_track:
+        if self.one_token_stream:
             all_events.sort(key=lambda x: x.time)
             all_events = self.__add_time_note_events(all_events)
             tok_sequence = TokSequence(events=all_events)
@@ -278,8 +278,8 @@ def tokens_to_midi(
         :param time_division: MIDI time division / resolution, in ticks/beat (of the MIDI to create).
         :return: the midi object (:class:`miditoolkit.MidiFile`).
         """
-        # Unsqueeze tokens in case of unique_track
-        if self.unique_track:  # ie single token seq
+        # Unsqueeze tokens in case of one_token_stream
+        if self.one_token_stream:  # ie single token seq
             tokens = [tokens]
         for i in range(len(tokens)):
             tokens[i] = tokens[i].tokens
@@ -299,7 +299,7 @@ def tokens_to_midi(
         previous_note_end = 0
         for si, seq in enumerate(tokens):
             # Set track / sequence program if needed
-            if not self.unique_track:
+            if not self.one_token_stream:
                 current_tick = 0
                 previous_note_end = 0
                 if programs is not None:
diff --git a/tests/test_bpe.py b/tests/test_bpe.py
index 7fa56cd6..cf886b41 100644
--- a/tests/test_bpe.py
+++ b/tests/test_bpe.py
@@ -78,7 +78,7 @@ def test_bpe_conversion(
 
         for file_path in files:
             tokens = tokenizer(file_path, apply_bpe_if_possible=False)
-            if not tokenizer.unique_track:
+            if not tokenizer.one_token_stream:
                 tokens = tokens[0]
             to_tok = tokenizer._bytes_to_tokens(tokens.bytes)
             to_id = tokenizer._tokens_to_ids(to_tok)
@@ -112,7 +112,7 @@ def test_bpe_conversion(
 
         for tokenization, tokenizer in zip(tokenizations, tokenizers):
             tokens_no_bpe = tokenizer(deepcopy(midi), apply_bpe_if_possible=False)
-            if not tokenizer.unique_track:
+            if not tokenizer.one_token_stream:
                 tokens_no_bpe = tokens_no_bpe[0]
             tokens_bpe = deepcopy(tokens_no_bpe)  # with BPE
 
@@ -144,7 +144,7 @@ def test_bpe_conversion(
         for i, file_path in enumerate(tqdm(files, desc="Testing BPE batched")):
             # Reads the midi
             midi = MidiFile(file_path)
-            if not tokenizer.unique_track:
+            if not tokenizer.one_token_stream:
                 samples_no_bpe.append(tokenizer(midi, apply_bpe_if_possible=False)[0])
             else:
                 samples_no_bpe.append(tokenizer(midi, apply_bpe_if_possible=False))
diff --git a/tests/test_io_formats.py b/tests/test_io_formats.py
index db13e847..3a7ddea3 100644
--- a/tests/test_io_formats.py
+++ b/tests/test_io_formats.py
@@ -49,8 +49,8 @@ def encode_decode_and_check(tokenizer: miditok.MIDITokenizer, midi: MidiFile):
         if track.is_drum:
             track.program = 0  # need to be done before sorting tracks per program
     # Sort and merge tracks if needed
-    # MIDI produced with unique_track contains tracks with different orders
-    if tokenizer.unique_track:
+    # MIDI produced with one_token_stream contains tracks with different orders
+    if tokenizer.one_token_stream:
         miditok.utils.merge_same_program_tracks(midi_to_compare.instruments)
     # reduce the duration of notes to long
     for track in midi_to_compare.instruments:
@@ -65,14 +65,14 @@ def encode_decode_and_check(tokenizer: miditok.MIDITokenizer, midi: MidiFile):
 
     # Convert the midi to tokens, and keeps the ids (integers)
     tokens = tokenizer(midi_to_compare)
-    if tokenizer.unique_track:
+    if tokenizer.one_token_stream:
         tokens = tokens.ids
     else:
         tokens = [stream.ids for stream in tokens]
 
     # Convert back token ids to a MIDI object
     kwargs = {"time_division": midi.ticks_per_beat}
-    if not tokenizer.unique_track:
+    if not tokenizer.one_token_stream:
         kwargs["programs"] = miditok.utils.get_midi_programs(midi_to_compare)
     try:
         decoded_midi = tokenizer(tokens, **kwargs)
@@ -116,7 +116,7 @@ def test_io_formats():
             encode_decode_and_check(tokenizer, midi) or at_least_one_error
         )
 
-        # If TSD, also test in use_programs / unique_track mode
+        # If TSD, also test in use_programs / one_token_stream mode
         if tokenization == "TSD":
             tokenizer_config = miditok.TokenizerConfig(**TOKENIZER_PARAMS)
             tokenizer_config.use_programs = True
diff --git a/tests/test_methods.py b/tests/test_methods.py
index 4d9c27bf..d0dfb3fc 100644
--- a/tests/test_methods.py
+++ b/tests/test_methods.py
@@ -207,7 +207,7 @@ def test_data_augmentation():
                 original_tokens, original_programs = file["ids"], file["programs"]
 
             # Compare them
-            if tokenizer.unique_track:
+            if tokenizer.one_token_stream:
                 original_tokens, aug_tokens = [original_tokens], [aug_tokens]
             for original_track, aug_track, (_, is_drum) in zip(
                 original_tokens, aug_tokens, original_programs
@@ -222,7 +222,7 @@ def test_data_augmentation():
                             pitch_offset = offsets[0]
                             # no offset for drum pitches
                             if (
-                                tokenizer.unique_track
+                                tokenizer.one_token_stream
                                 and idx > 0
                                 and tokenizer[original_track[idx - 1]] == "Program_-1"
                             ):
diff --git a/tests/test_multitrack.py b/tests/test_multitrack.py
index e75c696e..87d96e3e 100644
--- a/tests/test_multitrack.py
+++ b/tests/test_multitrack.py
@@ -105,8 +105,8 @@ def test_multitrack_midi_to_tokens_to_midi(
                     )
 
             # Sort and merge tracks if needed
-            # MIDI produced with unique_track contains tracks with different orders
-            if tokenizer.unique_track:
+            # MIDI produced with one_token_stream contains tracks with different orders
+            if tokenizer.one_token_stream:
                 miditok.utils.merge_same_program_tracks(midi_to_compare.instruments)
             # reduce the duration of notes to long
             for track in midi_to_compare.instruments:
@@ -134,7 +134,7 @@ def test_multitrack_midi_to_tokens_to_midi(
 
             # Checks types and values conformity following the rules
             tokens_types = tokenizer.tokens_errors(
-                tokens[0] if not tokenizer.unique_track else tokens
+                tokens[0] if not tokenizer.one_token_stream else tokens
             )
             if tokens_types != 0.0:
                 print(
diff --git a/tests/test_one_track.py b/tests/test_one_track.py
index 45a97d27..923c7de9 100644
--- a/tests/test_one_track.py
+++ b/tests/test_one_track.py
@@ -84,7 +84,7 @@ def test_one_track_midi_to_tokens_to_midi(
 
             # Convert the track in tokens
             tokens = tokenizer(midi)
-            if not tokenizer.unique_track:
+            if not tokenizer.one_token_stream:
                 tokens = tokens[0]
 
             # Checks types and values conformity following the rules
@@ -95,7 +95,7 @@ def test_one_track_midi_to_tokens_to_midi(
                 )
 
             # Convert back tokens into a track object
-            if not tokenizer.unique_track:
+            if not tokenizer.one_token_stream:
                 tokens = [tokens]
             new_midi = tokenizer.tokens_to_midi(
                 tokens, time_division=midi.ticks_per_beat