





"""simple docstring"""


import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest



class   lowercase_					(  __lowerCAmelCase				):
       snake_case__					:				Dict									=  (HeunDiscreteScheduler,)
       snake_case__					:				Tuple									=  10

       def 		_SCREAMING_SNAKE_CASE(      self ,       **__SCREAMING_SNAKE_CASE  ):






        '''simple docstring'''

        snake_case__     				=					{
            "num_train_timesteps": 1100,
            "beta_start": 0.0001,
            "beta_end": 0.02,
            "beta_schedule": "linear",
        }

        config.update(**__SCREAMING_SNAKE_CASE  )
        return config

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

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

       def 		_SCREAMING_SNAKE_CASE(      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=__SCREAMING_SNAKE_CASE ,       beta_end=__SCREAMING_SNAKE_CASE  )

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

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

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

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

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

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

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					self.scheduler_classes[0]
        snake_case__     				=					self.get_scheduler_config()
        snake_case__     				=					scheduler_class(**__SCREAMING_SNAKE_CASE  )

        scheduler.set_timesteps(self.num_inference_steps  )

        snake_case__     				=					self.dummy_model()
        snake_case__     				=					self.dummy_sample_deter * scheduler.init_noise_sigma
        snake_case__     				=					sample.to(__SCREAMING_SNAKE_CASE  )

        for i, t in enumerate(scheduler.timesteps  ):
         snake_case__     				=					scheduler.scale_model_input(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )

         snake_case__     				=					model(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )

         snake_case__     				=					scheduler.step(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )
         snake_case__     				=					output.prev_sample

        snake_case__     				=					torch.sum(torch.abs(__SCREAMING_SNAKE_CASE  )  )
        snake_case__     				=					torch.mean(torch.abs(__SCREAMING_SNAKE_CASE  )  )

        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 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					self.scheduler_classes[0]
        snake_case__     				=					self.get_scheduler_config(prediction_type="v_prediction"  )
        snake_case__     				=					scheduler_class(**__SCREAMING_SNAKE_CASE  )

        scheduler.set_timesteps(self.num_inference_steps  )

        snake_case__     				=					self.dummy_model()
        snake_case__     				=					self.dummy_sample_deter * scheduler.init_noise_sigma
        snake_case__     				=					sample.to(__SCREAMING_SNAKE_CASE  )

        for i, t in enumerate(scheduler.timesteps  ):
         snake_case__     				=					scheduler.scale_model_input(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )

         snake_case__     				=					model(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )

         snake_case__     				=					scheduler.step(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )
         snake_case__     				=					output.prev_sample

        snake_case__     				=					torch.sum(torch.abs(__SCREAMING_SNAKE_CASE  )  )
        snake_case__     				=					torch.mean(torch.abs(__SCREAMING_SNAKE_CASE  )  )

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

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					self.scheduler_classes[0]
        snake_case__     				=					self.get_scheduler_config()
        snake_case__     				=					scheduler_class(**__SCREAMING_SNAKE_CASE  )

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

        snake_case__     				=					self.dummy_model()
        snake_case__     				=					self.dummy_sample_deter.to(__SCREAMING_SNAKE_CASE  ) * scheduler.init_noise_sigma

        for t in scheduler.timesteps:
         snake_case__     				=					scheduler.scale_model_input(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )

         snake_case__     				=					model(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )

         snake_case__     				=					scheduler.step(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )
         snake_case__     				=					output.prev_sample

        snake_case__     				=					torch.sum(torch.abs(__SCREAMING_SNAKE_CASE  )  )
        snake_case__     				=					torch.mean(torch.abs(__SCREAMING_SNAKE_CASE  )  )

        if str(__SCREAMING_SNAKE_CASE  ).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(__SCREAMING_SNAKE_CASE  ).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 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					self.scheduler_classes[0]
        snake_case__     				=					self.get_scheduler_config()
        snake_case__     				=					scheduler_class(**__SCREAMING_SNAKE_CASE ,       use_karras_sigmas=__SCREAMING_SNAKE_CASE  )

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

        snake_case__     				=					self.dummy_model()
        snake_case__     				=					self.dummy_sample_deter.to(__SCREAMING_SNAKE_CASE  ) * scheduler.init_noise_sigma
        snake_case__     				=					sample.to(__SCREAMING_SNAKE_CASE  )

        for t in scheduler.timesteps:
         snake_case__     				=					scheduler.scale_model_input(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )

         snake_case__     				=					model(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )

         snake_case__     				=					scheduler.step(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  )
         snake_case__     				=					output.prev_sample

        snake_case__     				=					torch.sum(torch.abs(__SCREAMING_SNAKE_CASE  )  )
        snake_case__     				=					torch.mean(torch.abs(__SCREAMING_SNAKE_CASE  )  )

        assert abs(result_sum.item() - 0.00015  ) < 1e-2
        assert abs(result_mean.item() - 1.9869_5545_3503_4695e-07  ) < 1e-2








