
import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


lowercase_ 				=			logging.get_logger(__name__)

lowercase_ 				=			'''▁'''

lowercase_ 				=			{'''vocab_file''': '''sentencepiece.bpe.model'''}

lowercase_ 				=			{
    '''vocab_file''': {
        '''facebook/mbart-large-50-one-to-many-mmt''': (
            '''https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model'''
        ),
    }
}

lowercase_ 				=			{
    '''facebook/mbart-large-50-one-to-many-mmt''': 1_024,
}

# fmt: off
lowercase_ 				=			['''ar_AR''', '''cs_CZ''', '''de_DE''', '''en_XX''', '''es_XX''', '''et_EE''', '''fi_FI''', '''fr_XX''', '''gu_IN''', '''hi_IN''', '''it_IT''', '''ja_XX''', '''kk_KZ''', '''ko_KR''', '''lt_LT''', '''lv_LV''', '''my_MM''', '''ne_NP''', '''nl_XX''', '''ro_RO''', '''ru_RU''', '''si_LK''', '''tr_TR''', '''vi_VN''', '''zh_CN''', '''af_ZA''', '''az_AZ''', '''bn_IN''', '''fa_IR''', '''he_IL''', '''hr_HR''', '''id_ID''', '''ka_GE''', '''km_KH''', '''mk_MK''', '''ml_IN''', '''mn_MN''', '''mr_IN''', '''pl_PL''', '''ps_AF''', '''pt_XX''', '''sv_SE''', '''sw_KE''', '''ta_IN''', '''te_IN''', '''th_TH''', '''tl_XX''', '''uk_UA''', '''ur_PK''', '''xh_ZA''', '''gl_ES''', '''sl_SI''']





