

import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
  import torch
  from torch import nn

  from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
  from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import ViTImageProcessor







class        a__       :







  """simple docstring"""




  def __init__(			self    : List[Any]				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : Union[str, Any]=1_3				,				UpperCAmelCase__    : Union[str, Any]=3_0				,				UpperCAmelCase__    : List[str]=2				,				UpperCAmelCase__    : int=3				,				UpperCAmelCase__    : Union[str, Any]=True				,				UpperCAmelCase__    : int=True				,				UpperCAmelCase__    : Optional[int]=3_2				,				UpperCAmelCase__    : Union[str, Any]=5				,				UpperCAmelCase__    : Optional[int]=4				,				UpperCAmelCase__    : Dict=3_7				,				UpperCAmelCase__    : Union[str, Any]="gelu"				,				UpperCAmelCase__    : Tuple=0.1				,				UpperCAmelCase__    : Union[str, Any]=0.1				,				UpperCAmelCase__    : Any=1_0				,				UpperCAmelCase__    : Dict=0.02				,				UpperCAmelCase__    : Union[str, Any]=None				,				UpperCAmelCase__    : Tuple=2				,				)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Tuple       =   parent
        SCREAMING_SNAKE_CASE :      Dict       =   batch_size
        SCREAMING_SNAKE_CASE :      Tuple       =   image_size
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   patch_size
        SCREAMING_SNAKE_CASE :      Dict       =   num_channels
        SCREAMING_SNAKE_CASE :      Optional[int]       =   is_training
        SCREAMING_SNAKE_CASE :      str       =   use_labels
        SCREAMING_SNAKE_CASE :      int       =   hidden_size
        SCREAMING_SNAKE_CASE :      Any       =   num_hidden_layers
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   num_attention_heads
        SCREAMING_SNAKE_CASE :      str       =   intermediate_size
        SCREAMING_SNAKE_CASE :      Tuple       =   hidden_act
        SCREAMING_SNAKE_CASE :      str       =   hidden_dropout_prob
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   attention_probs_dropout_prob
        SCREAMING_SNAKE_CASE :      Any       =   type_sequence_label_size
        SCREAMING_SNAKE_CASE :      Any       =   initializer_range
        SCREAMING_SNAKE_CASE :      str       =   scope
        SCREAMING_SNAKE_CASE :      Tuple       =   encoder_stride

        # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
        SCREAMING_SNAKE_CASE :      Dict       =   (image_size // patch_size) ** 2
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   num_patches + 1


  def  _lowercase      (			self    : str		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      str       =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

        SCREAMING_SNAKE_CASE :      List[str]       =   None
        if self.use_labels:
              SCREAMING_SNAKE_CASE :      Dict       =   ids_tensor([self.batch_size]				,				self.type_sequence_label_size		)

        SCREAMING_SNAKE_CASE :      Any       =   self.get_config()

        return config, pixel_values, labels


  def  _lowercase      (			self    : Any		)					->str:
        """simple docstring"""







        return ViTConfig(
            image_size=self.image_size				,				patch_size=self.patch_size				,				num_channels=self.num_channels				,				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				,				is_decoder=UpperCAmelCase__				,				initializer_range=self.initializer_range				,				encoder_stride=self.encoder_stride				,				)


  def  _lowercase      (			self    : List[str]				,				UpperCAmelCase__    : int				,				UpperCAmelCase__    : List[Any]				,				UpperCAmelCase__    : Dict		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      int       =   ViTModel(config=UpperCAmelCase__		)
        model.to(UpperCAmelCase__		)
        model.eval()
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   model(UpperCAmelCase__		)
        self.parent.assertEqual(result.last_hidden_state.shape				,				(self.batch_size, self.seq_length, self.hidden_size)		)


  def  _lowercase      (			self    : List[str]				,				UpperCAmelCase__    : str				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : str		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   ViTForMaskedImageModeling(config=UpperCAmelCase__		)
        model.to(UpperCAmelCase__		)
        model.eval()
        SCREAMING_SNAKE_CASE :      Optional[int]       =   model(UpperCAmelCase__		)
        self.parent.assertEqual(
            result.reconstruction.shape				,				(self.batch_size, self.num_channels, self.image_size, self.image_size)		)

        # test greyscale images
        SCREAMING_SNAKE_CASE :      str       =   1
        SCREAMING_SNAKE_CASE :      Tuple       =   ViTForMaskedImageModeling(UpperCAmelCase__		)
        model.to(UpperCAmelCase__		)
        model.eval()

        SCREAMING_SNAKE_CASE :      Optional[int]       =   floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
        SCREAMING_SNAKE_CASE :      Tuple       =   model(UpperCAmelCase__		)
        self.parent.assertEqual(result.reconstruction.shape				,				(self.batch_size, 1, self.image_size, self.image_size)		)


  def  _lowercase      (			self    : Optional[Any]				,				UpperCAmelCase__    : List[Any]				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : int		)					->int:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   self.type_sequence_label_size
        SCREAMING_SNAKE_CASE :      str       =   ViTForImageClassification(UpperCAmelCase__		)
        model.to(UpperCAmelCase__		)
        model.eval()
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   model(UpperCAmelCase__				,				labels=UpperCAmelCase__		)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.type_sequence_label_size)		)

        # test greyscale images
        SCREAMING_SNAKE_CASE :      int       =   1
        SCREAMING_SNAKE_CASE :      Dict       =   ViTForImageClassification(UpperCAmelCase__		)
        model.to(UpperCAmelCase__		)
        model.eval()

        SCREAMING_SNAKE_CASE :      str       =   floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
        SCREAMING_SNAKE_CASE :      Optional[int]       =   model(UpperCAmelCase__		)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.type_sequence_label_size)		)




  def  _lowercase      (			self    : Optional[Any]		)					->Optional[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[int]       =   self.prepare_config_and_inputs()
        (
            (
            SCREAMING_SNAKE_CASE
        )     ,	(
            SCREAMING_SNAKE_CASE
        )     ,	(
            SCREAMING_SNAKE_CASE
        )     ,	
        ) :      Tuple       =   config_and_inputs
        SCREAMING_SNAKE_CASE :      Optional[int]       =   {"""pixel_values""": pixel_values}
        return config, inputs_dict







@require_torch
class        a__       (   UpperCAmelCase	,			UpperCAmelCase	,			unittest.TestCase       ):







  """simple docstring"""




  UpperCAmelCase__      : str			       =(
      (
          ViTModel,
          ViTForImageClassification,
          ViTForMaskedImageModeling,
      )
      if is_torch_available()
      else ()
  )
  UpperCAmelCase__      : Dict			       =(
      {"""feature-extraction""": ViTModel, """image-classification""": ViTForImageClassification}
      if is_torch_available()
      else {}
  )
  UpperCAmelCase__      : int			       =True

  UpperCAmelCase__      : str			       =False
  UpperCAmelCase__      : Any			       =False
  UpperCAmelCase__      : Optional[int]			       =False


  def  _lowercase      (			self    : Union[str, Any]		)					->List[str]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[int]       =   ViTModelTester(self		)
        SCREAMING_SNAKE_CASE :      Any       =   ConfigTester(self				,				config_class=UpperCAmelCase__				,				has_text_modality=UpperCAmelCase__				,				hidden_size=3_7		)


  def  _lowercase      (			self    : Dict		)					->Optional[int]:
        """simple docstring"""







        self.config_tester.run_common_tests()


  @unittest.skip(reason="""ViT does not use inputs_embeds"""		)
  def  _lowercase      (			self    : Tuple		)					->Tuple:
        """simple docstring"""







        pass


  def  _lowercase      (			self    : str		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      Tuple       =   self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              SCREAMING_SNAKE_CASE :      str       =   model_class(UpperCAmelCase__		)
              self.assertIsInstance(model.get_input_embeddings()				,				(nn.Module)		)
              SCREAMING_SNAKE_CASE :      List[Any]       =   model.get_output_embeddings()
              self.assertTrue(x is None or isinstance(UpperCAmelCase__				,				nn.Linear		)		)


  def  _lowercase      (			self    : List[str]		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      List[str]       =   self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              SCREAMING_SNAKE_CASE :      Any       =   model_class(UpperCAmelCase__		)
              SCREAMING_SNAKE_CASE :      int       =   inspect.signature(model.forward		)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              SCREAMING_SNAKE_CASE :      List[Any]       =   [*signature.parameters.keys()]

              SCREAMING_SNAKE_CASE :      int       =   ["""pixel_values"""]
              self.assertListEqual(arg_names[:1]				,				UpperCAmelCase__		)


  def  _lowercase      (			self    : Union[str, Any]		)					->List[str]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*UpperCAmelCase__		)


  def  _lowercase      (			self    : int		)					->Tuple:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[Any]       =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_image_modeling(*UpperCAmelCase__		)


  def  _lowercase      (			self    : Any		)					->int:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      str       =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase__		)




  @slow
  def  _lowercase      (			self    : Optional[int]		)					->str:
        """simple docstring"""







        for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              SCREAMING_SNAKE_CASE :      List[Any]       =   ViTModel.from_pretrained(UpperCAmelCase__		)
              self.assertIsNotNone(UpperCAmelCase__		)






def 					__lowercase      (     )     ->						Optional[Any]:
      SCREAMING_SNAKE_CASE :      Optional[int]       =   Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""			)
      return image







