

'''simple docstring'''






import shutil
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SpeechaTextTokenizer, is_speech_available
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_torch, require_torchaudio
from transformers.utils import FEATURE_EXTRACTOR_NAME

from .test_feature_extraction_speech_to_text import floats_list


if is_speech_available():
      from transformers import SpeechaTextFeatureExtractor, SpeechaTextProcessor


snake_case_       =       get_tests_dir('fixtures/test_sentencepiece.model')
@require_torch
@require_torchaudio
@require_sentencepiece
class 	SCREAMING_SNAKE_CASE__			(				unittest.TestCase      ):




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			str     =       tempfile.mkdtemp()

          SCREAMING_SNAKE_CASE_      :			List[Any]     =       ["<s>", "<pad>", "</s>", "<unk>", "▁This", "▁is", "▁a", "▁t", "est"]
          SCREAMING_SNAKE_CASE_      :			Tuple     =       dict(zip(lowercase__				,     range(len(lowercase__			)			)			)			)
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       Path(self.tmpdirname			)
          save_json(lowercase__				,     save_dir / VOCAB_FILES_NAMES["vocab_file"]			)
          if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
                copyfile(lowercase__				,     save_dir / VOCAB_FILES_NAMES["spm_file"]			)

          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       SpeechaTextTokenizer.from_pretrained(self.tmpdirname			)
          tokenizer.save_pretrained(self.tmpdirname			)

          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       {
              "feature_size": 24,
              "num_mel_bins": 24,
              "padding_value": 0.0,
              "sampling_rate": 1_6000,
              "return_attention_mask": False,
              "do_normalize": True,
          }
          save_json(lowercase__				,     save_dir / FEATURE_EXTRACTOR_NAME			)




    def   __lowerCamelCase      (					self				,     **lowercase__			):






          """simple docstring"""






          return SpeechaTextTokenizer.from_pretrained(self.tmpdirname				,     **lowercase__			)




    def   __lowerCamelCase      (					self				,     **lowercase__			):






          """simple docstring"""






          return SpeechaTextFeatureExtractor.from_pretrained(self.tmpdirname				,     **lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          shutil.rmtree(self.tmpdirname			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			str     =       self.get_tokenizer()
          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       self.get_feature_extractor()

          SCREAMING_SNAKE_CASE_      :			str     =       SpeechaTextProcessor(tokenizer=lowercase__				,     feature_extractor=lowercase__			)

          processor.save_pretrained(self.tmpdirname			)
          SCREAMING_SNAKE_CASE_      :			Tuple     =       SpeechaTextProcessor.from_pretrained(self.tmpdirname			)

          self.assertEqual(processor.tokenizer.get_vocab()				,     tokenizer.get_vocab()			)
          self.assertIsInstance(processor.tokenizer				,     lowercase__			)

          self.assertEqual(processor.feature_extractor.to_json_string()				,     feature_extractor.to_json_string()			)
          self.assertIsInstance(processor.feature_extractor				,     lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			List[str]     =       SpeechaTextProcessor(
              tokenizer=self.get_tokenizer()				,     feature_extractor=self.get_feature_extractor()			)
          processor.save_pretrained(self.tmpdirname			)

          SCREAMING_SNAKE_CASE_      :			str     =       self.get_tokenizer(bos_token="(BOS)"				,     eos_token="(EOS)"			)
          SCREAMING_SNAKE_CASE_      :			Dict     =       self.get_feature_extractor(do_normalize=lowercase__				,     padding_value=1.0			)

          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       SpeechaTextProcessor.from_pretrained(
              self.tmpdirname				,     bos_token="(BOS)"				,     eos_token="(EOS)"				,     do_normalize=lowercase__				,     padding_value=1.0			)

          self.assertEqual(processor.tokenizer.get_vocab()				,     tokenizer_add_kwargs.get_vocab()			)
          self.assertIsInstance(processor.tokenizer				,     lowercase__			)

          self.assertEqual(processor.feature_extractor.to_json_string()				,     feature_extractor_add_kwargs.to_json_string()			)
          self.assertIsInstance(processor.feature_extractor				,     lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Dict     =       self.get_feature_extractor()
          SCREAMING_SNAKE_CASE_      :			Tuple     =       self.get_tokenizer()

          SCREAMING_SNAKE_CASE_      :			Dict     =       SpeechaTextProcessor(tokenizer=lowercase__				,     feature_extractor=lowercase__			)

          SCREAMING_SNAKE_CASE_      :			Any     =       floats_list((3, 1000)			)

          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       feature_extractor(lowercase__				,     return_tensors="np"			)
          SCREAMING_SNAKE_CASE_      :			Any     =       processor(lowercase__				,     return_tensors="np"			)

          for key in input_feat_extract.keys():
                self.assertAlmostEqual(input_feat_extract[key].sum()				,     input_processor[key].sum()				,     delta=1e-2			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			str     =       self.get_feature_extractor()
          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.get_tokenizer()

          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       SpeechaTextProcessor(tokenizer=lowercase__				,     feature_extractor=lowercase__			)

          SCREAMING_SNAKE_CASE_      :			List[str]     =       "This is a test string"

          SCREAMING_SNAKE_CASE_      :			int     =       processor(text=lowercase__			)

          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       tokenizer(lowercase__			)

          for key in encoded_tok.keys():
                self.assertListEqual(encoded_tok[key]				,     encoded_processor[key]			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Dict     =       self.get_feature_extractor()
          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.get_tokenizer()

          SCREAMING_SNAKE_CASE_      :			List[Any]     =       SpeechaTextProcessor(tokenizer=lowercase__				,     feature_extractor=lowercase__			)

          SCREAMING_SNAKE_CASE_      :			Dict     =       [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

          SCREAMING_SNAKE_CASE_      :			int     =       processor.batch_decode(lowercase__			)
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       tokenizer.batch_decode(lowercase__			)

          self.assertListEqual(lowercase__				,     lowercase__			)







    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       self.get_feature_extractor()
          SCREAMING_SNAKE_CASE_      :			Any     =       self.get_tokenizer()

          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       SpeechaTextProcessor(tokenizer=lowercase__				,     feature_extractor=lowercase__			)

          self.assertListEqual(
              processor.model_input_names				,     feature_extractor.model_input_names				,     msg="`processor` and `feature_extractor` model input names do not match"				,     )




'''simple docstring'''






import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest
class 	SCREAMING_SNAKE_CASE__			(				_UpperCAmelCase      ):
    _A			    =							(HeunDiscreteScheduler,)
    _A			    =							10




    def   __lowerCamelCase      (					self				,     **lowercase__			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			List[Any]     =       {
              "num_train_timesteps": 1100,
              "beta_start": 0.0001,
              "beta_end": 0.02,
              "beta_schedule": "linear",
          }

          config.update(**lowercase__			)
          return config




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          for timesteps in [10, 50, 100, 1000]:
                self.check_over_configs(num_train_timesteps=lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          for beta_start, beta_end in zip([0.00001, 0.0001, 0.001]				,     [0.0002, 0.002, 0.02]			):
                self.check_over_configs(beta_start=lowercase__				,     beta_end=lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          for schedule in ["linear", "scaled_linear", "exp"]:
                self.check_over_configs(beta_schedule=lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          for clip_sample_range in [1.0, 2.0, 3.0]:
                self.check_over_configs(clip_sample_range=lowercase__				,     clip_sample=lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          for prediction_type in ["epsilon", "v_prediction", "sample"]:
                self.check_over_configs(prediction_type=lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Tuple     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.get_scheduler_config()
          SCREAMING_SNAKE_CASE_      :			List[str]     =       scheduler_class(**lowercase__			)

          scheduler.set_timesteps(self.num_inference_steps			)

          SCREAMING_SNAKE_CASE_      :			int     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			Dict     =       self.dummy_sample_deter * scheduler.init_noise_sigma
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       sample.to(lowercase__			)

          for i, t in enumerate(scheduler.timesteps			):
                SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       scheduler.scale_model_input(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Dict     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			str     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__			)
                SCREAMING_SNAKE_CASE_      :			Dict     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			List[Any]     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			List[str]     =       torch.mean(torch.abs(lowercase__			)			)

          if torch_device in ["cpu", "mps"]:
                assert abs(result_sum.item() - 0.1233			) < 1e-2
                assert abs(result_mean.item() - 0.0002			) < 1e-3
          else:
                # CUDA
                assert abs(result_sum.item() - 0.1233			) < 1e-2
                assert abs(result_mean.item() - 0.0002			) < 1e-3




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       self.get_scheduler_config(prediction_type="v_prediction"			)
          SCREAMING_SNAKE_CASE_      :			Any     =       scheduler_class(**lowercase__			)

          scheduler.set_timesteps(self.num_inference_steps			)

          SCREAMING_SNAKE_CASE_      :			Any     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       self.dummy_sample_deter * scheduler.init_noise_sigma
          SCREAMING_SNAKE_CASE_      :			List[str]     =       sample.to(lowercase__			)

          for i, t in enumerate(scheduler.timesteps			):
                SCREAMING_SNAKE_CASE_      :			Tuple     =       scheduler.scale_model_input(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			int     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__			)
                SCREAMING_SNAKE_CASE_      :			Optional[int]     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			int     =       torch.mean(torch.abs(lowercase__			)			)

          if torch_device in ["cpu", "mps"]:
                assert abs(result_sum.item() - 4.69_34e-07			) < 1e-2
                assert abs(result_mean.item() - 6.11_12e-10			) < 1e-3
          else:
                # CUDA
                assert abs(result_sum.item() - 4.6_93_42_86_50_17_09_72e-07			) < 1e-2
                assert abs(result_mean.item() - 0.0002			) < 1e-3




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       self.get_scheduler_config()
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       scheduler_class(**lowercase__			)

          scheduler.set_timesteps(self.num_inference_steps				,     device=lowercase__			)

          SCREAMING_SNAKE_CASE_      :			Any     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			str     =       self.dummy_sample_deter.to(lowercase__			) * scheduler.init_noise_sigma

          for t in scheduler.timesteps:
                SCREAMING_SNAKE_CASE_      :			Optional[int]     =       scheduler.scale_model_input(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Dict     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__			)
                SCREAMING_SNAKE_CASE_      :			List[Any]     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			str     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       torch.mean(torch.abs(lowercase__			)			)

          if str(lowercase__			).startswith("cpu"			):
                # The following sum varies between 148 and 156 on mps. Why?
                assert abs(result_sum.item() - 0.1233			) < 1e-2
                assert abs(result_mean.item() - 0.0002			) < 1e-3
          elif str(lowercase__			).startswith("mps"			):
                # Larger tolerance on mps
                assert abs(result_mean.item() - 0.0002			) < 1e-2
          else:
                # CUDA
                assert abs(result_sum.item() - 0.1233			) < 1e-2
                assert abs(result_mean.item() - 0.0002			) < 1e-3







    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			int     =       self.get_scheduler_config()
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       scheduler_class(**lowercase__				,     use_karras_sigmas=lowercase__			)

          scheduler.set_timesteps(self.num_inference_steps				,     device=lowercase__			)

          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       self.dummy_sample_deter.to(lowercase__			) * scheduler.init_noise_sigma
          SCREAMING_SNAKE_CASE_      :			Any     =       sample.to(lowercase__			)

          for t in scheduler.timesteps:
                SCREAMING_SNAKE_CASE_      :			Optional[int]     =       scheduler.scale_model_input(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			List[str]     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__			)
                SCREAMING_SNAKE_CASE_      :			Tuple     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			Any     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       torch.mean(torch.abs(lowercase__			)			)

          assert abs(result_sum.item() - 0.00015			) < 1e-2
          assert abs(result_mean.item() - 1.98_69_55_45_35_03_46_95e-07			) < 1e-2

