

import tempfile
import unittest

from transformers import TaConfig, is_torch_available
from transformers.testing_utils import (
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
    torch_device,
)

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


if is_torch_available():
					import torch

					from transformers import AutoTokenizer, UMTaForConditionalGeneration, UMTaForQuestionAnswering, UMTaModel

class 				a  :


						def __init__(       self						,					A_						,					A_=99						,					A_=13						,					A_=7						,					A_=9						,					A_=True						,					A_=True						,					A_=False						,					A_=32						,					A_=5						,					A_=4						,					A_=37						,					A_=8						,					A_=0.1						,					A_=0.0_02						,					A_=1						,					A_=0						,					A_=0						,					A_=None						,					A_=None						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		Any 		=    parent
										_UpperCAmelCase       :		Dict 		=    batch_size
										_UpperCAmelCase       :		Any 		=    encoder_seq_length
										_UpperCAmelCase       :		Tuple 		=    decoder_seq_length
										# For common tests
										_UpperCAmelCase       :		Any 		=    self.decoder_seq_length
										_UpperCAmelCase       :		List[Any] 		=    is_training
										_UpperCAmelCase       :		Dict 		=    use_attention_mask
										_UpperCAmelCase       :		List[Any] 		=    use_labels
										_UpperCAmelCase       :		str 		=    vocab_size
										_UpperCAmelCase       :		List[Any] 		=    hidden_size
										_UpperCAmelCase       :		Dict 		=    num_hidden_layers
										_UpperCAmelCase       :		Tuple 		=    num_attention_heads
										_UpperCAmelCase       :		Any 		=    d_ff
										_UpperCAmelCase       :		List[str] 		=    relative_attention_num_buckets
										_UpperCAmelCase       :		Optional[int] 		=    dropout_rate
										_UpperCAmelCase       :		List[Any] 		=    initializer_factor
										_UpperCAmelCase       :		List[Any] 		=    eos_token_id
										_UpperCAmelCase       :		Tuple 		=    pad_token_id
										_UpperCAmelCase       :		Union[str, Any] 		=    decoder_start_token_id
										_UpperCAmelCase       :		Union[str, Any] 		=    None
										_UpperCAmelCase       :		str 		=    decoder_layers


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return TaConfig.from_pretrained("google/umt5-base"	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_=None						,					A_=None						,					A_=None						,					A_=None						,					A_=None						,					):







										'''simple docstring'''


										if attention_mask is None:
														_UpperCAmelCase       :		Any 		=    input_ids.ne(config.pad_token_id	)
										if decoder_attention_mask is None:
														_UpperCAmelCase       :		List[Any] 		=    decoder_input_ids.ne(config.pad_token_id	)
										if head_mask is None:
														_UpperCAmelCase       :		List[str] 		=    torch.ones(config.num_hidden_layers						,					config.num_attention_heads						,					device=A_	)
										if decoder_head_mask is None:
														_UpperCAmelCase       :		List[str] 		=    torch.ones(config.num_decoder_layers						,					config.num_attention_heads						,					device=A_	)
										if cross_attn_head_mask is None:
														_UpperCAmelCase       :		str 		=    torch.ones(
														    config.num_decoder_layers						,					config.num_attention_heads						,					device=A_	)
										return {
										    "input_ids": input_ids,
										    "decoder_input_ids": decoder_input_ids,
										    "attention_mask": attention_mask,
										    "decoder_attention_mask": decoder_attention_mask,
										    "head_mask": head_mask,
										    "decoder_head_mask": decoder_head_mask,
										    "cross_attn_head_mask": cross_attn_head_mask,
										}


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[int] 		=    ids_tensor([self.batch_size, self.encoder_seq_length]						,					self.vocab_size	)
										_UpperCAmelCase       :		int 		=    ids_tensor([self.batch_size, self.decoder_seq_length]						,					self.vocab_size	)

										# we need to clamp the input ids here to avoid having pad token in between
										# this is because for NllbMoe the position_ids are prepared such that
										# all pad tokens have pos id = 2 and rest are between 2..seq_length
										# and the seq_length here is seq_length - num_pad_tokens
										# but when using past, there is no way of knowing if the past input ids had
										# pad tokens in them, which results in incorrect seq_lenth and which in turn results in
										# position_ids being off by num_pad_tokens in past input
										_UpperCAmelCase       :		Tuple 		=    input_ids.clamp(self.pad_token_id + 1	)
										_UpperCAmelCase       :		Optional[int] 		=    decoder_input_ids.clamp(self.pad_token_id + 1	)

										_UpperCAmelCase       :		Tuple 		=    self.get_config()
										_UpperCAmelCase       :		Tuple 		=    config.num_attention_heads
										_UpperCAmelCase       :		List[Any] 		=    self.prepare_inputs_dict(A_						,					A_						,					A_	)
										return config, input_dict


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		Tuple 		=    self.prepare_config_and_inputs()
										return config, inputs_dict


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return TaConfig(
										    vocab_size=166						,					d_model=self.hidden_size						,					d_ff=self.d_ff						,					d_kv=self.hidden_size // self.num_attention_heads						,					num_layers=self.num_hidden_layers						,					num_decoder_layers=self.decoder_layers						,					num_heads=self.num_attention_heads						,					relative_attention_num_buckets=self.relative_attention_num_buckets						,					dropout_rate=self.dropout_rate						,					initializer_factor=self.initializer_factor						,					eos_token_id=self.eos_token_id						,					bos_token_id=self.pad_token_id						,					pad_token_id=self.pad_token_id						,					decoder_start_token_id=self.decoder_start_token_id						,					)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return TaConfig(
										    vocab_size=self.vocab_size						,					d_model=self.hidden_size						,					d_ff=self.d_ff						,					d_kv=self.hidden_size // self.num_attention_heads						,					num_layers=self.num_hidden_layers						,					num_decoder_layers=self.decoder_layers						,					num_heads=self.num_attention_heads						,					relative_attention_num_buckets=self.relative_attention_num_buckets						,					dropout_rate=self.dropout_rate						,					initializer_factor=self.initializer_factor						,					eos_token_id=self.eos_token_id						,					bos_token_id=self.pad_token_id						,					pad_token_id=self.pad_token_id						,					decoder_start_token_id=self.decoder_start_token_id						,					)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    UMTaModel(config=A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		List[Any] 		=    model(
										    input_ids=A_						,					decoder_input_ids=A_						,					attention_mask=A_						,					decoder_attention_mask=A_						,					)
										_UpperCAmelCase       :		Union[str, Any] 		=    model(input_ids=A_						,					decoder_input_ids=A_	)
										_UpperCAmelCase       :		List[Any] 		=    result.last_hidden_state
										_UpperCAmelCase       :		Union[str, Any] 		=    result.past_key_values
										_UpperCAmelCase       :		Any 		=    result.encoder_last_hidden_state

										self.parent.assertEqual(encoder_output.size()						,					(self.batch_size, self.encoder_seq_length, self.hidden_size)	)
										self.parent.assertEqual(decoder_output.size()						,					(self.batch_size, self.decoder_seq_length, self.hidden_size)	)
										# There should be `num_layers` key value embeddings stored in decoder_past
										self.parent.assertEqual(len(A_	)						,					config.num_layers	)
										# There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past tuple
										self.parent.assertEqual(len(decoder_past[0]	)						,					4	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		Any 		=    UMTaModel(config=A_	).get_decoder().to(A_	).eval()
										# first forward pass
										_UpperCAmelCase       :		Dict 		=    model(A_						,					use_cache=A_	)
										_UpperCAmelCase       :		Optional[int] 		=    model(A_	)
										_UpperCAmelCase       :		Dict 		=    model(A_						,					use_cache=A_	)

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

										_UpperCAmelCase		,							_UpperCAmelCase       :		int 		=    outputs.to_tuple()

										# create hypothetical next token and extent to next_input_ids
										_UpperCAmelCase       :		Optional[Any] 		=    ids_tensor((self.batch_size, 1)						,					config.vocab_size	)

										# append to next input_ids and
										_UpperCAmelCase       :		List[Any] 		=    torch.cat([input_ids, next_tokens]						,					dim=-1	)

										_UpperCAmelCase       :		Dict 		=    model(A_	)["last_hidden_state"]
										_UpperCAmelCase       :		Any 		=    model(A_						,					past_key_values=A_	)["last_hidden_state"]

										# select random slice
										_UpperCAmelCase       :		Optional[Any] 		=    ids_tensor((1,)						,					output_from_past.shape[-1]	).item()
										_UpperCAmelCase       :		Tuple 		=    output_from_no_past[:, -1, random_slice_idx].detach()
										_UpperCAmelCase       :		List[str] 		=    output_from_past[:, 0, random_slice_idx].detach()

										# test that outputs are equal for slice
										self.parent.assertTrue(torch.allclose(A_						,					A_						,					atol=1e-3	)	)





						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[int] 		=    UMTaModel(config=A_	).to(A_	).half().eval()
										_UpperCAmelCase       :		Optional[int] 		=    model(**A_	)["last_hidden_state"]
										self.parent.assertFalse(torch.isnan(A_	).any().item()	)







@require_torch
class 				a  (	UpperCAmelCase							, UpperCAmelCase							, UpperCAmelCase							, unittest.TestCase				):
						_lowercase    						=			(
						    (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
						)
						_lowercase    						=			(UMTaForConditionalGeneration,) if is_torch_available() else ()
						_lowercase    						=			(
						    {
						        "conversational": UMTaForConditionalGeneration,
						        "feature-extraction": UMTaModel,
						        "summarization": UMTaForConditionalGeneration,
						        "text2text-generation": UMTaForConditionalGeneration,
						        "translation": UMTaForConditionalGeneration,
						        "question-answering": UMTaForQuestionAnswering,
						    }
						    if is_torch_available()
						    else {}
						)
						_lowercase    						=			True
						_lowercase    						=			False
						_lowercase    						=			False
						_lowercase    						=			True
						_lowercase    						=			True
						# The small UMT5 model needs higher percentages for CPU/MP tests
						_lowercase    						=			[0.8, 0.9]


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    UMTaModelTester(self	)


						@unittest.skip("Test has a segmentation fault on torch 1.8.0"	)
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    self.model_tester.prepare_config_and_inputs()
										_UpperCAmelCase       :		List[Any] 		=    UMTaModel(config_and_inputs[0]	).to(A_	)
										with tempfile.TemporaryDirectory() as tmpdirname:
														torch.onnx.export(
														    A_						,					(config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])						,					f'{tmpdirname}/t5_test.onnx'						,					export_params=A_						,					opset_version=9						,					input_names=["input_ids", "decoder_input_ids"]						,					)


						@unittest.skipIf(torch_device == "cpu"						,					"Cant do half precision"	)
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		str 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model_fpaa_forward(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		str 		=    ["encoder_attentions", "decoder_attentions", "cross_attentions"]
										_UpperCAmelCase       :		Dict 		=    self.model_tester.prepare_config_and_inputs()
										_UpperCAmelCase       :		List[str] 		=    config_and_inputs[0]
										_UpperCAmelCase       :		int 		=    UMTaForConditionalGeneration(A_	).eval()
										model.to(A_	)

										_UpperCAmelCase       :		str 		=    {
										    "head_mask": torch.zeros(config.num_layers						,					config.num_heads						,					device=A_	),
										    "decoder_head_mask": torch.zeros(config.num_decoder_layers						,					config.num_heads						,					device=A_	),
										    "cross_attn_head_mask": torch.zeros(config.num_decoder_layers						,					config.num_heads						,					device=A_	),
										}

										for attn_name, (name, mask) in zip(A_						,					head_masking.items()	):
														_UpperCAmelCase       :		Dict 		=    {name: mask}
														# Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
														if name == "head_mask":
																		_UpperCAmelCase       :		Dict 		=    torch.ones(
																		    config.num_decoder_layers						,					config.num_heads						,					device=A_	)

														_UpperCAmelCase       :		Optional[Any] 		=    model.generate(
														    config_and_inputs[1]["input_ids"]						,					num_beams=1						,					max_length=3						,					output_attentions=A_						,					return_dict_in_generate=A_						,					**A_						,					)
														# We check the state of decoder_attentions and cross_attentions just from the last step
														_UpperCAmelCase       :		Dict 		=    out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
														self.assertEqual(sum([w.sum().item() for w in attn_weights]	)						,					0.0	)





						@unittest.skip("Does not work on the tiny model as we keep hitting edge cases."	)
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										pass







@require_torch
@require_sentencepiece
@require_tokenizers
class 				a  (	unittest.TestCase				):


						@slow
						@unittest.skip(
						    "Unless we stop stripping left and right by default for all special tokens, the expected ids obtained here will not match the original ones. Wait for https://github.com/huggingface/transformers/pull/23909 to be merged"	)
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    UMTaForConditionalGeneration.from_pretrained("google/umt5-small"						,					return_dict=A_	).to(A_	)
										_UpperCAmelCase       :		Optional[Any] 		=    AutoTokenizer.from_pretrained("google/umt5-small"						,					use_fast=A_						,					legacy=A_	)
										_UpperCAmelCase       :		Any 		=    [
										    "Bonjour monsieur <extra_id_0> bien <extra_id_1>.",
										    "No se como puedo <extra_id_0>.",
										    "This is the reason why we <extra_id_0> them.",
										    "The <extra_id_0> walks in <extra_id_1>, seats",
										    "A <extra_id_0> walks into a bar and orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.",
										]
										_UpperCAmelCase       :		List[str] 		=    tokenizer(A_						,					return_tensors="pt"						,					padding=A_	).input_ids
										# fmt: off
										_UpperCAmelCase       :		Union[str, Any] 		=    torch.tensor(
										    [
										        [ 38530, 210703, 256299, 1410, 256298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
										        [   826, 321, 671, 25922, 256299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
										        [  1460, 339, 312, 19014, 10620, 758, 256299, 2355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
										        [   517, 256299, 14869, 281, 301, 256298, 275, 119983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
										        [   320, 256299, 14869, 281, 2234, 289, 2275, 333,61391, 289, 256298, 543, 256297, 168714, 329, 256296,274, 1],
										    ]	)
										# fmt: on
										torch.testing.assert_allclose(A_						,					A_	)

										_UpperCAmelCase       :		Tuple 		=    model.generate(input_ids.to(A_	)	)
										_UpperCAmelCase       :		int 		=    [
										    "<pad><extra_id_0> et<extra_id_1> [eod] <extra_id_2><extra_id_55>.. [eod] 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 <extra_id_56>ajšietosto<extra_id_56>lleux<extra_id_19><extra_id_6>ajšie</s>",
										    "<pad><extra_id_0>.<extra_id_1>.,<0x0A>...spech <0x0A><extra_id_20> <extra_id_21></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
										    "<pad><extra_id_0> are not going to be a part of the world. We are not going to be a part of<extra_id_1> and<extra_id_2><0x0A><extra_id_48>.<extra_id_48></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
										    "<pad><extra_id_0> door<extra_id_1>, the door<extra_id_2> 피해[/</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
										    "<pad><extra_id_0>nyone who<extra_id_1> drink<extra_id_2> a<extra_id_3> alcohol<extra_id_4> A<extra_id_5> A. This<extra_id_6> I<extra_id_7><extra_id_52><extra_id_53></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
										]
										_UpperCAmelCase       :		Any 		=    tokenizer.batch_decode(A_	)
										self.assertEqual(A_						,					A_	)





import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_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, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch
					from torch import nn

					from transformers import (
					    MODEL_MAPPING,
					    BeitForImageClassification,
					    BeitForMaskedImageModeling,
					    BeitForSemanticSegmentation,
					    BeitModel,
					)
					from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					import PIL
					from PIL import Image

					from transformers import BeitImageProcessor

class 				a  :


						def __init__(       self						,					A_						,					A_=100						,					A_=13						,					A_=30						,					A_=2						,					A_=3						,					A_=True						,					A_=True						,					A_=32						,					A_=4						,					A_=4						,					A_=37						,					A_="gelu"						,					A_=0.1						,					A_=0.1						,					A_=10						,					A_=0.02						,					A_=3						,					A_=None						,					A_=[0, 1, 2, 3]						,					):







										'''simple docstring'''


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

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


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

										_UpperCAmelCase       :		Optional[Any] 		=    None
										_UpperCAmelCase       :		str 		=    None
										if self.use_labels:
														_UpperCAmelCase       :		Union[str, Any] 		=    ids_tensor([self.batch_size]						,					self.type_sequence_label_size	)
														_UpperCAmelCase       :		Optional[int] 		=    ids_tensor([self.batch_size, self.image_size, self.image_size]						,					self.num_labels	)

										_UpperCAmelCase       :		Tuple 		=    self.get_config()

										return config, pixel_values, labels, pixel_labels


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return BeitConfig(
										    vocab_size=self.vocab_size						,					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						,					out_indices=self.out_indices						,					)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Any 		=    BeitModel(config=A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		int 		=    model(A_	)
										self.parent.assertEqual(result.last_hidden_state.shape						,					(self.batch_size, self.seq_length, self.hidden_size)	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Dict 		=    BeitForMaskedImageModeling(config=A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		Any 		=    model(A_	)
										self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.seq_length - 1, self.vocab_size)	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    self.type_sequence_label_size
										_UpperCAmelCase       :		Optional[int] 		=    BeitForImageClassification(A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		List[str] 		=    model(A_						,					labels=A_	)
										self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.type_sequence_label_size)	)

										# test greyscale images
										_UpperCAmelCase       :		Optional[Any] 		=    1
										_UpperCAmelCase       :		Tuple 		=    BeitForImageClassification(A_	)
										model.to(A_	)
										model.eval()

										_UpperCAmelCase       :		int 		=    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
										_UpperCAmelCase       :		Dict 		=    model(A_						,					labels=A_	)
										self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.type_sequence_label_size)	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Dict 		=    self.num_labels
										_UpperCAmelCase       :		Dict 		=    BeitForSemanticSegmentation(A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		str 		=    model(A_	)
										self.parent.assertEqual(
										    result.logits.shape						,					(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)	)
										_UpperCAmelCase       :		List[Any] 		=    model(A_						,					labels=A_	)
										self.parent.assertEqual(
										    result.logits.shape						,					(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)	)





						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    self.prepare_config_and_inputs()
										_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase       :		Union[str, Any] 		=    config_and_inputs
										_UpperCAmelCase       :		int 		=    {"pixel_values": pixel_values}
										return config, inputs_dict







@require_torch
class 				a  (	UpperCAmelCase							, UpperCAmelCase							, unittest.TestCase				):
						_lowercase    						=			(
						    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
						    if is_torch_available()
						    else ()
						)
						_lowercase    						=			(
						    {
						        "feature-extraction": BeitModel,
						        "image-classification": BeitForImageClassification,
						        "image-segmentation": BeitForSemanticSegmentation,
						    }
						    if is_torch_available()
						    else {}
						)

						_lowercase    						=			False
						_lowercase    						=			False
						_lowercase    						=			False


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    BeitModelTester(self	)
										_UpperCAmelCase       :		List[Any] 		=    ConfigTester(self						,					config_class=A_						,					has_text_modality=A_						,					hidden_size=37	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										self.config_tester.run_common_tests()


						@unittest.skip(reason="BEiT does not use inputs_embeds"	)
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										pass


						@require_torch_multi_gpu
						@unittest.skip(reason="BEiT has some layers using `add_module` which doesn't work well with `nn.DataParallel`"	)
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										pass


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		Dict 		=    self.model_tester.prepare_config_and_inputs_for_common()

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


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		int 		=    self.model_tester.prepare_config_and_inputs_for_common()

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

														_UpperCAmelCase       :		int 		=    ["pixel_values"]
														self.assertListEqual(arg_names[:1]						,					A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[Any] 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_semantic_segmentation(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										if not self.model_tester.is_training:
														return

										_UpperCAmelCase		,							_UpperCAmelCase       :		Union[str, Any] 		=    self.model_tester.prepare_config_and_inputs_for_common()
										_UpperCAmelCase       :		str 		=    True

										for model_class in self.all_model_classes:
														# we don't test BeitForMaskedImageModeling
														if model_class in [*get_values(A_	), BeitForMaskedImageModeling]:
																		continue

														_UpperCAmelCase       :		Union[str, Any] 		=    model_class(A_	)
														model.to(A_	)
														model.train()
														_UpperCAmelCase       :		str 		=    self._prepare_for_class(A_						,					A_						,					return_labels=A_	)
														_UpperCAmelCase       :		Optional[Any] 		=    model(**A_	).loss
														loss.backward()


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		Dict 		=    self.model_tester.prepare_config_and_inputs_for_common()
										if not self.model_tester.is_training:
														return

										_UpperCAmelCase       :		Tuple 		=    False
										_UpperCAmelCase       :		Any 		=    True

										for model_class in self.all_model_classes:
														# we don't test BeitForMaskedImageModeling
														if (
														    model_class in [*get_values(A_	), BeitForMaskedImageModeling]
														    or not model_class.supports_gradient_checkpointing
														):
																		continue

														_UpperCAmelCase       :		Optional[Any] 		=    model_class(A_	)
														model.gradient_checkpointing_enable()
														model.to(A_	)
														model.train()
														_UpperCAmelCase       :		Union[str, Any] 		=    self._prepare_for_class(A_						,					A_						,					return_labels=A_	)
														_UpperCAmelCase       :		Optional[Any] 		=    model(**A_	).loss
														loss.backward()


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		List[Any] 		=    self.model_tester.prepare_config_and_inputs_for_common()

										_UpperCAmelCase       :		Optional[Any] 		=    _config_zero_init(A_	)
										for model_class in self.all_model_classes:
														_UpperCAmelCase       :		Tuple 		=    model_class(config=A_	)
														for name, param in model.named_parameters():
																		# we skip lambda parameters as these require special initial values
																		# determined by config.layer_scale_init_value
																		if "lambda" in name:
																						continue
																		if param.requires_grad:
																						self.assertIn(
																						    ((param.data.mean() * 1e9).round() / 1e9).item()						,					[0.0, 1.0]						,					msg=f'Parameter {name} of model {model_class} seems not properly initialized'						,					)





						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_UpperCAmelCase       :		Dict 		=    BeitModel.from_pretrained(A_	)
														self.assertIsNotNone(A_	)







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





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


						@cached_property
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"	) if is_vision_available() else None


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    BeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"	).to(A_	)

										_UpperCAmelCase       :		Dict 		=    self.default_image_processor
										_UpperCAmelCase       :		Optional[Any] 		=    prepare_img()
										_UpperCAmelCase       :		Any 		=    image_processor(images=A_						,					return_tensors="pt"	).pixel_values.to(A_	)

										# prepare bool_masked_pos
										_UpperCAmelCase       :		Dict 		=    torch.ones((1, 196)						,					dtype=torch.bool	).to(A_	)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase       :		List[str] 		=    model(pixel_values=A_						,					bool_masked_pos=A_	)
										_UpperCAmelCase       :		Optional[int] 		=    outputs.logits

										# verify the logits
										_UpperCAmelCase       :		Optional[int] 		=    torch.Size((1, 196, 8192)	)
										self.assertEqual(logits.shape						,					A_	)

										_UpperCAmelCase       :		Tuple 		=    torch.tensor(
										    [[-3.24_37, 0.50_72, -13.91_74], [-3.24_56, 0.49_48, -13.94_01], [-3.20_33, 0.51_21, -13.85_50]]	).to(A_	)

										self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]						,					A_						,					atol=1e-2	)	)


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    BeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"	).to(A_	)

										_UpperCAmelCase       :		str 		=    self.default_image_processor
										_UpperCAmelCase       :		Union[str, Any] 		=    prepare_img()
										_UpperCAmelCase       :		List[Any] 		=    image_processor(images=A_						,					return_tensors="pt"	).to(A_	)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase       :		Tuple 		=    model(**A_	)
										_UpperCAmelCase       :		Dict 		=    outputs.logits

										# verify the logits
										_UpperCAmelCase       :		List[str] 		=    torch.Size((1, 1000)	)
										self.assertEqual(logits.shape						,					A_	)

										_UpperCAmelCase       :		int 		=    torch.tensor([-1.23_85, -1.09_87, -1.01_08]	).to(A_	)

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

										_UpperCAmelCase       :		List[Any] 		=    281
										self.assertEqual(logits.argmax(-1	).item()						,					A_	)


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[int] 		=    BeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"	).to(
										    A_	)

										_UpperCAmelCase       :		str 		=    self.default_image_processor
										_UpperCAmelCase       :		List[str] 		=    prepare_img()
										_UpperCAmelCase       :		List[Any] 		=    image_processor(images=A_						,					return_tensors="pt"	).to(A_	)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase       :		Any 		=    model(**A_	)
										_UpperCAmelCase       :		List[Any] 		=    outputs.logits

										# verify the logits
										_UpperCAmelCase       :		List[str] 		=    torch.Size((1, 21841)	)
										self.assertEqual(logits.shape						,					A_	)

										_UpperCAmelCase       :		Union[str, Any] 		=    torch.tensor([1.68_81, -0.27_87, 0.59_01]	).to(A_	)

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

										_UpperCAmelCase       :		List[str] 		=    2396
										self.assertEqual(logits.argmax(-1	).item()						,					A_	)


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Any 		=    BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"	)
										_UpperCAmelCase       :		str 		=    model.to(A_	)

										_UpperCAmelCase       :		int 		=    BeitImageProcessor(do_resize=A_						,					size=640						,					do_center_crop=A_	)

										_UpperCAmelCase       :		Union[str, Any] 		=    load_dataset("hf-internal-testing/fixtures_ade20k"						,					split="test"	)
										_UpperCAmelCase       :		Any 		=    Image.open(ds[0]["file"]	)
										_UpperCAmelCase       :		int 		=    image_processor(images=A_						,					return_tensors="pt"	).to(A_	)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase       :		Dict 		=    model(**A_	)
										_UpperCAmelCase       :		Optional[int] 		=    outputs.logits

										# verify the logits
										_UpperCAmelCase       :		str 		=    torch.Size((1, 150, 160, 160)	)
										self.assertEqual(logits.shape						,					A_	)

										_UpperCAmelCase       :		str 		=    version.parse(PIL.__version__	) < version.parse("9.0.0"	)

										if is_pillow_less_than_a:
														_UpperCAmelCase       :		int 		=    torch.tensor(
														    [
														        [[-4.92_25, -2.39_54, -3.05_22], [-2.88_22, -1.00_46, -1.75_61], [-2.95_49, -1.32_28, -2.13_47]],
														        [[-5.81_68, -3.41_29, -4.07_78], [-3.86_51, -2.22_14, -3.02_77], [-3.83_56, -2.46_43, -3.35_35]],
														        [[-0.00_78, 3.99_52, 4.07_54], [2.98_56, 4.69_44, 5.00_35], [3.24_13, 4.78_13, 4.99_69]],
														    ]						,					device=A_						,					)
										else:
														_UpperCAmelCase       :		Any 		=    torch.tensor(
														    [
														        [[-4.89_60, -2.36_88, -3.03_55], [-2.84_78, -0.98_36, -1.74_18], [-2.94_49, -1.33_32, -2.14_56]],
														        [[-5.80_81, -3.41_24, -4.10_06], [-3.85_61, -2.20_81, -3.03_23], [-3.83_65, -2.46_01, -3.36_69]],
														        [[-0.03_09, 3.98_68, 4.05_40], [2.96_40, 4.68_77, 4.99_76], [3.20_81, 4.76_90, 4.99_42]],
														    ]						,					device=A_						,					)

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





						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[int] 		=    BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"	)
										_UpperCAmelCase       :		Dict 		=    model.to(A_	)

										_UpperCAmelCase       :		Dict 		=    BeitImageProcessor(do_resize=A_						,					size=640						,					do_center_crop=A_	)

										_UpperCAmelCase       :		Optional[Any] 		=    load_dataset("hf-internal-testing/fixtures_ade20k"						,					split="test"	)
										_UpperCAmelCase       :		List[str] 		=    Image.open(ds[0]["file"]	)
										_UpperCAmelCase       :		Union[str, Any] 		=    image_processor(images=A_						,					return_tensors="pt"	).to(A_	)

										# forward pass
										with torch.no_grad():
														_UpperCAmelCase       :		List[Any] 		=    model(**A_	)

										_UpperCAmelCase       :		Tuple 		=    outputs.logits.detach().cpu()

										_UpperCAmelCase       :		str 		=    image_processor.post_process_semantic_segmentation(outputs=A_						,					target_sizes=[(500, 300)]	)
										_UpperCAmelCase       :		Dict 		=    torch.Size((500, 300)	)
										self.assertEqual(segmentation[0].shape						,					A_	)

										_UpperCAmelCase       :		Union[str, Any] 		=    image_processor.post_process_semantic_segmentation(outputs=A_	)
										_UpperCAmelCase       :		Union[str, Any] 		=    torch.Size((160, 160)	)
										self.assertEqual(segmentation[0].shape						,					A_	)


