
"""simple docstring"""





import inspect
import unittest
import warnings

from transformers import DeiTConfig
from transformers.models.auto import get_values
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 (
		    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
		    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
		    MODEL_MAPPING,
		    DeiTForImageClassification,
		    DeiTForImageClassificationWithTeacher,
		    DeiTForMaskedImageModeling,
		    DeiTModel,
		)
		from transformers.models.deit.modeling_deit import DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import DeiTImageProcessor



class __lowerCAmelCase					:




			"""simple docstring"""







			def __init__(					self    :       Optional[int]							, _snake_case    :       str							, _snake_case    :       Union[str, Any]=13							, _snake_case    :       Union[str, Any]=30							, _snake_case    :       str=2							, _snake_case    :       Optional[int]=3							, _snake_case    :       List[str]=True							, _snake_case    :       str=True							, _snake_case    :       List[str]=32							, _snake_case    :       str=5							, _snake_case    :       str=4							, _snake_case    :       Any=37							, _snake_case    :       Dict="gelu"							, _snake_case    :       str=0.1							, _snake_case    :       Optional[int]=0.1							, _snake_case    :       Any=10							, _snake_case    :       List[str]=0.0_2							, _snake_case    :       str=3							, _snake_case    :       str=None							, _snake_case    :       Dict=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 DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
								A_								=    (image_size // patch_size) ** 2
								A_								=    num_patches + 2




			def   lowerCamelCase__  (					self    :       Dict						)   ->			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   lowerCamelCase__  (					self    :       List[Any]						)   ->			Union[str, Any]:







								"""simple docstring"""
								return DeiTConfig(
								    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=_snake_case							, initializer_range=self.initializer_range							, encoder_stride=self.encoder_stride							, )




			def   lowerCamelCase__  (					self    :       Optional[int]							, _snake_case    :       str							, _snake_case    :       Optional[int]							, _snake_case    :       Tuple						)   ->			Dict:







								"""simple docstring"""
								A_								=    DeiTModel(config=_snake_case						)
								model.to(_snake_case						)
								model.eval()
								A_								=    model(_snake_case						)
								self.parent.assertEqual(result.last_hidden_state.shape							, (self.batch_size, self.seq_length, self.hidden_size)						)




			def   lowerCamelCase__  (					self    :       Optional[Any]							, _snake_case    :       Optional[int]							, _snake_case    :       List[str]							, _snake_case    :       Tuple						)   ->			Dict:







								"""simple docstring"""
								A_								=    DeiTForMaskedImageModeling(config=_snake_case						)
								model.to(_snake_case						)
								model.eval()
								A_								=    model(_snake_case						)
								self.parent.assertEqual(
								    result.reconstruction.shape							, (self.batch_size, self.num_channels, self.image_size, self.image_size)						)

								# test greyscale images
								A_								=    1
								A_								=    DeiTForMaskedImageModeling(_snake_case						)
								model.to(_snake_case						)
								model.eval()

								A_								=    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
								A_								=    model(_snake_case						)
								self.parent.assertEqual(result.reconstruction.shape							, (self.batch_size, 1, self.image_size, self.image_size)						)




			def   lowerCamelCase__  (					self    :       Any							, _snake_case    :       int							, _snake_case    :       str							, _snake_case    :       Optional[Any]						)   ->			List[str]:







								"""simple docstring"""
								A_								=    self.type_sequence_label_size
								A_								=    DeiTForImageClassification(_snake_case						)
								model.to(_snake_case						)
								model.eval()
								A_								=    model(_snake_case							, labels=_snake_case						)
								self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.type_sequence_label_size)						)

								# test greyscale images
								A_								=    1
								A_								=    DeiTForImageClassification(_snake_case						)
								model.to(_snake_case						)
								model.eval()

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






			def   lowerCamelCase__  (					self    :       Dict						)   ->			Any:







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