@require_torch
@require_vision
class        a__       (   unittest.TestCase       ):







  """simple docstring"""




  @cached_property
  def  _lowercase      (			self    : str		)					->Tuple:
        """simple docstring"""







        return ViTImageProcessor.from_pretrained("""google/vit-base-patch16-224"""		) if is_vision_available() else None


  @slow
  def  _lowercase      (			self    : Dict		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Tuple       =   ViTForImageClassification.from_pretrained("""google/vit-base-patch16-224"""		).to(UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   self.default_image_processor
        SCREAMING_SNAKE_CASE :      Optional[int]       =   prepare_img()
        SCREAMING_SNAKE_CASE :      Dict       =   image_processor(images=UpperCAmelCase__				,				return_tensors="""pt"""		).to(UpperCAmelCase__		)

        # forward pass
        with torch.no_grad():
              SCREAMING_SNAKE_CASE :      Optional[Any]       =   model(**UpperCAmelCase__		)

        # verify the logits
        SCREAMING_SNAKE_CASE :      Optional[int]       =   torch.Size((1, 1_0_0_0)		)
        self.assertEqual(outputs.logits.shape				,				UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      int       =   torch.tensor([-0.27_44, 0.82_15, -0.08_36]		).to(UpperCAmelCase__		)

        self.assertTrue(torch.allclose(outputs.logits[0, :3]				,				UpperCAmelCase__				,				atol=1e-4		)		)


  @slow
  def  _lowercase      (			self    : Any		)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[str]       =   ViTModel.from_pretrained("""facebook/dino-vits8"""		).to(UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Optional[Any]       =   ViTImageProcessor.from_pretrained("""facebook/dino-vits8"""				,				size=4_8_0		)
        SCREAMING_SNAKE_CASE :      Any       =   prepare_img()
        SCREAMING_SNAKE_CASE :      List[str]       =   image_processor(images=UpperCAmelCase__				,				return_tensors="""pt"""		)
        SCREAMING_SNAKE_CASE :      int       =   inputs.pixel_values.to(UpperCAmelCase__		)

        # forward pass
        with torch.no_grad():
              SCREAMING_SNAKE_CASE :      int       =   model(UpperCAmelCase__				,				interpolate_pos_encoding=UpperCAmelCase__		)

        # verify the logits
        SCREAMING_SNAKE_CASE :      List[str]       =   torch.Size((1, 3_6_0_1, 3_8_4)		)
        self.assertEqual(outputs.last_hidden_state.shape				,				UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Any       =   torch.tensor(
            [[4.23_40, 4.39_06, -6.66_92], [4.54_63, 1.89_28, -6.72_57], [4.44_29, 0.84_96, -5.85_85]]		).to(UpperCAmelCase__		)

        self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]				,				UpperCAmelCase__				,				atol=1e-4		)		)




  @slow
  @require_accelerate
  @require_torch_gpu
  def  _lowercase      (			self    : Tuple		)					->str:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[int]       =   ViTModel.from_pretrained("""facebook/dino-vits8"""				,				torch_dtype=torch.floataa				,				device_map="""auto"""		)
        SCREAMING_SNAKE_CASE :      List[str]       =   self.default_image_processor

        SCREAMING_SNAKE_CASE :      str       =   prepare_img()
        SCREAMING_SNAKE_CASE :      List[Any]       =   image_processor(images=UpperCAmelCase__				,				return_tensors="""pt"""		)
        SCREAMING_SNAKE_CASE :      List[Any]       =   inputs.pixel_values.to(UpperCAmelCase__		)

        # forward pass to make sure inference works in fp16
        with torch.no_grad():
              SCREAMING_SNAKE_CASE :      List[str]       =   model(UpperCAmelCase__		)







import os
import random
import sys

from . import cryptomath_module as cryptomath
from . import rabin_miller

UpperCAmelCase__ :     Optional[Any]				    =    3






def 					__lowercase      (     _A			)     ->						int:
      print("""Generating primitive root of p"""			)
      while True:
            SCREAMING_SNAKE_CASE :      int       =   random.randrange(3						,    _A			)
            if pow(_A						,    2						,    _A			) == 1:
                  continue
            if pow(_A						,    _A						,    _A			) == 1:
                  continue
            return g






def 					__lowercase      (     _A			)     ->						tuple[tuple[int, int, int, int], tuple[int, int]]:
      print("""Generating prime p..."""			)
      SCREAMING_SNAKE_CASE :      Optional[Any]       =   rabin_miller.generate_large_prime(_A			)  # select large prime number.
      SCREAMING_SNAKE_CASE :      Optional[Any]       =   primitive_root(_A			)  # one primitive root on modulo p.
      SCREAMING_SNAKE_CASE :      Optional[Any]       =   random.randrange(3						,    _A			)  # private_key -> have to be greater than 2 for safety.
      SCREAMING_SNAKE_CASE :      List[str]       =   cryptomath.find_mod_inverse(pow(_A						,    _A						,    _A			)						,    _A			)

      SCREAMING_SNAKE_CASE :      str       =   (key_size, e_a, e_a, p)
      SCREAMING_SNAKE_CASE :      Union[str, Any]       =   (key_size, d)

      return public_key, private_key






def 					__lowercase      (     _A						,    _A			)     ->						None:
      if os.path.exists(F"{name}_pubkey.txt"			) or os.path.exists(F"{name}_privkey.txt"			):
            print("""\nWARNING:"""			)
            print(
                F"\"{name}_pubkey.txt\" or \"{name}_privkey.txt\" already exists. \n"
                """Use a different name or delete these files and re-run this program."""			)
            sys.exit()

      SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      List[str]       =   generate_key(_A			)
      print(F"\nWriting public key to file {name}_pubkey.txt..."			)
      with open(F"{name}_pubkey.txt"						,    """w"""			) as fo:
            fo.write(F"{public_key[0]},{public_key[1]},{public_key[2]},{public_key[3]}"			)

      print(F"Writing private key to file {name}_privkey.txt..."			)
      with open(F"{name}_privkey.txt"						,    """w"""			) as fo:
            fo.write(F"{private_key[0]},{private_key[1]}"			)






def 					__lowercase      (     )     ->						None:
      print("""Making key files..."""			)
      make_key_files("""elgamal"""						,    2048			)
      print("""Key files generation successful"""			)


if __name__ == "__main__":
  main()




