Upload tokenization_song2midi.py

tokenization_song2midi.py

@@ -1,6 +1,6 @@
 import os
 from pathlib import Path
-from typing import Union
+from typing import Union, TYPE_CHECKING
 
 from transformers import BatchEncoding, PythonBackend
 from transformers.tokenization_utils_base import TruncationStrategy
@@ -8,7 +8,9 @@ from transformers.utils.generic import PaddingStrategy, TensorType
 
 try:
     from miditok import PerTok, TokSequence
-    from symusic import Score
+    import symusic
+    if TYPE_CHECKING:
+        from symusic.types import Score
 except ImportError:
     raise ImportError(
         "The `miditok` library is required for processing MIDI files. "
@@ -29,6 +31,11 @@ class Song2MIDIPerTokTokenizer(PythonBackend):
         **kwargs,
     ):
         self._tokenizer = PerTok(params=vocab_file)
+        
+        # PerTok as of miditok version 3.0.6.post1 does not load position token locations from the vocab file.
+        # use_position_toks workaround
+        if self._tokenizer.use_position_toks and not getattr(self._tokenizer, "position_locations", None):
+            self._tokenizer.position_locations = self._tokenizer._create_position_tok_locations()
 
         self._decoder = {value: key for key, value in self._tokenizer.vocab.items()}
 
@@ -49,8 +56,8 @@ class Song2MIDIPerTokTokenizer(PythonBackend):
 
     def _encode_plus(
         self,
-        text: Union["Score", Path, bytes, list[Union["Score", Path, bytes]], list[int]],
-        text_pair: Union["Score", Path, list[Union["Score", Path]], list[int], None] = None,
+        text: Union[Score, Path, bytes, list[Union[Score, Path, bytes]], list[int]],
+        text_pair: Union[Score, Path, list[Union[Score, Path]], list[int], None] = None,
         add_special_tokens: bool = True,
         padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
         truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
@@ -132,7 +139,7 @@ class Song2MIDIPerTokTokenizer(PythonBackend):
                 return []
             if isinstance(midi_input, (str, Path, Score, bytes)):
                 if isinstance(midi_input, bytes):
-                    midi_input = Score.from_midi(midi_input)
+                    midi_input = symusic.Score.from_midi(midi_input)
                 return self._tokenizer.encode(midi_input).ids
             if isinstance(midi_input, (list, tuple)) and midi_input:
                 if isinstance(midi_input[0], int):
@@ -186,7 +193,20 @@ class Song2MIDIPerTokTokenizer(PythonBackend):
             ]
 
         return " ".join(tokens)
+    
+    def decode_score(
+        self,
+        token_ids: int | list[int],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool | None = None,
+        **kwargs,
+    ) -> Score:
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
 
+        tok_sequence = TokSequence(ids=token_ids, are_ids_encoded=True)
+        return self._tokenizer.decode(tok_sequence)
+    
     def save_vocabulary(
         self, save_directory: str, filename_prefix: str | None = None
     ) -> tuple[str, ...]: