o hV}@sdZddlZddlZddlmZmZmZmZddlZ ddl m Z ddl m Z mZmZmZmZddlmZmZmZmZmZerCddlZeeZdd iZdd d iiZd d ZddZGdddeZ dS)z!Tokenization class for Pop2Piano.N)ListOptionalTupleUnion) BatchFeature) AddedToken BatchEncodingPaddingStrategyPreTrainedTokenizerTruncationStrategy) TensorTypeis_pretty_midi_availableloggingrequires_backendsto_numpyvocabz vocab.jsonzsweetcocoa/pop2pianoz@https://huggingface.co/sweetcocoa/pop2piano/blob/main/vocab.jsoncCs||7}|dur t||}|SN)minnumbercutoff_time_idx current_idxrj/var/www/html/ai/venv/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_idx onsets_readyrrrtoken_note_to_note5s r$cseZdZdZddgZeZeZ      dCfd d Z e d d Z ddZ de defddZdDde fddZdejde de de fddZ   dEdejdejde de de f d d!ZdFdejd#edefd$d%ZdGd'ejdejd(e fd)d*ZdFd+ed,eedeefd-d.Z " "dHd'eejeejfd/ee d0ee de!fd1d2Z" " "dHd'eejeejfd/ee d0ee de!fd3d4Z# 5 " " " " " 6dId'eejeejeeejfd7ee$ee%fd8ee$ee fd0ee d9ee d:ee$d;eeee&fdZ' 6dJd?e(d@e$fdAdBZ)Z*S)KPop2PianoTokenizeras 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. 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 pz/Pop2PianoTokenizer.__init__..) unk_token eos_token pad_token bos_tokenr) isinstancestrrrnum_barsopenjsonloadencoderitemsdecodersuper__init__) selfrrr<r6r7r8r9kwargsfile __class__rrrDXs$   zPop2PianoTokenizer.__init__cCs t|jS)z-Returns the vocabulary size of the tokenizer.)lenr@rErrr vocab_sizezs zPop2PianoTokenizer.vocab_sizecCst|jfi|jS)z(Returns the vocabulary of the tokenizer.)dictr@added_tokens_encoderrKrrr 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`). _TOKEN_TIME_Nr)rBgetr6splitjoinint)rErPtoken_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)r@rUrXr6)rEtokenrZrrr_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)axis)rangerJrelative_tokens_ids_to_notesnp concatenate) rEr`rarbrr index_tokens _start_idx_cutoff_time_idx_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`rarbr) offset_sec)ro notes_to_midi)rEr`rqrarbrr midirrr!relative_batch_tokens_ids_to_midisz4Pop2PianoTokenizer.relative_batch_tokens_ids_to_midiNrdc 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\)r1r^rKrr r5zCPop2PianoTokenizer.relative_tokens_ids_to_notes..rcSsg|]}dqSrrr1irrrrvcSsg|]}|dqS)NOTE)endswith)r1r2rrrrvr5rT TOKEN_SPECIALr]rTOKEN_VELOCITY TOKEN_NOTE)rrrrrr zToken type not understood!N)rfsumr@keysrr$r ValueError enumeratemaxrrJrharrayargsort)rEr`rdrwordsrrrr rZrpitch note_onsetcutoffr" note_orderrrKrrgsP *      $ z/Pop2PianoTokenizer.relative_tokens_ids_to_notesr rrc 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. pretty_midiig^@) resolution initial_tempor)program)velocityrstartend) rr PrettyMIDI InstrumentNoterr instrumentsremove_invalid_notes) rEr rqrrnew_pmnew_inst new_notesr!r"rrnew_noterrrrs1s      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>dumpsr@)rErrout_vocab_filerGrrrsave_vocabularyQs   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. rrcSs g|] }|j|j|j|jgqSr)rrrr)r1 each_noterrrrvs z2Pop2PianoTokenizer.encode_plus..rcNr)cSsg|]}gqSrrrwrrrrvryrTr]r}r~)idsnum_tokens_to_removerr&r)rr:rrrhrreshaperoundastypeint32rrfrrrJr_rXr DO_NOT_TRUNCATEtruncate_sequencesr )rEr rrrF max_time_idxtimesonsetoffsetrrr`rrxtime total_lenrSrrr encode_plushsH      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&)rfrJrrr )rEr rrrFencoded_batch_token_idsrxrrrbatch_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:rhndarrayndimlist"_get_padding_truncation_strategiesrrpad)rEr rrrrrrrrF is_batchedpadding_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 - .rc)r`rqrbr)r pretty_midi_objectsr )boolhasattrshaperrrhwhererrrXr7rrur<rr rrrr )rEr&rrattention_masks_present batch_idx notes_listpretty_midi_objects_listrdrjend_idxeach_tokens_idsrrrrpretty_midi_objectnoterrr batch_decodeRs        $  2   $   zPop2PianoTokenizer.batch_decode)r(r)r*r+r,r-)r])rr)rpr)r)NN)FNNNNNT)T)+__name__ __module__ __qualname____doc__model_input_namesrvocab_files_namesPRETRAINED_VOCAB_FILES_MAPpretrained_vocab_files_maprDpropertyrLrOrXrr\r_rhrrorufloatrgrsr;rrrrrrrr r rrrr r rrr __classcell__rrrHrr%Ds"  0 $<  J .         r%)!rr>rtypingrrrrnumpyrhfeature_extraction_utilsrtokenization_utilsrr r r r utilsr rrrrr get_loggerrrrrrr$r%rrrrs(