class     __UpperCamelCase (							a__       ):






     """simple docstring"""







     lowerCAmelCase_							   =			VOCAB_FILES_NAMES
     lowerCAmelCase_							   =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     lowerCAmelCase_							   =			PRETRAINED_VOCAB_FILES_MAP
     lowerCAmelCase_							   =			['''input_ids''', '''attention_mask''']

     lowerCAmelCase_							   =			[]
     lowerCAmelCase_							   =			[]



     def __init__(				self							:							Any     ,						_A							:							str     ,						_A							:							List[str]=None     ,						_A							:							List[Any]=None     ,						_A							:							Union[str, Any]="</s>"     ,						_A							:							Union[str, Any]="</s>"     ,						_A							:							int="<s>"     ,						_A							:							str="<unk>"     ,						_A							:							int="<pad>"     ,						_A							:							Any="<mask>"     ,						_A							:							List[Any] = None     ,						**_A							:							Optional[Any]     ,						):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Dict					=		AddedToken(_snake_case     ,						lstrip=_snake_case     ,						rstrip=_snake_case ) if isinstance(_snake_case     ,						_snake_case ) else mask_token

       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		{} if sp_model_kwargs is None else sp_model_kwargs

       __SCREAMING_SNAKE_CASE      :	List[str]					=		kwargs.get('''additional_special_tokens'''     ,						[] )
       kwargs["additional_special_tokens"] += [
           code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
       ]

       super().__init__(
           src_lang=_snake_case     ,						tgt_lang=_snake_case     ,						eos_token=_snake_case     ,						unk_token=_snake_case     ,						sep_token=_snake_case     ,						cls_token=_snake_case     ,						pad_token=_snake_case     ,						mask_token=_snake_case     ,						sp_model_kwargs=self.sp_model_kwargs     ,						**_snake_case     ,						)

       __SCREAMING_SNAKE_CASE      :	List[str]					=		spm.SentencePieceProcessor(**self.sp_model_kwargs )
       self.sp_model.Load(str(_snake_case ) )
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		vocab_file

       # Original fairseq vocab and spm vocab must be "aligned":
       # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
       # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
       # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
       # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

       # Mimic fairseq token-to-id alignment for the first 4 token
       __SCREAMING_SNAKE_CASE      :	Dict					=		{'''<s>''': 0, '''<pad>''': 1, '''</s>''': 2, '''<unk>''': 3}

       # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
       __SCREAMING_SNAKE_CASE      :	int					=		1

       __SCREAMING_SNAKE_CASE      :	str					=		len(self.sp_model )
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		{
           code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(_snake_case )
       }
       __SCREAMING_SNAKE_CASE      :	List[str]					=		{v: k for k, v in self.lang_code_to_id.items()}
       __SCREAMING_SNAKE_CASE      :	str					=		len(self.sp_model ) + len(self.lang_code_to_id ) + self.fairseq_offset

       self.fairseq_tokens_to_ids.update(self.lang_code_to_id )
       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		{v: k for k, v in self.fairseq_tokens_to_ids.items()}

       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		src_lang if src_lang is not None else '''en_XX'''
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		self.lang_code_to_id[self._src_lang]
       __SCREAMING_SNAKE_CASE      :	List[str]					=		tgt_lang
       self.set_src_lang_special_tokens(self._src_lang )



     @property
     def  UpperCAmelCase__							(				self							:							List[Any] ):







       """simple docstring"""




       return len(self.sp_model ) + len(self.lang_code_to_id ) + self.fairseq_offset + 1  # Plus 1 for the mask token



     @property
     def  UpperCAmelCase__							(				self							:							int ):







       """simple docstring"""




       return self._src_lang



     @src_lang.setter
     def  UpperCAmelCase__							(				self							:							List[str]     ,						_A							:							List[str] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	int					=		new_src_lang
       self.set_src_lang_special_tokens(self._src_lang )



     def __getstate__(				self							:							List[Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		self.__dict__.copy()
       __SCREAMING_SNAKE_CASE      :	Tuple					=		None
       return state



     def __setstate__(				self							:							Optional[int]     ,						_A							:							Tuple ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Dict					=		d

       # for backward compatibility
       if not hasattr(self     ,						'''sp_model_kwargs''' ):
         __SCREAMING_SNAKE_CASE      :	Tuple					=		{}

       __SCREAMING_SNAKE_CASE      :	List[Any]					=		spm.SentencePieceProcessor(**self.sp_model_kwargs )
       self.sp_model.Load(self.vocab_file )



     def  UpperCAmelCase__							(				self							:							int ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	List[str]					=		{self.convert_ids_to_tokens(_snake_case ): i for i in range(self.vocab_size )}
       vocab.update(self.added_tokens_encoder )
       return vocab



     def  UpperCAmelCase__							(				self							:							Optional[Any]     ,						_A							:							str ):







       """simple docstring"""




       return self.sp_model.encode(_snake_case     ,						out_type=_snake_case )



     def  UpperCAmelCase__							(				self							:							List[str]     ,						_A							:							str ):







       """simple docstring"""




       if token in self.fairseq_tokens_to_ids:
         return self.fairseq_tokens_to_ids[token]
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		self.sp_model.PieceToId(_snake_case )

       # Need to return unknown token if the SP model returned 0
       return spm_id + self.fairseq_offset if spm_id else self.unk_token_id



     def  UpperCAmelCase__							(				self							:							Dict     ,						_A							:							str ):







       """simple docstring"""




       if index in self.fairseq_ids_to_tokens:
         return self.fairseq_ids_to_tokens[index]
       return self.sp_model.IdToPiece(index - self.fairseq_offset )



     def  UpperCAmelCase__							(				self							:							Any     ,						_A							:							Optional[int] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	int					=		[]
       __SCREAMING_SNAKE_CASE      :	Tuple					=		''''''
       __SCREAMING_SNAKE_CASE      :	Tuple					=		False
       for token in tokens:
         # make sure that special tokens are not decoded using sentencepiece model
         if token in self.all_special_tokens:
           if not prev_is_special:
             out_string += " "
           out_string += self.sp_model.decode(_snake_case ) + token
           __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		True
           __SCREAMING_SNAKE_CASE      :	List[Any]					=		[]
         else:
           current_sub_tokens.append(_snake_case )
           __SCREAMING_SNAKE_CASE      :	Optional[int]					=		False
       out_string += self.sp_model.decode(_snake_case )
       return out_string.strip()



     def  UpperCAmelCase__							(				self							:							Optional[Any]     ,						_A							:							List[str]     ,						_A							:							int = None ):







       """simple docstring"""




       if not os.path.isdir(_snake_case ):
         logger.error(F'''Vocabulary path ({save_directory}) should be a directory''' )
         return
       __SCREAMING_SNAKE_CASE      :	Dict					=		os.path.join(
           _snake_case     ,						(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file'''] )

       if os.path.abspath(self.vocab_file ) != os.path.abspath(_snake_case ) and os.path.isfile(self.vocab_file ):
         copyfile(self.vocab_file     ,						_snake_case )
       elif not os.path.isfile(self.vocab_file ):
         with open(_snake_case     ,						'''wb''' ) as fi:
           __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		self.sp_model.serialized_model_proto()
           fi.write(_snake_case )

       return (out_vocab_file,)



     def  UpperCAmelCase__							(				self							:							Optional[int]     ,						_A							:							List[str]     ,						_A							:							Optional[int] = None     ,						_A							:							Tuple = False ):







       """simple docstring"""




       if already_has_special_tokens:
         return super().get_special_tokens_mask(
             token_ids_a=_snake_case     ,						token_ids_a=_snake_case     ,						already_has_special_tokens=_snake_case )

       __SCREAMING_SNAKE_CASE      :	Tuple					=		[1] * len(self.prefix_tokens )
       __SCREAMING_SNAKE_CASE      :	Dict					=		[1] * len(self.suffix_tokens )
       if token_ids_a is None:
         return prefix_ones + ([0] * len(_snake_case )) + suffix_ones
       return prefix_ones + ([0] * len(_snake_case )) + ([0] * len(_snake_case )) + suffix_ones



     def  UpperCAmelCase__							(				self							:							Tuple     ,						_A							:							int     ,						_A							:							List[str] = None ):







       """simple docstring"""




       if token_ids_a is None:
         return self.prefix_tokens + token_ids_a + self.suffix_tokens
       # We don't expect to process pairs, but leave the pair logic for API consistency
       return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens



     def  UpperCAmelCase__							(				self							:							Any     ,						_A							:							str     ,						_A							:							List[Any]     ,						_A							:							Union[str, Any]     ,						_A							:							str     ,						**_A							:							Any ):







       """simple docstring"""




       if src_lang is None or tgt_lang is None:
         raise ValueError('''Translation requires a `src_lang` and a `tgt_lang` for this model''' )
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		src_lang
       __SCREAMING_SNAKE_CASE      :	int					=		self(_snake_case     ,						add_special_tokens=_snake_case     ,						return_tensors=_snake_case     ,						**_snake_case )
       __SCREAMING_SNAKE_CASE      :	Any					=		self.convert_tokens_to_ids(_snake_case )
       __SCREAMING_SNAKE_CASE      :	int					=		tgt_lang_id
       return inputs



     def  UpperCAmelCase__							(				self							:							str     ,						_A							:							List[Any]     ,						_A							:							List[Any] = "en_XX"     ,						_A							:							List[Any] = None     ,						_A							:							Any = "ro_RO"     ,						**_A							:							Tuple     ,						):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Any					=		src_lang
       __SCREAMING_SNAKE_CASE      :	List[str]					=		tgt_lang
       return super().prepare_seqaseq_batch(_snake_case     ,						_snake_case     ,						**_snake_case )



     def  UpperCAmelCase__							(				self							:							Dict ):







       """simple docstring"""




       return self.set_src_lang_special_tokens(self.src_lang )



     def  UpperCAmelCase__							(				self							:							int ):







       """simple docstring"""




       return self.set_tgt_lang_special_tokens(self.tgt_lang )



     def  UpperCAmelCase__							(				self							:							Tuple     ,						_A							:							Tuple ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Dict					=		self.lang_code_to_id[src_lang]
       __SCREAMING_SNAKE_CASE      :	Any					=		[self.cur_lang_code_id]
       __SCREAMING_SNAKE_CASE      :	Dict					=		[self.eos_token_id]




     def  UpperCAmelCase__							(				self							:							List[str]     ,						_A							:							Union[str, Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	List[str]					=		self.lang_code_to_id[tgt_lang]
       __SCREAMING_SNAKE_CASE      :	Any					=		[self.cur_lang_code_id]
       __SCREAMING_SNAKE_CASE      :	int					=		[self.eos_token_id]







"""simple docstring"""
import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class       a							(   a__   , unittest.TestCase  ):
 snake_case__             =       DanceDiffusionPipeline
 snake_case__             =       UNCONDITIONAL_AUDIO_GENERATION_PARAMS
 snake_case__             =       PipelineTesterMixin.required_optional_params - {
     '''callback''',
     '''latents''',
     '''callback_steps''',
     '''output_type''',
     '''num_images_per_prompt''',
 }
 snake_case__             =       UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
 snake_case__             =       False
 snake_case__             =       False



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  torch.manual_seed(0     )
  lowerCAmelCase    				= UNetaDModel(
      block_out_channels=(32, 32, 64)				,       extra_in_channels=16				,       sample_size=5_12				,       sample_rate=1_60_00				,       in_channels=2				,       out_channels=2				,       flip_sin_to_cos=_snake_case				,       use_timestep_embedding=_snake_case				,       time_embedding_type='fourier'				,       mid_block_type='UNetMidBlock1D'				,       down_block_types=('DownBlock1DNoSkip', 'DownBlock1D', 'AttnDownBlock1D')				,       up_block_types=('AttnUpBlock1D', 'UpBlock1D', 'UpBlock1DNoSkip')				,       )
  lowerCAmelCase    				= IPNDMScheduler()

  lowerCAmelCase    				= {
      'unet': unet,
      'scheduler': scheduler,
  }
  return components



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case=0     ):







  """simple docstring"""






  if str(_snake_case     ).startswith('mps'     ):
   lowerCAmelCase    				= torch.manual_seed(_snake_case     )
  else:
   lowerCAmelCase    				= torch.Generator(device=_snake_case     ).manual_seed(_snake_case     )
  lowerCAmelCase    				= {
      'batch_size': 1,
      'generator': generator,
      'num_inference_steps': 4,
  }
  return inputs



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= 'cpu'  # ensure determinism for the device-dependent torch.Generator
  lowerCAmelCase    				= self.get_dummy_components()
  lowerCAmelCase    				= DanceDiffusionPipeline(**_snake_case     )
  lowerCAmelCase    				= pipe.to(_snake_case     )
  pipe.set_progress_bar_config(disable=_snake_case     )

  lowerCAmelCase    				= self.get_dummy_inputs(_snake_case     )
  lowerCAmelCase    				= pipe(**_snake_case     )
  lowerCAmelCase    				= output.audios

  lowerCAmelCase    				= audio[0, -3:, -3:]

  assert audio.shape == (1, 2, components["unet"].sample_size)
  lowerCAmelCase    				= np.array([-0.7_265, 1.0_000, -0.8_388, 0.1_175, 0.9_498, -1.0_000]     )
  assert np.abs(audio_slice.flatten() - expected_slice     ).max() < 1E-2



 @skip_mps
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  return super().test_save_load_local()



 @skip_mps
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3     )



 @skip_mps
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  return super().test_save_load_optional_components()



 @skip_mps
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  return super().test_attention_slicing_forward_pass()



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  super().test_inference_batch_single_identical(expected_max_diff=3E-3     )







@slow
@require_torch_gpu
class       a							(   unittest.TestCase  ):



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  super().tearDown()
  gc.collect()
  torch.cuda.empty_cache()



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= torch_device

  lowerCAmelCase    				= DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k'     )
  lowerCAmelCase    				= pipe.to(_snake_case     )
  pipe.set_progress_bar_config(disable=_snake_case     )

  lowerCAmelCase    				= torch.manual_seed(0     )
  lowerCAmelCase    				= pipe(generator=_snake_case				,       num_inference_steps=1_00				,       audio_length_in_s=4.096     )
  lowerCAmelCase    				= output.audios

  lowerCAmelCase    				= audio[0, -3:, -3:]

  assert audio.shape == (1, 2, pipe.unet.sample_size)
  lowerCAmelCase    				= np.array([-0.0_192, -0.0_231, -0.0_318, -0.0_059, 0.0_002, -0.0_020]     )

  assert np.abs(audio_slice.flatten() - expected_slice     ).max() < 1E-2



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= torch_device

  lowerCAmelCase    				= DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k'				,       torch_dtype=torch.floataa     )
  lowerCAmelCase    				= pipe.to(_snake_case     )
  pipe.set_progress_bar_config(disable=_snake_case     )

  lowerCAmelCase    				= torch.manual_seed(0     )
  lowerCAmelCase    				= pipe(generator=_snake_case				,       num_inference_steps=1_00				,       audio_length_in_s=4.096     )
  lowerCAmelCase    				= output.audios

  lowerCAmelCase    				= audio[0, -3:, -3:]

  assert audio.shape == (1, 2, pipe.unet.sample_size)
  lowerCAmelCase    				= np.array([-0.0_367, -0.0_488, -0.0_771, -0.0_525, -0.0_444, -0.0_341]     )

  assert np.abs(audio_slice.flatten() - expected_slice     ).max() < 1E-2