"""simple docstring"""


import unittest

from transformers import SpeechaTextaConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers.models.speech_to_text_a.modeling_speech_to_text_a import (
       SpeechaTextaDecoder,
       SpeechaTextaForCausalLM,
   )



@require_torch
class   lowercase_					:

       def __init__(      self ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE=99 ,       __SCREAMING_SNAKE_CASE=13 ,       __SCREAMING_SNAKE_CASE=16 ,       __SCREAMING_SNAKE_CASE=7 ,       __SCREAMING_SNAKE_CASE=True ,       __SCREAMING_SNAKE_CASE=True ,       __SCREAMING_SNAKE_CASE=True ,       __SCREAMING_SNAKE_CASE=False ,       __SCREAMING_SNAKE_CASE=True ,       __SCREAMING_SNAKE_CASE=2 ,       __SCREAMING_SNAKE_CASE=32 ,       __SCREAMING_SNAKE_CASE=4 ,       __SCREAMING_SNAKE_CASE=4 ,       __SCREAMING_SNAKE_CASE=30 ,       __SCREAMING_SNAKE_CASE=0 ,       __SCREAMING_SNAKE_CASE=1 ,       __SCREAMING_SNAKE_CASE=2 ,       __SCREAMING_SNAKE_CASE=None ,       ):






        '''simple docstring'''

        snake_case__     				=					parent
        snake_case__     				=					batch_size
        snake_case__     				=					decoder_seq_length
        # For common tests
        snake_case__     				=					self.decoder_seq_length
        snake_case__     				=					is_training
        snake_case__     				=					use_attention_mask
        snake_case__     				=					use_labels

        snake_case__     				=					vocab_size
        snake_case__     				=					d_model
        snake_case__     				=					d_model
        snake_case__     				=					decoder_layers
        snake_case__     				=					decoder_layers
        snake_case__     				=					decoder_ffn_dim
        snake_case__     				=					decoder_attention_heads
        snake_case__     				=					decoder_attention_heads
        snake_case__     				=					eos_token_id
        snake_case__     				=					bos_token_id
        snake_case__     				=					pad_token_id
        snake_case__     				=					decoder_start_token_id
        snake_case__     				=					use_cache
        snake_case__     				=					max_position_embeddings

        snake_case__     				=					None
        snake_case__     				=					decoder_seq_length
        snake_case__     				=					2
        snake_case__     				=					1

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					ids_tensor([self.batch_size, self.decoder_seq_length] ,       self.vocab_size  )

        snake_case__     				=					None
        if self.use_attention_mask:
         snake_case__     				=					ids_tensor([self.batch_size, self.decoder_seq_length] ,       vocab_size=2  )

        snake_case__     				=					None
        if self.use_labels:
         snake_case__     				=					ids_tensor([self.batch_size, self.decoder_seq_length] ,       self.vocab_size  )

        snake_case__     				=					SpeechaTextaConfig(
            vocab_size=self.vocab_size ,       d_model=self.d_model ,       decoder_layers=self.decoder_layers ,       decoder_ffn_dim=self.decoder_ffn_dim ,       decoder_attention_heads=self.decoder_attention_heads ,       eos_token_id=self.eos_token_id ,       bos_token_id=self.bos_token_id ,       use_cache=self.use_cache ,       pad_token_id=self.pad_token_id ,       decoder_start_token_id=self.decoder_start_token_id ,       max_position_embeddings=self.max_position_embeddings ,       )

        return (
            config,
            input_ids,
            attention_mask,
            lm_labels,
        )

       def 		_SCREAMING_SNAKE_CASE(      self ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       ):






        '''simple docstring'''

        snake_case__     				=					True
        snake_case__     				=					SpeechaTextaDecoder(config=__SCREAMING_SNAKE_CASE  ).to(__SCREAMING_SNAKE_CASE  ).eval()
        snake_case__     				=					input_ids[:2]

        input_ids[input_ids == 0] += 1
        # first forward pass
        snake_case__     				=					model(__SCREAMING_SNAKE_CASE ,       use_cache=__SCREAMING_SNAKE_CASE  )
        snake_case__     				=					model(__SCREAMING_SNAKE_CASE  )
        snake_case__     				=					model(__SCREAMING_SNAKE_CASE ,       use_cache=__SCREAMING_SNAKE_CASE  )

        self.parent.assertTrue(len(__SCREAMING_SNAKE_CASE  ) == len(__SCREAMING_SNAKE_CASE  )  )
        self.parent.assertTrue(len(__SCREAMING_SNAKE_CASE  ) == len(__SCREAMING_SNAKE_CASE  ) + 1  )

        snake_case__     				=					outputs["past_key_values"]

        # create hypothetical next token and extent to next_input_ids
        snake_case__     				=					ids_tensor((2, 1) ,       config.vocab_size - 1  ) + 1

        # append to next input_ids and
        snake_case__     				=					torch.cat([input_ids, next_tokens] ,       dim=-1  )

        print(__SCREAMING_SNAKE_CASE  )

        snake_case__     				=					model(__SCREAMING_SNAKE_CASE  )["last_hidden_state"]
        snake_case__     				=					model(__SCREAMING_SNAKE_CASE ,       past_key_values=__SCREAMING_SNAKE_CASE  )["last_hidden_state"]

        # select random slice
        snake_case__     				=					ids_tensor((1,) ,       output_from_past.shape[-1]  ).item()
        snake_case__     				=					output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
        snake_case__     				=					output_from_past[:, 0, random_slice_idx].detach()

        # test that outputs are equal for slice
        assert torch.allclose(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       atol=1e-3  )

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					self.prepare_config_and_inputs()
        (
            (
            snake_case__
        )				, (
            snake_case__
        )				, (
            snake_case__
        )				, (
            snake_case__
        )				, 
        )     				=					config_and_inputs

        snake_case__     				=					{
            "input_ids": input_ids,
            "attention_mask": attention_mask,
        }
        return config, inputs_dict



@require_torch
class   lowercase_					(  __lowerCAmelCase					,__lowerCAmelCase					,__lowerCAmelCase					,unittest.TestCase				):
       snake_case__					:				Tuple									=  (SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
       snake_case__					:				Optional[Any]									=  (SpeechaTextaForCausalLM,) if is_torch_available() else ()
       snake_case__					:				List[Any]									=  {"""text-generation""": SpeechaTextaForCausalLM} if is_torch_available() else {}
       snake_case__					:				Dict									=  True
       snake_case__					:				Dict									=  False

       def 		_SCREAMING_SNAKE_CASE(      self ,       ):






        '''simple docstring'''

        snake_case__     				=					SpeechaTextaStandaloneDecoderModelTester(self ,       is_training=__SCREAMING_SNAKE_CASE  )
        snake_case__     				=					ConfigTester(self ,       config_class=__SCREAMING_SNAKE_CASE  )

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        pass

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        pass

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        pass

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        self.config_tester.run_common_tests()

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_decoder_model_past(*__SCREAMING_SNAKE_CASE  )

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        return