@require_torch
class __lowerCAmelCase					(       _lowercase  ,					_lowercase  ,					unittest.TestCase   ):




			"""simple docstring"""







			snake_case						   =					(
			    (
			        DeiTModel,
			        DeiTForImageClassification,
			        DeiTForImageClassificationWithTeacher,
			        DeiTForMaskedImageModeling,
			    )
			    if is_torch_available()
			    else ()
			)
			snake_case						   =					(
			    {
			        "feature-extraction": DeiTModel,
			        "image-classification": (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
			    }
			    if is_torch_available()
			    else {}
			)

			snake_case						   =					False
			snake_case						   =					False
			snake_case						   =					False




			def   lowerCamelCase__  (					self    :       List[str]						)   ->			Union[str, Any]:







								"""simple docstring"""
								A_								=    DeiTModelTester(self						)
								A_								=    ConfigTester(self							, config_class=_snake_case							, has_text_modality=_snake_case							, hidden_size=37						)




			def   lowerCamelCase__  (					self    :       Dict						)   ->			int:







								"""simple docstring"""
								self.config_tester.run_common_tests()




			@unittest.skip(reason="DeiT does not use inputs_embeds"						)
			def   lowerCamelCase__  (					self    :       Union[str, Any]						)   ->			Tuple:







								"""simple docstring"""
								pass




			def   lowerCamelCase__  (					self    :       Optional[Any]						)   ->			str:







								"""simple docstring"""
								A_			,     A_								=    self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
													A_								=    model_class(_snake_case						)
													self.assertIsInstance(model.get_input_embeddings()							, (nn.Module)						)
													A_								=    model.get_output_embeddings()
													self.assertTrue(x is None or isinstance(_snake_case							, nn.Linear						)						)




			def   lowerCamelCase__  (					self    :       Union[str, Any]						)   ->			Union[str, Any]:







								"""simple docstring"""
								A_			,     A_								=    self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
													A_								=    model_class(_snake_case						)
													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]							, _snake_case						)




			def   lowerCamelCase__  (					self    :       int						)   ->			str:







								"""simple docstring"""
								A_								=    self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case						)




			def   lowerCamelCase__  (					self    :       Dict						)   ->			str:







								"""simple docstring"""
								A_								=    self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_image_modeling(*_snake_case						)




			def   lowerCamelCase__  (					self    :       List[Any]						)   ->			Optional[Any]:







								"""simple docstring"""
								A_								=    self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*_snake_case						)




			def   lowerCamelCase__  (					self    :       List[str]							, _snake_case    :       List[str]							, _snake_case    :       str							, _snake_case    :       str=False						)   ->			int:







								"""simple docstring"""
								A_								=    super()._prepare_for_class(_snake_case							, _snake_case							, return_labels=_snake_case						)

								if return_labels:
													if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
																		del inputs_dict["labels"]

								return inputs_dict




			def   lowerCamelCase__  (					self    :       List[str]						)   ->			Optional[Any]:







								"""simple docstring"""
								if not self.model_tester.is_training:
													return

								A_			,     A_								=    self.model_tester.prepare_config_and_inputs_for_common()
								A_								=    True

								for model_class in self.all_model_classes:
													# DeiTForImageClassificationWithTeacher supports inference-only
													if (
													    model_class in get_values(_snake_case						)
													    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
													):
																		continue
													A_								=    model_class(_snake_case						)
													model.to(_snake_case						)
													model.train()
													A_								=    self._prepare_for_class(_snake_case							, _snake_case							, return_labels=_snake_case						)
													A_								=    model(**_snake_case						).loss
													loss.backward()




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			Tuple:







								"""simple docstring"""
								A_			,     A_								=    self.model_tester.prepare_config_and_inputs_for_common()
								if not self.model_tester.is_training:
													return

								A_								=    False
								A_								=    True

								for model_class in self.all_model_classes:
													if model_class in get_values(_snake_case						) or not model_class.supports_gradient_checkpointing:
																		continue
													# DeiTForImageClassificationWithTeacher supports inference-only
													if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
																		continue
													A_								=    model_class(_snake_case						)
													model.gradient_checkpointing_enable()
													model.to(_snake_case						)
													model.train()
													A_								=    self._prepare_for_class(_snake_case							, _snake_case							, return_labels=_snake_case						)
													A_								=    model(**_snake_case						).loss
													loss.backward()




			def   lowerCamelCase__  (					self    :       Any						)   ->			Optional[Any]:







								"""simple docstring"""
								A_			,     A_								=    self.model_tester.prepare_config_and_inputs_for_common()

								A_								=    [
								    {"title": "multi_label_classification", "num_labels": 2, "dtype": torch.float},
								    {"title": "single_label_classification", "num_labels": 1, "dtype": torch.long},
								    {"title": "regression", "num_labels": 1, "dtype": torch.float},
								]

								for model_class in self.all_model_classes:
													if (
													    model_class
													    not in [
													        *get_values(_snake_case						),
													        *get_values(_snake_case						),
													    ]
													    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
													):
																		continue

													for problem_type in problem_types:
																		with self.subTest(msg=F'Testing {model_class} with {problem_type["title"]}'						):
																							A_								=    problem_type["title"]
																							A_								=    problem_type["num_labels"]

																							A_								=    model_class(_snake_case						)
																							model.to(_snake_case						)
																							model.train()

																							A_								=    self._prepare_for_class(_snake_case							, _snake_case							, return_labels=_snake_case						)

																							if problem_type["num_labels"] > 1:
																												A_								=    inputs["labels"].unsqueeze(1						).repeat(1							, problem_type["num_labels"]						)

																							A_								=    inputs["labels"].to(problem_type["dtype"]						)

																							# This tests that we do not trigger the warning form PyTorch "Using a target size that is different
																							# to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
																							# they have the same size." which is a symptom something in wrong for the regression problem.
																							# See https://github.com/huggingface/transformers/issues/11780
																							with warnings.catch_warnings(record=_snake_case						) as warning_list:
																												A_								=    model(**_snake_case						).loss
																							for w in warning_list:
																												if "Using a target size that is different to the input size" in str(w.message						):
																																	raise ValueError(
																																	    F'Something is going wrong in the regression problem: intercepted {w.message}'						)

																							loss.backward()






			@slow
			def   lowerCamelCase__  (					self    :       str						)   ->			int:







								"""simple docstring"""
								for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													A_								=    DeiTModel.from_pretrained(_snake_case						)
													self.assertIsNotNone(_snake_case						)





