


"""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     __lowerCamelCase  (					a__  ,					unittest.TestCase				):

     '''simple docstring'''
     A_       :	List[Any]    =		DanceDiffusionPipeline
     A_       :	List[str]    =		UNCONDITIONAL_AUDIO_GENERATION_PARAMS
     A_       :	Optional[int]    =		PipelineTesterMixin.required_optional_params - {
         'callback',
         'latents',
         'callback_steps',
         'output_type',
         'num_images_per_prompt',
     }
     A_       :	List[str]    =		UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
     A_       :	int    =		False
     A_       :	str    =		False


     def 							_UpperCAmelCase      (      self							)       ->							List[str]:
          torch.manual_seed(0							)
          _a    	=				UNetaDModel(
              block_out_channels=(32, 32, 64)    ,				extra_in_channels=16    ,				sample_size=512    ,				sample_rate=16000    ,				in_channels=2    ,				out_channels=2    ,				flip_sin_to_cos=__UpperCAmelCase    ,				use_timestep_embedding=__UpperCAmelCase    ,				time_embedding_type='''fourier'''    ,				mid_block_type='''UNetMidBlock1D'''    ,				down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')    ,				up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')    ,				)
          _a    	=				IPNDMScheduler()

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


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase=0							)       ->							List[Any]:
          if str(__UpperCAmelCase							).startswith('''mps'''							):
               _a    	=				torch.manual_seed(__UpperCAmelCase							)
          else:
               _a    	=				torch.Generator(device=__UpperCAmelCase							).manual_seed(__UpperCAmelCase							)
          _a    	=				{
              '''batch_size''': 1,
              '''generator''': generator,
              '''num_inference_steps''': 4,
          }
          return inputs


     def 							_UpperCAmelCase      (      self							)       ->							Any:
          _a    	=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
          _a    	=				self.get_dummy_components()
          _a    	=				DanceDiffusionPipeline(**__UpperCAmelCase							)
          _a    	=				pipe.to(__UpperCAmelCase							)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

          _a    	=				self.get_dummy_inputs(__UpperCAmelCase							)
          _a    	=				pipe(**__UpperCAmelCase							)
          _a    	=				output.audios

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

          assert audio.shape == (1, 2, components["unet"].sample_size)
          _a    	=				np.array([-0.7265, 1.0000, -0.8388, 0.1175, 0.9498, -1.0000]							)
          assert np.abs(audio_slice.flatten() - expected_slice							).max() < 1e-2


     @skip_mps
     def 							_UpperCAmelCase      (      self							)       ->							Tuple:
          return super().test_save_load_local()


     @skip_mps
     def 							_UpperCAmelCase      (      self							)       ->							Optional[int]:
          return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3							)


     @skip_mps
     def 							_UpperCAmelCase      (      self							)       ->							List[Any]:
          return super().test_save_load_optional_components()


     @skip_mps
     def 							_UpperCAmelCase      (      self							)       ->							int:
          return super().test_attention_slicing_forward_pass()





     def 							_UpperCAmelCase      (      self							)       ->							List[Any]:
          super().test_inference_batch_single_identical(expected_max_diff=3e-3							)







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

     '''simple docstring'''


     def 							_UpperCAmelCase      (      self							)       ->							Optional[int]:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()


     def 							_UpperCAmelCase      (      self							)       ->							List[str]:
          _a    	=				torch_device

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

          _a    	=				torch.manual_seed(0							)
          _a    	=				pipe(generator=__UpperCAmelCase    ,				num_inference_steps=100    ,				audio_length_in_s=4.096							)
          _a    	=				output.audios

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

          assert audio.shape == (1, 2, pipe.unet.sample_size)
          _a    	=				np.array([-0.0192, -0.0231, -0.0318, -0.0059, 0.0002, -0.0020]							)

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





     def 							_UpperCAmelCase      (      self							)       ->							Optional[int]:
          _a    	=				torch_device

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

          _a    	=				torch.manual_seed(0							)
          _a    	=				pipe(generator=__UpperCAmelCase    ,				num_inference_steps=100    ,				audio_length_in_s=4.096							)
          _a    	=				output.audios

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

          assert audio.shape == (1, 2, pipe.unet.sample_size)
          _a    	=				np.array([-0.0367, -0.0488, -0.0771, -0.0525, -0.0444, -0.0341]							)

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




"""simple docstring"""






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

import sentencepiece as spm

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


__snake_case		=       logging.get_logger(__name__)
__snake_case		=       {'''vocab_file''': '''spiece.model'''}

__snake_case		=       {
    '''vocab_file''': {
        '''albert-base-v1''': '''https://huggingface.co/albert-base-v1/resolve/main/spiece.model''',
        '''albert-large-v1''': '''https://huggingface.co/albert-large-v1/resolve/main/spiece.model''',
        '''albert-xlarge-v1''': '''https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model''',
        '''albert-xxlarge-v1''': '''https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model''',
        '''albert-base-v2''': '''https://huggingface.co/albert-base-v2/resolve/main/spiece.model''',
        '''albert-large-v2''': '''https://huggingface.co/albert-large-v2/resolve/main/spiece.model''',
        '''albert-xlarge-v2''': '''https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model''',
        '''albert-xxlarge-v2''': '''https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model''',
    }
}

__snake_case		=       {
    '''albert-base-v1''': 512,
    '''albert-large-v1''': 512,
    '''albert-xlarge-v1''': 512,
    '''albert-xxlarge-v1''': 512,
    '''albert-base-v2''': 512,
    '''albert-large-v2''': 512,
    '''albert-xlarge-v2''': 512,
    '''albert-xxlarge-v2''': 512,
}

__snake_case		=       '''▁'''


class     __lowerCamelCase  (					a__				):

     '''simple docstring'''
     A_       :	Optional[int]    =		VOCAB_FILES_NAMES
     A_       :	int    =		PRETRAINED_VOCAB_FILES_MAP
     A_       :	Optional[int]    =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase=True    ,				__UpperCAmelCase=True    ,				__UpperCAmelCase=False    ,				__UpperCAmelCase="[CLS]"    ,				__UpperCAmelCase="[SEP]"    ,				__UpperCAmelCase="<unk>"    ,				__UpperCAmelCase="[SEP]"    ,				__UpperCAmelCase="<pad>"    ,				__UpperCAmelCase="[CLS]"    ,				__UpperCAmelCase="[MASK]"    ,				__UpperCAmelCase = None    ,				**__UpperCAmelCase    ,				)       ->							None:
          # Mask token behave like a normal word, i.e. include the space before it and
          # is included in the raw text, there should be a match in a non-normalized sentence.
          _a    	=				(
              AddedToken(__UpperCAmelCase    ,				lstrip=__UpperCAmelCase    ,				rstrip=__UpperCAmelCase    ,				normalized=__UpperCAmelCase							)
              if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							)
              else mask_token
          )

          _a    	=				{} if sp_model_kwargs is None else sp_model_kwargs

          super().__init__(
              do_lower_case=__UpperCAmelCase    ,				remove_space=__UpperCAmelCase    ,				keep_accents=__UpperCAmelCase    ,				bos_token=__UpperCAmelCase    ,				eos_token=__UpperCAmelCase    ,				unk_token=__UpperCAmelCase    ,				sep_token=__UpperCAmelCase    ,				pad_token=__UpperCAmelCase    ,				cls_token=__UpperCAmelCase    ,				mask_token=__UpperCAmelCase    ,				sp_model_kwargs=self.sp_model_kwargs    ,				**__UpperCAmelCase    ,				)

          _a    	=				do_lower_case
          _a    	=				remove_space
          _a    	=				keep_accents
          _a    	=				vocab_file

          _a    	=				spm.SentencePieceProcessor(**self.sp_model_kwargs							)
          self.sp_model.Load(__UpperCAmelCase							)


     @property
     def 							_UpperCAmelCase      (      self							)       ->							str:
          return len(self.sp_model							)


     def 							_UpperCAmelCase      (      self							)       ->							Optional[Any]:
          _a    	=				{self.convert_ids_to_tokens(__UpperCAmelCase							): i for i in range(self.vocab_size							)}
          vocab.update(self.added_tokens_encoder							)
          return vocab


     def __getstate__(      self							)       ->							int:
          _a    	=				self.__dict__.copy()
          _a    	=				None
          return state


     def __setstate__(      self    ,				__UpperCAmelCase							)       ->							Tuple:
          _a    	=				d

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

          _a    	=				spm.SentencePieceProcessor(**self.sp_model_kwargs							)
          self.sp_model.Load(self.vocab_file							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[str]:
          if self.remove_space:
               _a    	=				''' '''.join(inputs.strip().split()							)
          else:
               _a    	=				inputs
          _a    	=				outputs.replace('''``'''    ,				'''"'''							).replace('''\'\''''    ,				'''"'''							)

          if not self.keep_accents:
               _a    	=				unicodedata.normalize('''NFKD'''    ,				__UpperCAmelCase							)
               _a    	=				''''''.join([c for c in outputs if not unicodedata.combining(__UpperCAmelCase							)]							)
          if self.do_lower_case:
               _a    	=				outputs.lower()

          return outputs


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[str]:
          _a    	=				self.preprocess_text(__UpperCAmelCase							)
          _a    	=				self.sp_model.encode(__UpperCAmelCase    ,				out_type=__UpperCAmelCase							)
          _a    	=				[]
          for piece in pieces:
               if len(__UpperCAmelCase							) > 1 and piece[-1] == str(''','''							) and piece[-2].isdigit():
                    _a    	=				self.sp_model.EncodeAsPieces(piece[:-1].replace(__UpperCAmelCase    ,				''''''							)							)
                    if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
                         if len(cur_pieces[0]							) == 1:
                              _a    	=				cur_pieces[1:]
                         else:
                              _a    	=				cur_pieces[0][1:]
                    cur_pieces.append(piece[-1]							)
                    new_pieces.extend(__UpperCAmelCase							)
               else:
                    new_pieces.append(__UpperCAmelCase							)

          return new_pieces


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							str:
          return self.sp_model.PieceToId(__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							int:
          return self.sp_model.IdToPiece(__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Optional[Any]:
          _a    	=				[]
          _a    	=				''''''
          _a    	=				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(__UpperCAmelCase							) + token
                    _a    	=				True
                    _a    	=				[]
               else:
                    current_sub_tokens.append(__UpperCAmelCase							)
                    _a    	=				False
          out_string += self.sp_model.decode(__UpperCAmelCase							)
          return out_string.strip()


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							List[int]:
          _a    	=				[self.sep_token_id]
          _a    	=				[self.cls_token_id]
          if token_ids_a is None:
               return cls + token_ids_a + sep
          return cls + token_ids_a + sep + token_ids_a + sep


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = False							)       ->							List[int]:

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

          if token_ids_a is not None:
               return [1] + ([0] * len(__UpperCAmelCase							)) + [1] + ([0] * len(__UpperCAmelCase							)) + [1]
          return [1] + ([0] * len(__UpperCAmelCase							)) + [1]


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							List[int]:
          _a    	=				[self.sep_token_id]
          _a    	=				[self.cls_token_id]

          if token_ids_a is None:
               return len(cls + token_ids_a + sep							) * [0]
          return len(cls + token_ids_a + sep							) * [0] + len(token_ids_a + sep							) * [1]





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							Tuple[str]:
          if not os.path.isdir(__UpperCAmelCase							):
               logger.error(F'Vocabulary path ({save_directory}) should be a directory'							)
               return
          _a    	=				os.path.join(
              __UpperCAmelCase    ,				(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']							)

          if os.path.abspath(self.vocab_file							) != os.path.abspath(__UpperCAmelCase							) and os.path.isfile(self.vocab_file							):
               copyfile(self.vocab_file    ,				__UpperCAmelCase							)
          elif not os.path.isfile(self.vocab_file							):
               with open(__UpperCAmelCase    ,				'''wb'''							) as fi:
                    _a    	=				self.sp_model.serialized_model_proto()
                    fi.write(__UpperCAmelCase							)

          return (out_vocab_file,)