

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      UpperCamelCase :







		"""simple docstring"""



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


						'''simple docstring'''
						A	 =      parent
						A	 =      batch_size
						A	 =      image_size
						A	 =      patch_size
						A	 =      num_channels
						A	 =      is_training
						A	 =      use_labels
						A	 =      hidden_size
						A	 =      num_hidden_layers
						A	 =      num_attention_heads
						A	 =      intermediate_size
						A	 =      hidden_act
						A	 =      hidden_dropout_prob
						A	 =      attention_probs_dropout_prob
						A	 =      type_sequence_label_size
						A	 =      initializer_range
						A	 =      scope
						A	 =      encoder_stride

						# in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
						A	 =      (image_size // patch_size) ** 2
						A	 =      num_patches + 1
		def       A(  self		:       Any		)    ->   int:


						'''simple docstring'''
						A	 =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

						A	 =      None
						if self.use_labels:
										A	 =      ids_tensor([self.batch_size]						,self.type_sequence_label_size		)

						A	 =      self.get_config()

						return config, pixel_values, labels
		def       A(  self		:       List[Any]		)    ->   Optional[Any]:


						'''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=_a						,initializer_range=self.initializer_range						,encoder_stride=self.encoder_stride						,)
		def       A(  self		:       Dict						,_SCREAMING_SNAKE_CASE		:       Optional[int]						,_SCREAMING_SNAKE_CASE		:       List[Any]						,_SCREAMING_SNAKE_CASE		:       Any		)    ->   Union[str, Any]:


						'''simple docstring'''
						A	 =      ViTModel(config=_a		)
						model.to(_a		)
						model.eval()
						A	 =      model(_a		)
						self.parent.assertEqual(result.last_hidden_state.shape						,(self.batch_size, self.seq_length, self.hidden_size)		)
		def       A(  self		:       List[Any]						,_SCREAMING_SNAKE_CASE		:       List[str]						,_SCREAMING_SNAKE_CASE		:       Any						,_SCREAMING_SNAKE_CASE		:       Optional[int]		)    ->   int:


						'''simple docstring'''
						A	 =      ViTForMaskedImageModeling(config=_a		)
						model.to(_a		)
						model.eval()
						A	 =      model(_a		)
						self.parent.assertEqual(
						    result.reconstruction.shape						,(self.batch_size, self.num_channels, self.image_size, self.image_size)		)

						# test greyscale images
						A	 =      1
						A	 =      ViTForMaskedImageModeling(_a		)
						model.to(_a		)
						model.eval()

						A	 =      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
						A	 =      model(_a		)
						self.parent.assertEqual(result.reconstruction.shape						,(self.batch_size, 1, self.image_size, self.image_size)		)
		def       A(  self		:       List[Any]						,_SCREAMING_SNAKE_CASE		:       str						,_SCREAMING_SNAKE_CASE		:       Optional[int]						,_SCREAMING_SNAKE_CASE		:       Dict		)    ->   Optional[Any]:


						'''simple docstring'''
						A	 =      self.type_sequence_label_size
						A	 =      ViTForImageClassification(_a		)
						model.to(_a		)
						model.eval()
						A	 =      model(_a						,labels=_a		)
						self.parent.assertEqual(result.logits.shape						,(self.batch_size, self.type_sequence_label_size)		)

						# test greyscale images
						A	 =      1
						A	 =      ViTForImageClassification(_a		)
						model.to(_a		)
						model.eval()

						A	 =      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
						A	 =      model(_a		)
						self.parent.assertEqual(result.logits.shape						,(self.batch_size, self.type_sequence_label_size)		)







		def       A(  self		:       Optional[int]		)    ->   Any:


						'''simple docstring'''
						A	 =      self.prepare_config_and_inputs()
						(
						    A					
						)	 =      config_and_inputs
						A	 =      {"""pixel_values""": pixel_values}
						return config, inputs_dict


@require_torch
class      UpperCamelCase (				__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       unittest.TestCase ):







		"""simple docstring"""



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

		snake_case    =   False
		snake_case    =   False
		snake_case    =   False
		def       A(  self		:       int		)    ->   Optional[int]:


						'''simple docstring'''
						A	 =      ViTModelTester(self		)
						A	 =      ConfigTester(self						,config_class=_a						,has_text_modality=_a						,hidden_size=3_7		)
		def       A(  self		:       Union[str, Any]		)    ->   Any:


						'''simple docstring'''
						self.config_tester.run_common_tests()
		@unittest.skip(reason='ViT does not use inputs_embeds'		)
		def       A(  self		:       int		)    ->   List[Any]:


						'''simple docstring'''
						pass
		def       A(  self		:       int		)    ->   Union[str, Any]:


						'''simple docstring'''
						A	 =      self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										A	 =      model_class(_a		)
										self.assertIsInstance(model.get_input_embeddings()						,(nn.Module)		)
										A	 =      model.get_output_embeddings()
										self.assertTrue(x is None or isinstance(_a						,nn.Linear		)		)
		def       A(  self		:       Optional[Any]		)    ->   Any:


						'''simple docstring'''
						A	 =      self.model_tester.prepare_config_and_inputs_for_common()

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

										A	 =      ["""pixel_values"""]
										self.assertListEqual(arg_names[:1]						,_a		)
		def       A(  self		:       Optional[int]		)    ->   List[str]:


						'''simple docstring'''
						A	 =      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*_a		)
		def       A(  self		:       Union[str, Any]		)    ->   str:


						'''simple docstring'''
						A	 =      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_image_modeling(*_a		)
		def       A(  self		:       str		)    ->   Tuple:


						'''simple docstring'''
						A	 =      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*_a		)







		@slow
		def       A(  self		:       Optional[Any]		)    ->   List[Any]:


						'''simple docstring'''




						for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										A	 =      ViTModel.from_pretrained(_a		)
										self.assertIsNotNone(_a		)







def      snake_case		(							):
				A	 =      Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
				return image



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







		"""simple docstring"""



		@cached_property
		def       A(  self		:       Optional[int]		)    ->   Tuple:


						'''simple docstring'''
						return ViTImageProcessor.from_pretrained('google/vit-base-patch16-224'		) if is_vision_available() else None
		@slow
		def       A(  self		:       str		)    ->   Tuple:


						'''simple docstring'''
						A	 =      ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'		).to(_a		)

						A	 =      self.default_image_processor
						A	 =      prepare_img()
						A	 =      image_processor(images=_a						,return_tensors='pt'		).to(_a		)

						# forward pass
						with torch.no_grad():
										A	 =      model(**_a		)

						# verify the logits
						A	 =      torch.Size((1, 1_0_0_0)		)
						self.assertEqual(outputs.logits.shape						,_a		)

						A	 =      torch.tensor([-0.2744, 0.8215, -0.0836]		).to(_a		)

						self.assertTrue(torch.allclose(outputs.logits[0, :3]						,_a						,atol=1E-4		)		)
		@slow
		def       A(  self		:       List[str]		)    ->   List[str]:


						'''simple docstring'''
						A	 =      ViTModel.from_pretrained('facebook/dino-vits8'		).to(_a		)

						A	 =      ViTImageProcessor.from_pretrained('facebook/dino-vits8'						,size=4_8_0		)
						A	 =      prepare_img()
						A	 =      image_processor(images=_a						,return_tensors='pt'		)
						A	 =      inputs.pixel_values.to(_a		)

						# forward pass
						with torch.no_grad():
										A	 =      model(_a						,interpolate_pos_encoding=_a		)

						# verify the logits
						A	 =      torch.Size((1, 3_6_0_1, 3_8_4)		)
						self.assertEqual(outputs.last_hidden_state.shape						,_a		)

						A	 =      torch.tensor(
						    [[4.2340, 4.3906, -6.6692], [4.5463, 1.8928, -6.7257], [4.4429, 0.8496, -5.8585]]		).to(_a		)

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







		@slow
		@require_accelerate
		@require_torch_gpu
		def       A(  self		:       Tuple		)    ->   Optional[int]:


						'''simple docstring'''
						A	 =      ViTModel.from_pretrained('facebook/dino-vits8'						,torch_dtype=torch.floataa						,device_map='auto'		)
						A	 =      self.default_image_processor

						A	 =      prepare_img()
						A	 =      image_processor(images=_a						,return_tensors='pt'		)
						A	 =      inputs.pixel_values.to(_a		)

						# forward pass to make sure inference works in fp16
						with torch.no_grad():
										A	 =      model(_a		)













from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin
class      UpperCamelCase (				snake_case__			,       snake_case__			,       snake_case__ ):







		"""simple docstring"""



		snake_case    =   [R"h\.\d+\.attn\.bias", R"h\.\d+\.attn\.masked_bias"]
		@register_to_config
		def __init__(  self		:       Any						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       Optional[int] = None						,_SCREAMING_SNAKE_CASE		:       int = 5_0_2_5_7						,_SCREAMING_SNAKE_CASE		:       int = 1_0_2_4						,_SCREAMING_SNAKE_CASE		:       int = 7_6_8						,_SCREAMING_SNAKE_CASE		:       int = 1_2						,_SCREAMING_SNAKE_CASE		:       int = 1_2						,_SCREAMING_SNAKE_CASE		:       Optional[int] = None						,_SCREAMING_SNAKE_CASE		:       str = "gelu_new"						,_SCREAMING_SNAKE_CASE		:       float = 0.1						,_SCREAMING_SNAKE_CASE		:       float = 0.1						,_SCREAMING_SNAKE_CASE		:       float = 0.1						,_SCREAMING_SNAKE_CASE		:       float = 1E-5						,_SCREAMING_SNAKE_CASE		:       float = 0.02						,_SCREAMING_SNAKE_CASE		:       bool = True						,_SCREAMING_SNAKE_CASE		:       bool = True						,_SCREAMING_SNAKE_CASE		:       bool = False						,_SCREAMING_SNAKE_CASE		:       bool = False						,)    ->   Tuple:


						'''simple docstring'''
						super().__init__()

						A	 =      prefix_length

						if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
										raise ValueError(
										    f'`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and'
										    f' `n_embd`: {n_embd} are not equal.'		)

						A	 =      prefix_inner_dim
						A	 =      prefix_hidden_dim

						A	 =      (
						    nn.Linear(self.prefix_inner_dim						,self.prefix_hidden_dim		)
						    if self.prefix_hidden_dim is not None
						    else nn.Identity()
						)
						A	 =      (
						    nn.Linear(self.prefix_hidden_dim						,_SCREAMING_SNAKE_CASE		) if self.prefix_hidden_dim is not None else nn.Identity()
						)

						A	 =      GPTaConfig(
						    vocab_size=_SCREAMING_SNAKE_CASE						,n_positions=_SCREAMING_SNAKE_CASE						,n_embd=_SCREAMING_SNAKE_CASE						,n_layer=_SCREAMING_SNAKE_CASE						,n_head=_SCREAMING_SNAKE_CASE						,n_inner=_SCREAMING_SNAKE_CASE						,activation_function=_SCREAMING_SNAKE_CASE						,resid_pdrop=_SCREAMING_SNAKE_CASE						,embd_pdrop=_SCREAMING_SNAKE_CASE						,attn_pdrop=_SCREAMING_SNAKE_CASE						,layer_norm_epsilon=_SCREAMING_SNAKE_CASE						,initializer_range=_SCREAMING_SNAKE_CASE						,scale_attn_weights=_SCREAMING_SNAKE_CASE						,use_cache=_SCREAMING_SNAKE_CASE						,scale_attn_by_inverse_layer_idx=_SCREAMING_SNAKE_CASE						,reorder_and_upcast_attn=_SCREAMING_SNAKE_CASE						,)
						A	 =      GPTaLMHeadModel(_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       Optional[int]						,_SCREAMING_SNAKE_CASE		:       torch.Tensor						,_SCREAMING_SNAKE_CASE		:       torch.Tensor						,_SCREAMING_SNAKE_CASE		:       Optional[torch.Tensor] = None						,_SCREAMING_SNAKE_CASE		:       Optional[torch.Tensor] = None						,)    ->   Tuple:


						'''simple docstring'''
						A	 =      self.transformer.transformer.wte(_SCREAMING_SNAKE_CASE		)
						A	 =      self.encode_prefix(_SCREAMING_SNAKE_CASE		)
						A	 =      self.decode_prefix(_SCREAMING_SNAKE_CASE		)
						A	 =      torch.cat((prefix_embeds, embedding_text)						,dim=1		)

						if labels is not None:
										A	 =      self.get_dummy_token(input_ids.shape[0]						,input_ids.device		)
										A	 =      torch.cat((dummy_token, input_ids)						,dim=1		)
						A	 =      self.transformer(inputs_embeds=_SCREAMING_SNAKE_CASE						,labels=_SCREAMING_SNAKE_CASE						,attention_mask=_SCREAMING_SNAKE_CASE		)
						if self.prefix_hidden_dim is not None:
										return out, hidden
						else:
										return out
		def       A(  self		:       Any						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       torch.device		)    ->   torch.Tensor:


						'''simple docstring'''
						return torch.zeros(_SCREAMING_SNAKE_CASE						,self.prefix_length						,dtype=torch.intaa						,device=_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       List[str]						,_SCREAMING_SNAKE_CASE		:       int		)    ->   Optional[Any]:


						'''simple docstring'''
						return self.encode_prefix(_SCREAMING_SNAKE_CASE		)
		@torch.no_grad()
		def       A(  self		:       List[str]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       Union[str, Any]		)    ->   str:


						'''simple docstring'''
						A	 =      torch.split(_SCREAMING_SNAKE_CASE						,1						,dim=0		)
						A	 =      []
						A	 =      []
						for feature in features:
										A	 =      self.decode_prefix(feature.to(_SCREAMING_SNAKE_CASE		)		)  # back to the clip feature
										# Only support beam search for now
										A ,		A	 =      self.generate_beam(
										    input_embeds=_SCREAMING_SNAKE_CASE						,device=_SCREAMING_SNAKE_CASE						,eos_token_id=_SCREAMING_SNAKE_CASE		)
										generated_tokens.append(output_tokens[0]		)
										generated_seq_lengths.append(seq_lengths[0]		)
						A	 =      torch.stack(_SCREAMING_SNAKE_CASE		)
						A	 =      torch.stack(_SCREAMING_SNAKE_CASE		)
						return generated_tokens, generated_seq_lengths







		@torch.no_grad()
		def       A(  self		:       Dict						,_SCREAMING_SNAKE_CASE		:       List[str]=None						,_SCREAMING_SNAKE_CASE		:       str=None						,_SCREAMING_SNAKE_CASE		:       Any=None						,_SCREAMING_SNAKE_CASE		:       int = 5						,_SCREAMING_SNAKE_CASE		:       int = 6_7						,_SCREAMING_SNAKE_CASE		:       float = 1.0						,_SCREAMING_SNAKE_CASE		:       Optional[int] = None						,)    ->   str:


						'''simple docstring'''
						A	 =      eos_token_id
						A	 =      None
						A	 =      None
						A	 =      torch.ones(_SCREAMING_SNAKE_CASE						,device=_SCREAMING_SNAKE_CASE						,dtype=torch.int		)
						A	 =      torch.zeros(_SCREAMING_SNAKE_CASE						,device=_SCREAMING_SNAKE_CASE						,dtype=torch.bool		)

						if input_embeds is not None:
										A	 =      input_embeds
						else:
										A	 =      self.transformer.transformer.wte(_SCREAMING_SNAKE_CASE		)

						for i in range(_SCREAMING_SNAKE_CASE		):
										A	 =      self.transformer(inputs_embeds=_SCREAMING_SNAKE_CASE		)
										A	 =      outputs.logits
										A	 =      logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
										A	 =      logits.softmax(-1		).log()

										if scores is None:
														A ,		A	 =      logits.topk(_SCREAMING_SNAKE_CASE						,-1		)
														A	 =      generated.expand(_SCREAMING_SNAKE_CASE						,*generated.shape[1:]		)
														A ,		A	 =      next_tokens.permute(1						,0		), scores.squeeze(0		)
														if tokens is None:
																		A	 =      next_tokens
														else:
																		A	 =      tokens.expand(_SCREAMING_SNAKE_CASE						,*tokens.shape[1:]		)
																		A	 =      torch.cat((tokens, next_tokens)						,dim=1		)
										else:
														A	 =      -float(np.inf		)
														A	 =      0
														A	 =      scores[:, None] + logits
														seq_lengths[~is_stopped] += 1
														A	 =      scores_sum / seq_lengths[:, None]
														A ,		A	 =      scores_sum_average.view(-1		).topk(_SCREAMING_SNAKE_CASE						,-1		)
														A	 =      next_tokens // scores_sum.shape[1]
														A	 =      seq_lengths[next_tokens_source]
														A	 =      next_tokens % scores_sum.shape[1]
														A	 =      next_tokens.unsqueeze(1		)
														A	 =      tokens[next_tokens_source]
														A	 =      torch.cat((tokens, next_tokens)						,dim=1		)
														A	 =      generated[next_tokens_source]
														A	 =      scores_sum_average * seq_lengths
														A	 =      is_stopped[next_tokens_source]

										A	 =      self.transformer.transformer.wte(next_tokens.squeeze()		).view(generated.shape[0]						,1						,-1		)
										A	 =      torch.cat((generated, next_token_embed)						,dim=1		)
										A	 =      is_stopped + next_tokens.eq(_SCREAMING_SNAKE_CASE		).squeeze()
										if is_stopped.all():
														break

						A	 =      scores / seq_lengths
						A	 =      scores.argsort(descending=_SCREAMING_SNAKE_CASE		)
						# tokens tensors are already padded to max_seq_length
						A	 =      [tokens[i] for i in order]
						A	 =      torch.stack(_SCREAMING_SNAKE_CASE						,dim=0		)
						A	 =      torch.tensor([seq_lengths[i] for i in order]						,dtype=seq_lengths.dtype		)
						return output_texts, seq_lengths





