o Zh8@sdZddlZddlZddlmZmZmZmZddlZ ddl m Z ddl m Z mZmZmZmZddlmZmZmZmZmZddlmZerIddlZeeZd d iZd d Zd dZ eddGdddeZ!dgZ"dS)z!Tokenization class for Pop2Piano.N)ListOptionalTupleUnion) BatchFeature) AddedToken BatchEncodingPaddingStrategyPreTrainedTokenizerTruncationStrategy) TensorTypeis_pretty_midi_availableloggingrequires_backendsto_numpy)requiresvocabz vocab.jsoncCs||7}|dur t||}|SN)minnumbercutoff_time_idx current_idxrc/var/www/auris/lib/python3.10/site-packages/transformers/models/pop2piano/tokenization_pop2piano.pytoken_time_to_note(s rc CsZ||dur'||}||kr%|}|||||g|dkrdn|}|||<|S|||<|S)Nr)append) rcurrent_velocitydefault_velocitynote_onsets_readyrnotes onset_idx offset_idxZ onsets_readyrrrtoken_note_to_note0s r$) pretty_midiZtorch)backendscseZdZdZddgZeZ      dCfd d Zed d Z ddZ de de fddZ dDde fddZdejde de de fddZ   dEdejdejde de de f d d!Z "dFdejd#edeefd$d%ZdGd'ejdejd(e fd)d*ZdFd+ed,eedeefd-d.Z " "dHd'eejeejfd/eed0ee defd1d2Z " "dHd'eejeejfd/eed0ee defd3d4Z! 5 " " " " " 6dId'eejeejeeejfd7ee"ee#fd8ee"eefd0ee d9ee d:ee"d;eeee$fdZ% 6dJd?e&d@e"fdAdBZ'Z(S)KPop2PianoTokenizera 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. token_idsattention_maskM-1102c  st|tr t|dddn|}t|trt|dddn|}t|tr(t|dddn|}t|tr6t|dddn|}||_||_t|d} t| |_Wdn1sTwYdd|j D|_ t j d||||d|dS)NF)lstriprstriprbcSsi|]\}}||qSrr).0kvrrr uz/Pop2PianoTokenizer.__init__..) unk_token eos_token pad_token bos_tokenr) isinstancestrrrnum_barsopenjsonloadencoderitemsdecodersuper__init__) selfrrr>r8r9r:r;kwargsfile __class__rrrF]s$   zPop2PianoTokenizer.__init__cCs t|jS)z-Returns the vocabulary size of the tokenizer.)lenrBrGrrr vocab_sizes zPop2PianoTokenizer.vocab_sizecCst|jfi|jS)z(Returns the vocabulary of the tokenizer.)dictrBZadded_tokens_encoderrMrrr get_vocabszPop2PianoTokenizer.get_vocabtoken_idreturncCsH|j||jd}|d}d|ddt|d}}||gS)a? 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`). Z _TOKEN_TIME_Nr)rDgetr8splitjoinint)rGrQZtoken_type_value token_typevaluerrr_convert_id_to_tokens  z'Pop2PianoTokenizer._convert_id_to_token TOKEN_TIMEcCs|j|d|t|jS)a 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. rS)rBrUrXr8)rGtokenrYrrr_convert_token_to_idsz'Pop2PianoTokenizer._convert_token_to_idtokensbeat_offset_idxbars_per_batchrc Csd}tt|D]1}||}|||d}||} |j||| d} t| dkr)q|dur0| }qtj|| fdd}q|dur@gS|S)a 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. N) start_idxrr)Zaxis)rangerLrelative_tokens_ids_to_notesnpZ concatenate) rGr_r`rarr!index_tokensZ _start_idxZ_cutoff_time_idxZ_notesrrr"relative_batch_tokens_ids_to_notess$ z5Pop2PianoTokenizer.relative_batch_tokens_ids_to_notesr beatstepcCs:|durdn|}|j||||d}|j||||d}|S)al 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. Nr)r_r`rar) offset_sec)ri notes_to_midi)rGr_rkr`rarr!Zmidirrr!relative_batch_tokens_ids_to_midisz4Pop2PianoTokenizer.relative_batch_tokens_ids_to_midiNrcc s^fdd|D}|}d}ddttddjDdD}g}|D]5\} } | dkr5| dkr4n(q&| dkrAt| ||d }q&| d krH| }q&| d krXt| |j|||d }q&td t|D]'\} } | dur|durq| d} nt || d} t || }| | || jgq`t |dkrgSt |}|dddfd|dddf}||}|S)a 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. csg|]}|qSr)r[)r3r]rMrr r7zCPop2PianoTokenizer.relative_tokens_ids_to_notes..rcSsg|]}dqSrrr3irrrro cSsg|]}|dqS)NOTE)endswith)r3r4rrrro r7rTZ TOKEN_SPECIALr\rTOKEN_VELOCITY TOKEN_NOTE)rrrr rr!zToken type not understood!N)rdsumrBkeysrr$r ValueError enumeratemaxrrLrfarrayZargsort)rGr_rcrwordsrrr r!rYrpitchZ note_onsetcutoffr#Z note_orderrrMrresP*      $ z/Pop2PianoTokenizer.relative_tokens_ids_to_notesr!rlc Cst|dgtjddd}tjdd}g}|D]\}}} } tj| | ||||||d} || q||_|j|||S)a 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. r%ig^@) resolutionZ initial_tempor)program)velocityrstartend) rr%Z PrettyMIDIZ InstrumentNoterr! instrumentsZremove_invalid_notes) rGr!rkrlZnew_pmZnew_instZ new_notesr"r#rrZnew_noterrrrm8s      z Pop2PianoTokenizer.notes_to_midisave_directoryfilename_prefixcCstj|std|ddStj||r|dndtd}t|d}|t |j Wd|fS1s=wY|fS)a} 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. zVocabulary path (z) should be a directoryN-rw) ospathisdirloggererrorrWVOCAB_FILES_NAMESr?writer@dumpsrB)rGrrZout_vocab_filerIrrrsave_vocabularyXs   z"Pop2PianoTokenizer.save_vocabularytruncation_strategy max_lengthcKs~t|dgt|dtjrtdd|Ddd}t|tj }|ddddf }d dt |d D}|D]\}}} } || | | g|| | dgq>g} d} t |D]9\} }t|dkrjq_| || d |D]"\} } t| dk} | | kr| } | || d | || d quq_t| }|tjkr|r||kr|jd| |||d|\} }}td| iS)a 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. r%rcSs g|] }|j|j|j|jgqSr)rrrr)r3Z each_noterrrros z2Pop2PianoTokenizer.encode_plus..rbNr+cSsg|]}gqSrrrprrrrorrrTr\rurv)ZidsZnum_tokens_to_removerr(r)rr<r%rrfr}ZreshaperoundZastypeZint32r|rdrr{rLr^rXr ZDO_NOT_TRUNCATEZtruncate_sequencesr )rGr!rrrHZ max_time_idxtimesZonsetoffsetrrr_rrqtimeZ total_lenrSrrr encode_plusosH      zPop2PianoTokenizer.encode_pluscKsHg}tt|D]}||j||f||d|dqtd|iS)a 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. )rrr()rdrLrrr )rGr!rrrHZencoded_batch_token_idsrqrrrbatch_encode_pluss z$Pop2PianoTokenizer.batch_encode_plusFTpadding truncationpad_to_multiple_ofreturn_attention_maskreturn_tensorsverbosec  Kst|tjr |jdknt|dt} |jd|||||d| \} } }} | r;|dur,dn|}|jd|| |d| } n |jd|| |d| } |j| | |||||d} | S) a 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. rr)rrrrrNT)r!rr)rrrrrrr) r<rfndarrayndimlistZ"_get_padding_truncation_strategiesrrpad)rGr!rrrrrrrrHZ is_batchedZpadding_strategyrr(rrr__call__sH$R  zPop2PianoTokenizer.__call__feature_extractor_output return_midic Cstt|dot|dot|d}|s |djddkr td|r~t|ddddfdk|djdksE|djd|d jdkr_td |jdd |djdd |d jd|djd|jdkr}td |djdd|jdn'|djddks|d jddkrtd|djdd|d jdd|rt|ddddfdkd}n|jdg}g}g}d}t|D]\} } ||| } | dddtt| t |j kddf} |d| } |d | } |r+|d| }|d| }| dtt|dkdd} | dtt|dkdd} t | } t | } t | } |j | | |j |j ddd}|jdjD]}|j| d7_|j| d7_||qL|||| d7}q|r}t||dStd|iS)aF 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`. r)attention_mask_beatsteps$attention_mask_extrapolated_beatstep beatstepsrrTzattention_mask, attention_mask_beatsteps and attention_mask_extrapolated_beatstep must be present for batched inputs! But one of them were not present.Nextrapolated_beatstepzbLength mistamtch between token_ids, beatsteps and extrapolated_beatstep! Found token_ids length - z, beatsteps shape - z$ and extrapolated_beatsteps shape - z!Found attention_mask of length - z but token_ids of length - zLength mistamtch of beatsteps and extrapolated_beatstep! Since attention_mask is not present the number of examples must be 1, But found beatsteps length - z", extrapolated_beatsteps length - .rb)r_rkrar)r!Zpretty_midi_objectsr!)boolhasattrshaperzrxrfwherer{r|rXr9rrnr>rr!rrrr )rGr(rrZattention_masks_presentZ batch_idxZ notes_listZpretty_midi_objects_listrcrgZend_idxZeach_tokens_idsrrrrZpretty_midi_objectZnoterrr batch_decodeYs        $  2   $   zPop2PianoTokenizer.batch_decode)r*r+r,r-r.r/)r\)rr+rjr)r)NN)FNNNNNT)T))__name__ __module__ __qualname____doc__Zmodel_input_namesrZvocab_files_namesrFpropertyrNrPrXrr[r^rfrrirnfloatrrermr=rrrrr%rr r rrrr r rrr __classcell__rrrJrr'?s"  0 % >  J .         r')#rr@rtypingrrrrnumpyrfZfeature_extraction_utilsrZtokenization_utilsrr r r r utilsr rrrrZutils.import_utilsrr%Z get_loggerrrrrr$r'__all__rrrrs0