def 				A_  ():

					'''simple docstring'''





					A_								=    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"					)
					return image




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




			"""simple docstring"""







			@cached_property
			def   lowerCamelCase__  (					self    :       str						)   ->			Optional[Any]:







								"""simple docstring"""
								return (
								    DeiTImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224"						)
								    if is_vision_available()
								    else None
								)




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







								"""simple docstring"""
								A_								=    DeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224"						).to(
								    _snake_case						)

								A_								=    self.default_image_processor
								A_								=    prepare_img()
								A_								=    image_processor(images=_snake_case							, return_tensors="pt"						).to(_snake_case						)

								# forward pass
								with torch.no_grad():
													A_								=    model(**_snake_case						)

								# verify the logits
								A_								=    torch.Size((1, 1_000)						)
								self.assertEqual(outputs.logits.shape							, _snake_case						)

								A_								=    torch.tensor([-1.0_2_6_6, 0.1_9_1_2, -1.2_8_6_1]						).to(_snake_case						)

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






			@slow
			@require_accelerate
			@require_torch_gpu
			def   lowerCamelCase__  (					self    :       Dict						)   ->			Union[str, Any]:







								"""simple docstring"""
								A_								=    DeiTModel.from_pretrained(
								    "facebook/deit-base-distilled-patch16-224"							, torch_dtype=torch.floataa							, device_map="auto"						)
								A_								=    self.default_image_processor

								A_								=    prepare_img()
								A_								=    image_processor(images=_snake_case							, return_tensors="pt"						)
								A_								=    inputs.pixel_values.to(_snake_case						)

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






"""simple docstring"""





import argparse
import json
import os

import torch

