



"""simple docstring"""






import unittest

from transformers import YosoConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        YosoForMaskedLM,
        YosoForMultipleChoice,
        YosoForQuestionAnswering,
        YosoForSequenceClassification,
        YosoForTokenClassification,
        YosoModel,
    )
    from transformers.models.yoso.modeling_yoso import YOSO_PRETRAINED_MODEL_ARCHIVE_LIST




class    __snake_case:




    def __init__(   self    ,							A_    ,							A_=13    ,							A_=7    ,							A_=True    ,							A_=True    ,							A_=True    ,							A_=True    ,							A_=99    ,							A_=32    ,							A_=5    ,							A_=4    ,							A_=37    ,							A_="gelu"    ,							A_=0.1    ,							A_=0.1    ,							A_=512    ,							A_=16    ,							A_=2    ,							A_=0.02    ,							A_=3    ,							A_=4    ,							A_=None    ,							):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			parent
      _SCREAMING_SNAKE_CASE        =			batch_size
      _SCREAMING_SNAKE_CASE        =			seq_length
      _SCREAMING_SNAKE_CASE        =			is_training
      _SCREAMING_SNAKE_CASE        =			use_input_mask
      _SCREAMING_SNAKE_CASE        =			use_token_type_ids
      _SCREAMING_SNAKE_CASE        =			use_labels
      _SCREAMING_SNAKE_CASE        =			vocab_size
      _SCREAMING_SNAKE_CASE        =			hidden_size
      _SCREAMING_SNAKE_CASE        =			num_hidden_layers
      _SCREAMING_SNAKE_CASE        =			num_attention_heads
      _SCREAMING_SNAKE_CASE        =			intermediate_size
      _SCREAMING_SNAKE_CASE        =			hidden_act
      _SCREAMING_SNAKE_CASE        =			hidden_dropout_prob
      _SCREAMING_SNAKE_CASE        =			attention_probs_dropout_prob
      _SCREAMING_SNAKE_CASE        =			max_position_embeddings
      _SCREAMING_SNAKE_CASE        =			type_vocab_size
      _SCREAMING_SNAKE_CASE        =			type_sequence_label_size
      _SCREAMING_SNAKE_CASE        =			initializer_range
      _SCREAMING_SNAKE_CASE        =			num_labels
      _SCREAMING_SNAKE_CASE        =			num_choices
      _SCREAMING_SNAKE_CASE        =			scope




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			ids_tensor([self.batch_size, self.seq_length]    ,							self.vocab_size			)

      _SCREAMING_SNAKE_CASE        =			None
      if self.use_input_mask:
        _SCREAMING_SNAKE_CASE        =			random_attention_mask([self.batch_size, self.seq_length]			)

      _SCREAMING_SNAKE_CASE        =			None
      if self.use_token_type_ids:
        _SCREAMING_SNAKE_CASE        =			ids_tensor([self.batch_size, self.seq_length]    ,							self.type_vocab_size			)

      _SCREAMING_SNAKE_CASE        =			None
      _SCREAMING_SNAKE_CASE        =			None
      _SCREAMING_SNAKE_CASE        =			None
      if self.use_labels:
        _SCREAMING_SNAKE_CASE        =			ids_tensor([self.batch_size]    ,							self.type_sequence_label_size			)
        _SCREAMING_SNAKE_CASE        =			ids_tensor([self.batch_size, self.seq_length]    ,							self.num_labels			)
        _SCREAMING_SNAKE_CASE        =			ids_tensor([self.batch_size]    ,							self.num_choices			)

      _SCREAMING_SNAKE_CASE        =			self.get_config()

      return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




    def   A   (   self			):





      '''simple docstring'''



      return YosoConfig(
          vocab_size=self.vocab_size    ,							hidden_size=self.hidden_size    ,							num_hidden_layers=self.num_hidden_layers    ,							num_attention_heads=self.num_attention_heads    ,							intermediate_size=self.intermediate_size    ,							hidden_act=self.hidden_act    ,							hidden_dropout_prob=self.hidden_dropout_prob    ,							attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,							max_position_embeddings=self.max_position_embeddings    ,							type_vocab_size=self.type_vocab_size    ,							is_decoder=A_    ,							initializer_range=self.initializer_range    ,							)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.get_config()
      _SCREAMING_SNAKE_CASE        =			300
      return config




    def   A   (   self			):





      '''simple docstring'''



      (
          (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     
      )        =			self.prepare_config_and_inputs()

      _SCREAMING_SNAKE_CASE        =			True
      _SCREAMING_SNAKE_CASE        =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
      _SCREAMING_SNAKE_CASE        =			ids_tensor([self.batch_size, self.seq_length]    ,							vocab_size=2			)

      return (
          config,
          input_ids,
          token_type_ids,
          input_mask,
          sequence_labels,
          token_labels,
          choice_labels,
          encoder_hidden_states,
          encoder_attention_mask,
      )




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			YosoModel(config=A_			)
      model.to(A_			)
      model.eval()
      _SCREAMING_SNAKE_CASE        =			model(A_    ,							attention_mask=A_    ,							token_type_ids=A_			)
      _SCREAMING_SNAKE_CASE        =			model(A_    ,							token_type_ids=A_			)
      _SCREAMING_SNAKE_CASE        =			model(A_			)
      self.parent.assertEqual(result.last_hidden_state.shape    ,							(self.batch_size, self.seq_length, self.hidden_size)			)




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			True
      _SCREAMING_SNAKE_CASE        =			YosoModel(A_			)
      model.to(A_			)
      model.eval()
      _SCREAMING_SNAKE_CASE        =			model(
          A_    ,							attention_mask=A_    ,							token_type_ids=A_    ,							encoder_hidden_states=A_    ,							encoder_attention_mask=A_    ,							)
      _SCREAMING_SNAKE_CASE        =			model(
          A_    ,							attention_mask=A_    ,							token_type_ids=A_    ,							encoder_hidden_states=A_    ,							)
      _SCREAMING_SNAKE_CASE        =			model(A_    ,							attention_mask=A_    ,							token_type_ids=A_			)
      self.parent.assertEqual(result.last_hidden_state.shape    ,							(self.batch_size, self.seq_length, self.hidden_size)			)




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			YosoForMaskedLM(config=A_			)
      model.to(A_			)
      model.eval()
      _SCREAMING_SNAKE_CASE        =			model(A_    ,							attention_mask=A_    ,							token_type_ids=A_    ,							labels=A_			)
      self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, self.seq_length, self.vocab_size)			)




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			YosoForQuestionAnswering(config=A_			)
      model.to(A_			)
      model.eval()
      _SCREAMING_SNAKE_CASE        =			model(
          A_    ,							attention_mask=A_    ,							token_type_ids=A_    ,							start_positions=A_    ,							end_positions=A_    ,							)
      self.parent.assertEqual(result.start_logits.shape    ,							(self.batch_size, self.seq_length)			)
      self.parent.assertEqual(result.end_logits.shape    ,							(self.batch_size, self.seq_length)			)




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.num_labels
      _SCREAMING_SNAKE_CASE        =			YosoForSequenceClassification(A_			)
      model.to(A_			)
      model.eval()
      _SCREAMING_SNAKE_CASE        =			model(A_    ,							attention_mask=A_    ,							token_type_ids=A_    ,							labels=A_			)
      self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, self.num_labels)			)




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.num_labels
      _SCREAMING_SNAKE_CASE        =			YosoForTokenClassification(config=A_			)
      model.to(A_			)
      model.eval()
      _SCREAMING_SNAKE_CASE        =			model(A_    ,							attention_mask=A_    ,							token_type_ids=A_    ,							labels=A_			)
      self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, self.seq_length, self.num_labels)			)




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.num_choices
      _SCREAMING_SNAKE_CASE        =			YosoForMultipleChoice(config=A_			)
      model.to(A_			)
      model.eval()
      _SCREAMING_SNAKE_CASE        =			input_ids.unsqueeze(1			).expand(-1    ,							self.num_choices    ,							-1			).contiguous()
      _SCREAMING_SNAKE_CASE        =			token_type_ids.unsqueeze(1			).expand(-1    ,							self.num_choices    ,							-1			).contiguous()
      _SCREAMING_SNAKE_CASE        =			input_mask.unsqueeze(1			).expand(-1    ,							self.num_choices    ,							-1			).contiguous()
      _SCREAMING_SNAKE_CASE        =			model(
          A_    ,							attention_mask=A_    ,							token_type_ids=A_    ,							labels=A_    ,							)
      self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, self.num_choices)			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.prepare_config_and_inputs()
      (
          (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     
      )        =			config_and_inputs
      _SCREAMING_SNAKE_CASE        =			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
      return config, inputs_dict




@require_torch
class    __snake_case(					__A			,			__A			,			unittest.TestCase       ):
    _A       = (
        (
            YosoModel,
            YosoForMaskedLM,
            YosoForMultipleChoice,
            YosoForQuestionAnswering,
            YosoForSequenceClassification,
            YosoForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    _A       = False
    _A       = False
    _A       = False

    _A       = ()
    _A       = (
        {
            '''feature-extraction''': YosoModel,
            '''fill-mask''': YosoForMaskedLM,
            '''question-answering''': YosoForQuestionAnswering,
            '''text-classification''': YosoForSequenceClassification,
            '''token-classification''': YosoForTokenClassification,
            '''zero-shot''': YosoForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			YosoModelTester(self			)
      _SCREAMING_SNAKE_CASE        =			ConfigTester(self    ,							config_class=A_    ,							hidden_size=37			)




    def   A   (   self			):





      '''simple docstring'''



      self.config_tester.run_common_tests()




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*A_			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
        _SCREAMING_SNAKE_CASE        =			type
        self.model_tester.create_and_check_model(*A_			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*A_			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_multiple_choice(*A_			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*A_			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*A_			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*A_			)




    @slow
    def   A   (   self			):





      '''simple docstring'''



      for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _SCREAMING_SNAKE_CASE        =			YosoModel.from_pretrained(A_			)
        self.assertIsNotNone(A_			)




    def   A   (   self			):





      '''simple docstring'''



      return




@require_torch
class    __snake_case(					unittest.TestCase       ):




    @slow
    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			YosoModel.from_pretrained('''uw-madison/yoso-4096'''			)
      _SCREAMING_SNAKE_CASE        =			torch.tensor([[0, 1, 2, 3, 4, 5]]			)

      with torch.no_grad():
        _SCREAMING_SNAKE_CASE        =			model(A_			)[0]

      _SCREAMING_SNAKE_CASE        =			torch.Size((1, 6, 768)			)
      self.assertEqual(output.shape    ,							A_			)

      _SCREAMING_SNAKE_CASE        =			torch.tensor(
          [[[-0.0_611, 0.1_242, 0.0_840], [0.0_280, -0.0_048, 0.1_125], [0.0_106, 0.0_226, 0.0_751]]]			)

      self.assertTrue(torch.allclose(output[:, :3, :3]    ,							A_    ,							atol=1e-4			)			)




    @slow
    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			YosoForMaskedLM.from_pretrained('''uw-madison/yoso-4096'''			)
      _SCREAMING_SNAKE_CASE        =			torch.tensor([[0, 1, 2, 3, 4, 5]]			)

      with torch.no_grad():
        _SCREAMING_SNAKE_CASE        =			model(A_			)[0]

      _SCREAMING_SNAKE_CASE        =			50_265

      _SCREAMING_SNAKE_CASE        =			torch.Size((1, 6, vocab_size)			)
      self.assertEqual(output.shape    ,							A_			)

      _SCREAMING_SNAKE_CASE        =			torch.tensor(
          [[[-2.1_313, -3.7_285, -2.2_407], [-2.7_047, -3.3_314, -2.6_408], [0.0_629, -2.5_166, -0.3_356]]]			)

      self.assertTrue(torch.allclose(output[:, :3, :3]    ,							A_    ,							atol=1e-4			)			)




    @slow
    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			YosoForMaskedLM.from_pretrained('''uw-madison/yoso-4096'''			)
      _SCREAMING_SNAKE_CASE        =			torch.arange(4_096			).unsqueeze(0			)

      with torch.no_grad():
        _SCREAMING_SNAKE_CASE        =			model(A_			)[0]

      _SCREAMING_SNAKE_CASE        =			50_265

      _SCREAMING_SNAKE_CASE        =			torch.Size((1, 4_096, vocab_size)			)
      self.assertEqual(output.shape    ,							A_			)

      _SCREAMING_SNAKE_CASE        =			torch.tensor(
          [[[-2.3_914, -4.3_742, -5.0_956], [-4.0_988, -4.2_384, -7.0_406], [-3.1_427, -3.7_192, -6.6_800]]]			)

      self.assertTrue(torch.allclose(output[:, :3, :3]    ,							A_    ,							atol=1e-4			)			)





"""simple docstring"""






import unittest

import numpy as np
import torch

from diffusers import DDPMPipeline, DDPMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device


enable_full_determinism()




class    __snake_case(					unittest.TestCase       ):




    @property
    def   A   (   self			):





      '''simple docstring'''



      torch.manual_seed(0			)
      _SCREAMING_SNAKE_CASE        =			UNetaDModel(
          block_out_channels=(32, 64)    ,							layers_per_block=2    ,							sample_size=32    ,							in_channels=3    ,							out_channels=3    ,							down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')    ,							up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')    ,							)
      return model




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			'''cpu'''
      _SCREAMING_SNAKE_CASE        =			self.dummy_uncond_unet
      _SCREAMING_SNAKE_CASE        =			DDPMScheduler()

      _SCREAMING_SNAKE_CASE        =			DDPMPipeline(unet=A_    ,							scheduler=A_			)
      ddpm.to(A_			)
      ddpm.set_progress_bar_config(disable=A_			)

      _SCREAMING_SNAKE_CASE        =			torch.Generator(device=A_			).manual_seed(0			)
      _SCREAMING_SNAKE_CASE        =			ddpm(generator=A_    ,							num_inference_steps=2    ,							output_type='''numpy'''			).images

      _SCREAMING_SNAKE_CASE        =			torch.Generator(device=A_			).manual_seed(0			)
      _SCREAMING_SNAKE_CASE        =			ddpm(generator=A_    ,							num_inference_steps=2    ,							output_type='''numpy'''    ,							return_dict=A_			)[0]

      _SCREAMING_SNAKE_CASE        =			image[0, -3:, -3:, -1]
      _SCREAMING_SNAKE_CASE        =			image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      _SCREAMING_SNAKE_CASE        =			np.array(
          [9.956e-01, 5.785e-01, 4.675e-01, 9.930e-01, 0.0, 1.000, 1.199e-03, 2.648e-04, 5.101e-04]			)

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




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.dummy_uncond_unet
      _SCREAMING_SNAKE_CASE        =			DDPMScheduler(prediction_type='''sample'''			)

      _SCREAMING_SNAKE_CASE        =			DDPMPipeline(unet=A_    ,							scheduler=A_			)
      ddpm.to(A_			)
      ddpm.set_progress_bar_config(disable=A_			)

      _SCREAMING_SNAKE_CASE        =			torch.manual_seed(0			)
      _SCREAMING_SNAKE_CASE        =			ddpm(generator=A_    ,							num_inference_steps=2    ,							output_type='''numpy'''			).images

      _SCREAMING_SNAKE_CASE        =			torch.manual_seed(0			)
      _SCREAMING_SNAKE_CASE        =			ddpm(generator=A_    ,							num_inference_steps=2    ,							output_type='''numpy'''			)[0]

      _SCREAMING_SNAKE_CASE        =			image[0, -3:, -3:, -1]
      _SCREAMING_SNAKE_CASE        =			image_eps[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      _SCREAMING_SNAKE_CASE        =			1e-2 if torch_device != '''mps''' else 3e-2
      assert np.abs(image_slice.flatten() - image_eps_slice.flatten()			).max() < tolerance




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




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			'''google/ddpm-cifar10-32'''

      _SCREAMING_SNAKE_CASE        =			UNetaDModel.from_pretrained(A_			)
      _SCREAMING_SNAKE_CASE        =			DDPMScheduler.from_pretrained(A_			)

      _SCREAMING_SNAKE_CASE        =			DDPMPipeline(unet=A_    ,							scheduler=A_			)
      ddpm.to(A_			)
      ddpm.set_progress_bar_config(disable=A_			)

      _SCREAMING_SNAKE_CASE        =			torch.manual_seed(0			)
      _SCREAMING_SNAKE_CASE        =			ddpm(generator=A_    ,							output_type='''numpy'''			).images

      _SCREAMING_SNAKE_CASE        =			image[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      _SCREAMING_SNAKE_CASE        =			np.array([0.4_200, 0.3_588, 0.1_939, 0.3_847, 0.3_382, 0.2_647, 0.4_155, 0.3_582, 0.3_385]			)
      assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-2
