717 lines
32 KiB
Python
717 lines
32 KiB
Python
# coding=utf-8
|
|
# Copyright 2023 The Pop2Piano Authors and The HuggingFace Inc. team.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Tokenization class for Pop2Piano."""
|
|
|
|
import json
|
|
import os
|
|
from typing import List, Optional, Tuple, Union
|
|
|
|
import numpy as np
|
|
|
|
from ...feature_extraction_utils import BatchFeature
|
|
from ...tokenization_utils import AddedToken, BatchEncoding, PaddingStrategy, PreTrainedTokenizer, TruncationStrategy
|
|
from ...utils import TensorType, is_pretty_midi_available, logging, requires_backends, to_numpy
|
|
|
|
|
|
if is_pretty_midi_available():
|
|
import pretty_midi
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
|
|
VOCAB_FILES_NAMES = {
|
|
"vocab": "vocab.json",
|
|
}
|
|
|
|
|
|
def token_time_to_note(number, cutoff_time_idx, current_idx):
|
|
current_idx += number
|
|
if cutoff_time_idx is not None:
|
|
current_idx = min(current_idx, cutoff_time_idx)
|
|
|
|
return current_idx
|
|
|
|
|
|
def token_note_to_note(number, current_velocity, default_velocity, note_onsets_ready, current_idx, notes):
|
|
if note_onsets_ready[number] is not None:
|
|
# offset with onset
|
|
onset_idx = note_onsets_ready[number]
|
|
if onset_idx < current_idx:
|
|
# Time shift after previous note_on
|
|
offset_idx = current_idx
|
|
notes.append([onset_idx, offset_idx, number, default_velocity])
|
|
onsets_ready = None if current_velocity == 0 else current_idx
|
|
note_onsets_ready[number] = onsets_ready
|
|
else:
|
|
note_onsets_ready[number] = current_idx
|
|
return notes
|
|
|
|
|
|
class Pop2PianoTokenizer(PreTrainedTokenizer):
|
|
"""
|
|
Constructs a Pop2Piano tokenizer. This tokenizer does not require training.
|
|
|
|
This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
|
|
this superclass for more information regarding those methods.
|
|
|
|
Args:
|
|
vocab (`str`):
|
|
Path to the vocab file which contains the vocabulary.
|
|
default_velocity (`int`, *optional*, defaults to 77):
|
|
Determines the default velocity to be used while creating midi Notes.
|
|
num_bars (`int`, *optional*, defaults to 2):
|
|
Determines cutoff_time_idx in for each token.
|
|
unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"-1"`):
|
|
The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
|
|
token instead.
|
|
eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to 1):
|
|
The end of sequence token.
|
|
pad_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to 0):
|
|
A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by
|
|
attention mechanisms or loss computation.
|
|
bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to 2):
|
|
The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
|
|
"""
|
|
|
|
model_input_names = ["token_ids", "attention_mask"]
|
|
vocab_files_names = VOCAB_FILES_NAMES
|
|
|
|
def __init__(
|
|
self,
|
|
vocab,
|
|
default_velocity=77,
|
|
num_bars=2,
|
|
unk_token="-1",
|
|
eos_token="1",
|
|
pad_token="0",
|
|
bos_token="2",
|
|
**kwargs,
|
|
):
|
|
unk_token = AddedToken(unk_token, lstrip=False, rstrip=False) if isinstance(unk_token, str) else unk_token
|
|
eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
|
|
pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
|
|
bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
|
|
|
|
self.default_velocity = default_velocity
|
|
self.num_bars = num_bars
|
|
|
|
# Load the vocab
|
|
with open(vocab, "rb") as file:
|
|
self.encoder = json.load(file)
|
|
|
|
# create mappings for encoder
|
|
self.decoder = {v: k for k, v in self.encoder.items()}
|
|
|
|
super().__init__(
|
|
unk_token=unk_token,
|
|
eos_token=eos_token,
|
|
pad_token=pad_token,
|
|
bos_token=bos_token,
|
|
**kwargs,
|
|
)
|
|
|
|
@property
|
|
def vocab_size(self):
|
|
"""Returns the vocabulary size of the tokenizer."""
|
|
return len(self.encoder)
|
|
|
|
def get_vocab(self):
|
|
"""Returns the vocabulary of the tokenizer."""
|
|
return dict(self.encoder, **self.added_tokens_encoder)
|
|
|
|
def _convert_id_to_token(self, token_id: int) -> list:
|
|
"""
|
|
Decodes the token ids generated by the transformer into notes.
|
|
|
|
Args:
|
|
token_id (`int`):
|
|
This denotes the ids generated by the transformers to be converted to Midi tokens.
|
|
|
|
Returns:
|
|
`List`: A list consists of token_type (`str`) and value (`int`).
|
|
"""
|
|
|
|
token_type_value = self.decoder.get(token_id, f"{self.unk_token}_TOKEN_TIME")
|
|
token_type_value = token_type_value.split("_")
|
|
token_type, value = "_".join(token_type_value[1:]), int(token_type_value[0])
|
|
|
|
return [token_type, value]
|
|
|
|
def _convert_token_to_id(self, token, token_type="TOKEN_TIME") -> int:
|
|
"""
|
|
Encodes the Midi tokens to transformer generated token ids.
|
|
|
|
Args:
|
|
token (`int`):
|
|
This denotes the token value.
|
|
token_type (`str`):
|
|
This denotes the type of the token. There are four types of midi tokens such as "TOKEN_TIME",
|
|
"TOKEN_VELOCITY", "TOKEN_NOTE" and "TOKEN_SPECIAL".
|
|
|
|
Returns:
|
|
`int`: returns the id of the token.
|
|
"""
|
|
return self.encoder.get(f"{token}_{token_type}", int(self.unk_token))
|
|
|
|
def relative_batch_tokens_ids_to_notes(
|
|
self,
|
|
tokens: np.ndarray,
|
|
beat_offset_idx: int,
|
|
bars_per_batch: int,
|
|
cutoff_time_idx: int,
|
|
):
|
|
"""
|
|
Converts relative tokens to notes which are then used to generate pretty midi object.
|
|
|
|
Args:
|
|
tokens (`numpy.ndarray`):
|
|
Tokens to be converted to notes.
|
|
beat_offset_idx (`int`):
|
|
Denotes beat offset index for each note in generated Midi.
|
|
bars_per_batch (`int`):
|
|
A parameter to control the Midi output generation.
|
|
cutoff_time_idx (`int`):
|
|
Denotes the cutoff time index for each note in generated Midi.
|
|
"""
|
|
|
|
notes = None
|
|
|
|
for index in range(len(tokens)):
|
|
_tokens = tokens[index]
|
|
_start_idx = beat_offset_idx + index * bars_per_batch * 4
|
|
_cutoff_time_idx = cutoff_time_idx + _start_idx
|
|
_notes = self.relative_tokens_ids_to_notes(
|
|
_tokens,
|
|
start_idx=_start_idx,
|
|
cutoff_time_idx=_cutoff_time_idx,
|
|
)
|
|
|
|
if len(_notes) == 0:
|
|
pass
|
|
elif notes is None:
|
|
notes = _notes
|
|
else:
|
|
notes = np.concatenate((notes, _notes), axis=0)
|
|
|
|
if notes is None:
|
|
return []
|
|
return notes
|
|
|
|
def relative_batch_tokens_ids_to_midi(
|
|
self,
|
|
tokens: np.ndarray,
|
|
beatstep: np.ndarray,
|
|
beat_offset_idx: int = 0,
|
|
bars_per_batch: int = 2,
|
|
cutoff_time_idx: int = 12,
|
|
):
|
|
"""
|
|
Converts tokens to Midi. This method calls `relative_batch_tokens_ids_to_notes` method to convert batch tokens
|
|
to notes then uses `notes_to_midi` method to convert them to Midi.
|
|
|
|
Args:
|
|
tokens (`numpy.ndarray`):
|
|
Denotes tokens which alongside beatstep will be converted to Midi.
|
|
beatstep (`np.ndarray`):
|
|
We get beatstep from feature extractor which is also used to get Midi.
|
|
beat_offset_idx (`int`, *optional*, defaults to 0):
|
|
Denotes beat offset index for each note in generated Midi.
|
|
bars_per_batch (`int`, *optional*, defaults to 2):
|
|
A parameter to control the Midi output generation.
|
|
cutoff_time_idx (`int`, *optional*, defaults to 12):
|
|
Denotes the cutoff time index for each note in generated Midi.
|
|
"""
|
|
beat_offset_idx = 0 if beat_offset_idx is None else beat_offset_idx
|
|
notes = self.relative_batch_tokens_ids_to_notes(
|
|
tokens=tokens,
|
|
beat_offset_idx=beat_offset_idx,
|
|
bars_per_batch=bars_per_batch,
|
|
cutoff_time_idx=cutoff_time_idx,
|
|
)
|
|
midi = self.notes_to_midi(notes, beatstep, offset_sec=beatstep[beat_offset_idx])
|
|
return midi
|
|
|
|
# Taken from the original code
|
|
# Please see https://github.com/sweetcocoa/pop2piano/blob/fac11e8dcfc73487513f4588e8d0c22a22f2fdc5/midi_tokenizer.py#L257
|
|
def relative_tokens_ids_to_notes(self, tokens: np.ndarray, start_idx: float, cutoff_time_idx: float = None):
|
|
"""
|
|
Converts relative tokens to notes which will then be used to create Pretty Midi objects.
|
|
|
|
Args:
|
|
tokens (`numpy.ndarray`):
|
|
Relative Tokens which will be converted to notes.
|
|
start_idx (`float`):
|
|
A parameter which denotes the starting index.
|
|
cutoff_time_idx (`float`, *optional*):
|
|
A parameter used while converting tokens to notes.
|
|
"""
|
|
words = [self._convert_id_to_token(token) for token in tokens]
|
|
|
|
current_idx = start_idx
|
|
current_velocity = 0
|
|
note_onsets_ready = [None for i in range(sum([k.endswith("NOTE") for k in self.encoder.keys()]) + 1)]
|
|
notes = []
|
|
for token_type, number in words:
|
|
if token_type == "TOKEN_SPECIAL":
|
|
if number == 1:
|
|
break
|
|
elif token_type == "TOKEN_TIME":
|
|
current_idx = token_time_to_note(
|
|
number=number, cutoff_time_idx=cutoff_time_idx, current_idx=current_idx
|
|
)
|
|
elif token_type == "TOKEN_VELOCITY":
|
|
current_velocity = number
|
|
|
|
elif token_type == "TOKEN_NOTE":
|
|
notes = token_note_to_note(
|
|
number=number,
|
|
current_velocity=current_velocity,
|
|
default_velocity=self.default_velocity,
|
|
note_onsets_ready=note_onsets_ready,
|
|
current_idx=current_idx,
|
|
notes=notes,
|
|
)
|
|
else:
|
|
raise ValueError("Token type not understood!")
|
|
|
|
for pitch, note_onset in enumerate(note_onsets_ready):
|
|
# force offset if no offset for each pitch
|
|
if note_onset is not None:
|
|
if cutoff_time_idx is None:
|
|
cutoff = note_onset + 1
|
|
else:
|
|
cutoff = max(cutoff_time_idx, note_onset + 1)
|
|
|
|
offset_idx = max(current_idx, cutoff)
|
|
notes.append([note_onset, offset_idx, pitch, self.default_velocity])
|
|
|
|
if len(notes) == 0:
|
|
return []
|
|
else:
|
|
notes = np.array(notes)
|
|
note_order = notes[:, 0] * 128 + notes[:, 1]
|
|
notes = notes[note_order.argsort()]
|
|
return notes
|
|
|
|
def notes_to_midi(self, notes: np.ndarray, beatstep: np.ndarray, offset_sec: int = 0.0):
|
|
"""
|
|
Converts notes to Midi.
|
|
|
|
Args:
|
|
notes (`numpy.ndarray`):
|
|
This is used to create Pretty Midi objects.
|
|
beatstep (`numpy.ndarray`):
|
|
This is the extrapolated beatstep that we get from feature extractor.
|
|
offset_sec (`int`, *optional*, defaults to 0.0):
|
|
This represents the offset seconds which is used while creating each Pretty Midi Note.
|
|
"""
|
|
|
|
requires_backends(self, ["pretty_midi"])
|
|
|
|
new_pm = pretty_midi.PrettyMIDI(resolution=384, initial_tempo=120.0)
|
|
new_inst = pretty_midi.Instrument(program=0)
|
|
new_notes = []
|
|
|
|
for onset_idx, offset_idx, pitch, velocity in notes:
|
|
new_note = pretty_midi.Note(
|
|
velocity=velocity,
|
|
pitch=pitch,
|
|
start=beatstep[onset_idx] - offset_sec,
|
|
end=beatstep[offset_idx] - offset_sec,
|
|
)
|
|
new_notes.append(new_note)
|
|
new_inst.notes = new_notes
|
|
new_pm.instruments.append(new_inst)
|
|
new_pm.remove_invalid_notes()
|
|
return new_pm
|
|
|
|
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
|
|
"""
|
|
Saves the tokenizer's vocabulary dictionary to the provided save_directory.
|
|
|
|
Args:
|
|
save_directory (`str`):
|
|
A path to the directory where to saved. It will be created if it doesn't exist.
|
|
filename_prefix (`Optional[str]`, *optional*):
|
|
A prefix to add to the names of the files saved by the tokenizer.
|
|
"""
|
|
if not os.path.isdir(save_directory):
|
|
logger.error(f"Vocabulary path ({save_directory}) should be a directory")
|
|
return
|
|
|
|
# Save the encoder.
|
|
out_vocab_file = os.path.join(
|
|
save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab"]
|
|
)
|
|
with open(out_vocab_file, "w") as file:
|
|
file.write(json.dumps(self.encoder))
|
|
|
|
return (out_vocab_file,)
|
|
|
|
def encode_plus(
|
|
self,
|
|
notes: Union[np.ndarray, List[pretty_midi.Note]],
|
|
truncation_strategy: Optional[TruncationStrategy] = None,
|
|
max_length: Optional[int] = None,
|
|
**kwargs,
|
|
) -> BatchEncoding:
|
|
r"""
|
|
This is the `encode_plus` method for `Pop2PianoTokenizer`. It converts the midi notes to the transformer
|
|
generated token ids. It only works on a single batch, to process multiple batches please use
|
|
`batch_encode_plus` or `__call__` method.
|
|
|
|
Args:
|
|
notes (`numpy.ndarray` of shape `[sequence_length, 4]` or `list` of `pretty_midi.Note` objects):
|
|
This represents the midi notes. If `notes` is a `numpy.ndarray`:
|
|
- Each sequence must have 4 values, they are `onset idx`, `offset idx`, `pitch` and `velocity`.
|
|
If `notes` is a `list` containing `pretty_midi.Note` objects:
|
|
- Each sequence must have 4 attributes, they are `start`, `end`, `pitch` and `velocity`.
|
|
truncation_strategy ([`~tokenization_utils_base.TruncationStrategy`], *optional*):
|
|
Indicates the truncation strategy that is going to be used during truncation.
|
|
max_length (`int`, *optional*):
|
|
Maximum length of the returned list and optionally padding length (see above).
|
|
|
|
Returns:
|
|
`BatchEncoding` containing the tokens ids.
|
|
"""
|
|
|
|
requires_backends(self, ["pretty_midi"])
|
|
|
|
# check if notes is a pretty_midi object or not, if yes then extract the attributes and put them into a numpy
|
|
# array.
|
|
if isinstance(notes[0], pretty_midi.Note):
|
|
notes = np.array(
|
|
[[each_note.start, each_note.end, each_note.pitch, each_note.velocity] for each_note in notes]
|
|
).reshape(-1, 4)
|
|
|
|
# to round up all the values to the closest int values.
|
|
notes = np.round(notes).astype(np.int32)
|
|
max_time_idx = notes[:, :2].max()
|
|
|
|
times = [[] for i in range((max_time_idx + 1))]
|
|
for onset, offset, pitch, velocity in notes:
|
|
times[onset].append([pitch, velocity])
|
|
times[offset].append([pitch, 0])
|
|
|
|
tokens = []
|
|
current_velocity = 0
|
|
for i, time in enumerate(times):
|
|
if len(time) == 0:
|
|
continue
|
|
tokens.append(self._convert_token_to_id(i, "TOKEN_TIME"))
|
|
for pitch, velocity in time:
|
|
velocity = int(velocity > 0)
|
|
if current_velocity != velocity:
|
|
current_velocity = velocity
|
|
tokens.append(self._convert_token_to_id(velocity, "TOKEN_VELOCITY"))
|
|
tokens.append(self._convert_token_to_id(pitch, "TOKEN_NOTE"))
|
|
|
|
total_len = len(tokens)
|
|
|
|
# truncation
|
|
if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE and max_length and total_len > max_length:
|
|
tokens, _, _ = self.truncate_sequences(
|
|
ids=tokens,
|
|
num_tokens_to_remove=total_len - max_length,
|
|
truncation_strategy=truncation_strategy,
|
|
**kwargs,
|
|
)
|
|
|
|
return BatchEncoding({"token_ids": tokens})
|
|
|
|
def batch_encode_plus(
|
|
self,
|
|
notes: Union[np.ndarray, List[pretty_midi.Note]],
|
|
truncation_strategy: Optional[TruncationStrategy] = None,
|
|
max_length: Optional[int] = None,
|
|
**kwargs,
|
|
) -> BatchEncoding:
|
|
r"""
|
|
This is the `batch_encode_plus` method for `Pop2PianoTokenizer`. It converts the midi notes to the transformer
|
|
generated token ids. It works on multiple batches by calling `encode_plus` multiple times in a loop.
|
|
|
|
Args:
|
|
notes (`numpy.ndarray` of shape `[batch_size, sequence_length, 4]` or `list` of `pretty_midi.Note` objects):
|
|
This represents the midi notes. If `notes` is a `numpy.ndarray`:
|
|
- Each sequence must have 4 values, they are `onset idx`, `offset idx`, `pitch` and `velocity`.
|
|
If `notes` is a `list` containing `pretty_midi.Note` objects:
|
|
- Each sequence must have 4 attributes, they are `start`, `end`, `pitch` and `velocity`.
|
|
truncation_strategy ([`~tokenization_utils_base.TruncationStrategy`], *optional*):
|
|
Indicates the truncation strategy that is going to be used during truncation.
|
|
max_length (`int`, *optional*):
|
|
Maximum length of the returned list and optionally padding length (see above).
|
|
|
|
Returns:
|
|
`BatchEncoding` containing the tokens ids.
|
|
"""
|
|
|
|
encoded_batch_token_ids = []
|
|
for i in range(len(notes)):
|
|
encoded_batch_token_ids.append(
|
|
self.encode_plus(
|
|
notes[i],
|
|
truncation_strategy=truncation_strategy,
|
|
max_length=max_length,
|
|
**kwargs,
|
|
)["token_ids"]
|
|
)
|
|
|
|
return BatchEncoding({"token_ids": encoded_batch_token_ids})
|
|
|
|
def __call__(
|
|
self,
|
|
notes: Union[
|
|
np.ndarray,
|
|
List[pretty_midi.Note],
|
|
List[List[pretty_midi.Note]],
|
|
],
|
|
padding: Union[bool, str, PaddingStrategy] = False,
|
|
truncation: Union[bool, str, TruncationStrategy] = None,
|
|
max_length: Optional[int] = None,
|
|
pad_to_multiple_of: Optional[int] = None,
|
|
return_attention_mask: Optional[bool] = None,
|
|
return_tensors: Optional[Union[str, TensorType]] = None,
|
|
verbose: bool = True,
|
|
**kwargs,
|
|
) -> BatchEncoding:
|
|
r"""
|
|
This is the `__call__` method for `Pop2PianoTokenizer`. It converts the midi notes to the transformer generated
|
|
token ids.
|
|
|
|
Args:
|
|
notes (`numpy.ndarray` of shape `[batch_size, max_sequence_length, 4]` or `list` of `pretty_midi.Note` objects):
|
|
This represents the midi notes.
|
|
|
|
If `notes` is a `numpy.ndarray`:
|
|
- Each sequence must have 4 values, they are `onset idx`, `offset idx`, `pitch` and `velocity`.
|
|
If `notes` is a `list` containing `pretty_midi.Note` objects:
|
|
- Each sequence must have 4 attributes, they are `start`, `end`, `pitch` and `velocity`.
|
|
padding (`bool`, `str` or [`~file_utils.PaddingStrategy`], *optional*, defaults to `False`):
|
|
Activates and controls padding. Accepts the following values:
|
|
|
|
- `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
|
|
sequence if provided).
|
|
- `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
|
|
acceptable input length for the model if that argument is not provided.
|
|
- `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
|
|
lengths).
|
|
truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):
|
|
Activates and controls truncation. Accepts the following values:
|
|
|
|
- `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or
|
|
to the maximum acceptable input length for the model if that argument is not provided. This will
|
|
truncate token by token, removing a token from the longest sequence in the pair if a pair of
|
|
sequences (or a batch of pairs) is provided.
|
|
- `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the
|
|
maximum acceptable input length for the model if that argument is not provided. This will only
|
|
truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
|
|
- `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the
|
|
maximum acceptable input length for the model if that argument is not provided. This will only
|
|
truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
|
|
- `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths
|
|
greater than the model maximum admissible input size).
|
|
max_length (`int`, *optional*):
|
|
Controls the maximum length to use by one of the truncation/padding parameters. If left unset or set to
|
|
`None`, this will use the predefined model maximum length if a maximum length is required by one of the
|
|
truncation/padding parameters. If the model has no specific maximum input length (like XLNet)
|
|
truncation/padding to a maximum length will be deactivated.
|
|
pad_to_multiple_of (`int`, *optional*):
|
|
If set will pad the sequence to a multiple of the provided value. This is especially useful to enable
|
|
the use of Tensor Cores on NVIDIA hardware with compute capability `>= 7.5` (Volta).
|
|
return_attention_mask (`bool`, *optional*):
|
|
Whether to return the attention mask. If left to the default, will return the attention mask according
|
|
to the specific tokenizer's default, defined by the `return_outputs` attribute.
|
|
|
|
[What are attention masks?](../glossary#attention-mask)
|
|
return_tensors (`str` or [`~file_utils.TensorType`], *optional*):
|
|
If set, will return tensors instead of list of python integers. Acceptable values are:
|
|
|
|
- `'tf'`: Return TensorFlow `tf.constant` objects.
|
|
- `'pt'`: Return PyTorch `torch.Tensor` objects.
|
|
- `'np'`: Return Numpy `np.ndarray` objects.
|
|
verbose (`bool`, *optional*, defaults to `True`):
|
|
Whether or not to print more information and warnings.
|
|
|
|
Returns:
|
|
`BatchEncoding` containing the token_ids.
|
|
"""
|
|
|
|
# check if it is batched or not
|
|
# it is batched if its a list containing a list of `pretty_midi.Notes` where the outer list contains all the
|
|
# batches and the inner list contains all Notes for a single batch. Otherwise if np.ndarray is passed it will be
|
|
# considered batched if it has shape of `[batch_size, seqence_length, 4]` or ndim=3.
|
|
is_batched = notes.ndim == 3 if isinstance(notes, np.ndarray) else isinstance(notes[0], list)
|
|
|
|
# get the truncation and padding strategy
|
|
padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(
|
|
padding=padding,
|
|
truncation=truncation,
|
|
max_length=max_length,
|
|
pad_to_multiple_of=pad_to_multiple_of,
|
|
verbose=verbose,
|
|
**kwargs,
|
|
)
|
|
|
|
if is_batched:
|
|
# If the user has not explicitly mentioned `return_attention_mask` as False, we change it to True
|
|
return_attention_mask = True if return_attention_mask is None else return_attention_mask
|
|
token_ids = self.batch_encode_plus(
|
|
notes=notes,
|
|
truncation_strategy=truncation_strategy,
|
|
max_length=max_length,
|
|
**kwargs,
|
|
)
|
|
else:
|
|
token_ids = self.encode_plus(
|
|
notes=notes,
|
|
truncation_strategy=truncation_strategy,
|
|
max_length=max_length,
|
|
**kwargs,
|
|
)
|
|
|
|
# since we already have truncated sequnences we are just left to do padding
|
|
token_ids = self.pad(
|
|
token_ids,
|
|
padding=padding_strategy,
|
|
max_length=max_length,
|
|
pad_to_multiple_of=pad_to_multiple_of,
|
|
return_attention_mask=return_attention_mask,
|
|
return_tensors=return_tensors,
|
|
verbose=verbose,
|
|
)
|
|
|
|
return token_ids
|
|
|
|
def batch_decode(
|
|
self,
|
|
token_ids,
|
|
feature_extractor_output: BatchFeature,
|
|
return_midi: bool = True,
|
|
):
|
|
r"""
|
|
This is the `batch_decode` method for `Pop2PianoTokenizer`. It converts the token_ids generated by the
|
|
transformer to midi_notes and returns them.
|
|
|
|
Args:
|
|
token_ids (`Union[np.ndarray, torch.Tensor, tf.Tensor]`):
|
|
Output token_ids of `Pop2PianoConditionalGeneration` model.
|
|
feature_extractor_output (`BatchFeature`):
|
|
Denotes the output of `Pop2PianoFeatureExtractor.__call__`. It must contain `"beatstep"` and
|
|
`"extrapolated_beatstep"`. Also `"attention_mask_beatsteps"` and
|
|
`"attention_mask_extrapolated_beatstep"`
|
|
should be present if they were returned by the feature extractor.
|
|
return_midi (`bool`, *optional*, defaults to `True`):
|
|
Whether to return midi object or not.
|
|
Returns:
|
|
If `return_midi` is True:
|
|
- `BatchEncoding` containing both `notes` and `pretty_midi.pretty_midi.PrettyMIDI` objects.
|
|
If `return_midi` is False:
|
|
- `BatchEncoding` containing `notes`.
|
|
"""
|
|
|
|
# check if they have attention_masks(attention_mask, attention_mask_beatsteps, attention_mask_extrapolated_beatstep) or not
|
|
attention_masks_present = bool(
|
|
hasattr(feature_extractor_output, "attention_mask")
|
|
and hasattr(feature_extractor_output, "attention_mask_beatsteps")
|
|
and hasattr(feature_extractor_output, "attention_mask_extrapolated_beatstep")
|
|
)
|
|
|
|
# if we are processing batched inputs then we must need attention_masks
|
|
if not attention_masks_present and feature_extractor_output["beatsteps"].shape[0] > 1:
|
|
raise ValueError(
|
|
"attention_mask, attention_mask_beatsteps and attention_mask_extrapolated_beatstep must be present "
|
|
"for batched inputs! But one of them were not present."
|
|
)
|
|
|
|
# check for length mismatch between inputs_embeds, beatsteps and extrapolated_beatstep
|
|
if attention_masks_present:
|
|
# since we know about the number of examples in token_ids from attention_mask
|
|
if (
|
|
sum(feature_extractor_output["attention_mask"][:, 0] == 0)
|
|
!= feature_extractor_output["beatsteps"].shape[0]
|
|
or feature_extractor_output["beatsteps"].shape[0]
|
|
!= feature_extractor_output["extrapolated_beatstep"].shape[0]
|
|
):
|
|
raise ValueError(
|
|
"Length mistamtch between token_ids, beatsteps and extrapolated_beatstep! Found "
|
|
f"token_ids length - {token_ids.shape[0]}, beatsteps shape - {feature_extractor_output['beatsteps'].shape[0]} "
|
|
f"and extrapolated_beatsteps shape - {feature_extractor_output['extrapolated_beatstep'].shape[0]}"
|
|
)
|
|
if feature_extractor_output["attention_mask"].shape[0] != token_ids.shape[0]:
|
|
raise ValueError(
|
|
f"Found attention_mask of length - {feature_extractor_output['attention_mask'].shape[0]} but token_ids of length - {token_ids.shape[0]}"
|
|
)
|
|
else:
|
|
# if there is no attention mask present then it's surely a single example
|
|
if (
|
|
feature_extractor_output["beatsteps"].shape[0] != 1
|
|
or feature_extractor_output["extrapolated_beatstep"].shape[0] != 1
|
|
):
|
|
raise ValueError(
|
|
"Length mistamtch of beatsteps and extrapolated_beatstep! Since attention_mask is not present the number of examples must be 1, "
|
|
f"But found beatsteps length - {feature_extractor_output['beatsteps'].shape[0]}, extrapolated_beatsteps length - {feature_extractor_output['extrapolated_beatstep'].shape[0]}."
|
|
)
|
|
|
|
if attention_masks_present:
|
|
# check for zeros(since token_ids are seperated by zero arrays)
|
|
batch_idx = np.where(feature_extractor_output["attention_mask"][:, 0] == 0)[0]
|
|
else:
|
|
batch_idx = [token_ids.shape[0]]
|
|
|
|
notes_list = []
|
|
pretty_midi_objects_list = []
|
|
start_idx = 0
|
|
for index, end_idx in enumerate(batch_idx):
|
|
each_tokens_ids = token_ids[start_idx:end_idx]
|
|
# check where the whole example ended by searching for eos_token_id and getting the upper bound
|
|
each_tokens_ids = each_tokens_ids[:, : np.max(np.where(each_tokens_ids == int(self.eos_token))[1]) + 1]
|
|
beatsteps = feature_extractor_output["beatsteps"][index]
|
|
extrapolated_beatstep = feature_extractor_output["extrapolated_beatstep"][index]
|
|
|
|
# if attention mask is present then mask out real array/tensor
|
|
if attention_masks_present:
|
|
attention_mask_beatsteps = feature_extractor_output["attention_mask_beatsteps"][index]
|
|
attention_mask_extrapolated_beatstep = feature_extractor_output[
|
|
"attention_mask_extrapolated_beatstep"
|
|
][index]
|
|
beatsteps = beatsteps[: np.max(np.where(attention_mask_beatsteps == 1)[0]) + 1]
|
|
extrapolated_beatstep = extrapolated_beatstep[
|
|
: np.max(np.where(attention_mask_extrapolated_beatstep == 1)[0]) + 1
|
|
]
|
|
|
|
each_tokens_ids = to_numpy(each_tokens_ids)
|
|
beatsteps = to_numpy(beatsteps)
|
|
extrapolated_beatstep = to_numpy(extrapolated_beatstep)
|
|
|
|
pretty_midi_object = self.relative_batch_tokens_ids_to_midi(
|
|
tokens=each_tokens_ids,
|
|
beatstep=extrapolated_beatstep,
|
|
bars_per_batch=self.num_bars,
|
|
cutoff_time_idx=(self.num_bars + 1) * 4,
|
|
)
|
|
|
|
for note in pretty_midi_object.instruments[0].notes:
|
|
note.start += beatsteps[0]
|
|
note.end += beatsteps[0]
|
|
notes_list.append(note)
|
|
|
|
pretty_midi_objects_list.append(pretty_midi_object)
|
|
start_idx += end_idx + 1 # 1 represents the zero array
|
|
|
|
if return_midi:
|
|
return BatchEncoding({"notes": notes_list, "pretty_midi_objects": pretty_midi_objects_list})
|
|
|
|
return BatchEncoding({"notes": notes_list})
|