from transformers import LukeConfig, LukeModel, LukeTokenizer, RobertaTokenizer
from transformers.tokenization_utils_base import AddedToken
@torch.no_grad()
def 				A_  (__a					, __a					, __a					, __a					, __a					):

					'''simple docstring'''





					with open(__a					) as metadata_file:
										A_								=    json.load(__a					)
					A_								=    LukeConfig(use_entity_aware_attention=__a					, **metadata["model_config"]					)

					# Load in the weights from the checkpoint_path
					A_								=    torch.load(__a					, map_location="cpu"					)

					# Load the entity vocab file
					A_								=    load_entity_vocab(__a					)

					A_								=    RobertaTokenizer.from_pretrained(metadata["model_config"]["bert_model_name"]					)

					# Add special tokens to the token vocabulary for downstream tasks
					A_								=    AddedToken("<ent>"					, lstrip=__a					, rstrip=__a					)
					A_								=    AddedToken("<ent2>"					, lstrip=__a					, rstrip=__a					)
					tokenizer.add_special_tokens({"additional_special_tokens": [entity_token_a, entity_token_a]}					)
					config.vocab_size += 2

					print(f'Saving tokenizer to {pytorch_dump_folder_path}'					)
					tokenizer.save_pretrained(__a					)
					with open(os.path.join(__a					, LukeTokenizer.vocab_files_names["entity_vocab_file"]					)					, "w"					) as f:
										json.dump(__a					, __a					)

					A_								=    LukeTokenizer.from_pretrained(__a					)

					# Initialize the embeddings of the special tokens
					A_								=    state_dict["embeddings.word_embeddings.weight"]
					A_								=    word_emb[tokenizer.convert_tokens_to_ids(["@"]					)[0]].unsqueeze(0					)
					A_								=    word_emb[tokenizer.convert_tokens_to_ids(["#"]					)[0]].unsqueeze(0					)
					A_								=    torch.cat([word_emb, ent_emb, enta_emb]					)

					# Initialize the query layers of the entity-aware self-attention mechanism
					for layer_index in range(config.num_hidden_layers					):
										for matrix_name in ["query.weight", "query.bias"]:
															A_								=    f'encoder.layer.{layer_index}.attention.self.'
															A_								=    state_dict[prefix + matrix_name]
															A_								=    state_dict[prefix + matrix_name]
															A_								=    state_dict[prefix + matrix_name]

    # Initialize the embedding of the [MASK2] entity using that of the [MASK] entity for downstream tasks
					A_								=    state_dict["entity_embeddings.entity_embeddings.weight"]
					A_								=    entity_emb[entity_vocab["[MASK]"]]

					A_								=    LukeModel(config=__a					).eval()

					A_			,     A_								=    model.load_state_dict(__a					, strict=__a					)
					if not (len(__a					) == 1 and missing_keys[0] == "embeddings.position_ids"):
										raise ValueError(f'Missing keys {", ".join(__a					)}. Expected only missing embeddings.position_ids'					)
					if not (all(key.startswith("entity_predictions"					) or key.startswith("lm_head"					) for key in unexpected_keys					)):
										raise ValueError(
										    "Unexpected keys"
										    f' {", ".join([key for key in unexpected_keys if not (key.startswith("entity_predictions"					) or key.startswith("lm_head"					))]					)}'					)

					# Check outputs
					A_								=    LukeTokenizer.from_pretrained(__a					, task="entity_classification"					)

					A_								=    (
					    "Top seed Ana Ivanovic said on Thursday she could hardly believe her luck as a fortuitous netcord helped the"
					    " new world number one avoid a humiliating second- round exit at Wimbledon ."
					)
					A_								=    (39, 42)
					A_								=    tokenizer(__a					, entity_spans=[span]					, add_prefix_space=__a					, return_tensors="pt"					)

					A_								=    model(**__a					)

					# Verify word hidden states
					if model_size == "large":
										A_								=    torch.Size((1, 42, 1024)					)
										A_								=    torch.tensor(
										    [[0.0133, 0.0865, 0.0095], [0.3093, -0.2576, -0.7418], [-0.1720, -0.2117, -0.2869]]					)
					else:  # base
										A_								=    torch.Size((1, 42, 768)					)
										A_								=    torch.tensor([[0.0037, 0.1368, -0.0091], [0.1099, 0.3329, -0.1095], [0.0765, 0.5335, 0.1179]]					)

					if not (outputs.last_hidden_state.shape == expected_shape):
										raise ValueError(
										    f'Outputs.last_hidden_state.shape is {outputs.last_hidden_state.shape}, Expected shape is {expected_shape}'					)
					if not torch.allclose(outputs.last_hidden_state[0, :3, :3]					, __a					, atol=1e-4					):
										raise ValueError

					# Verify entity hidden states
					if model_size == "large":
										A_								=    torch.Size((1, 1, 1024)					)
										A_								=    torch.tensor([[0.0466, -0.0106, -0.0179]]					)
					else:  # base
										A_								=    torch.Size((1, 1, 768)					)
										A_								=    torch.tensor([[0.1457, 0.1044, 0.0174]]					)

					if not (outputs.entity_last_hidden_state.shape != expected_shape):
										raise ValueError(
										    f'Outputs.entity_last_hidden_state.shape is {outputs.entity_last_hidden_state.shape}, Expected shape is'
										    f' {expected_shape}'					)
					if not torch.allclose(outputs.entity_last_hidden_state[0, :3, :3]					, __a					, atol=1e-4					):
										raise ValueError

					# Finally, save our PyTorch model and tokenizer
					print("Saving PyTorch model to {}".format(__a					)					)
					model.save_pretrained(__a					)
def 				A_  (__a					):

					'''simple docstring'''





					A_								=    {}
					with open(__a					, "r"					, encoding="utf-8"					) as f:
										for index, line in enumerate(__a					):
															A_			,     A_								=    line.rstrip().split("\t"					)
															A_								=    index

					return entity_vocab


if __name__ == "__main__":
		UpperCamelCase_       :    Dict						      =			argparse.ArgumentParser()
		# Required parameters
		parser.add_argument('''--checkpoint_path''', type=str, help='''Path to a pytorch_model.bin file.''')
		parser.add_argument(
		    '''--metadata_path''', default=None, type=str, help='''Path to a metadata.json file, defining the configuration.'''
		)
		parser.add_argument(
		    '''--entity_vocab_path''',
		    default=None,
		    type=str,
		    help='''Path to an entity_vocab.tsv file, containing the entity vocabulary.''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to where to dump the output PyTorch model.'''
		)
		parser.add_argument(
		    '''--model_size''', default='''base''', type=str, choices=['''base''', '''large'''], help='''Size of the model to be converted.'''
		)
		UpperCamelCase_       :    List[str]						      =			parser.parse_args()
		convert_luke_checkpoint(
		    args.checkpoint_path,
		    args.metadata_path,
		    args.entity_vocab_path,
		    args.pytorch_dump_folder_path,
		    args.model_size,
		)




