





'''simple docstring'''






import argparse
import json

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm.models import create_model

from transformers import (
    BeitImageProcessor,
    DataaVecVisionConfig,
    DataaVecVisionForImageClassification,
    DataaVecVisionModel,
)





def      lowercase  (  lowerCAmelCase :						Dict   , lowerCAmelCase :						List[str]=False   , lowerCAmelCase :						List[Any]=False   , lowerCAmelCase :						int="data2vec."):



				"""simple docstring"""


				_A :				str				=  '''backbone.''' if is_semantic else ''''''

				_A :				List[str]				=  []
				for i in range(config.num_hidden_layers):
								# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
								rename_keys.append(
								    (f"""{prefix}blocks.{i}.norm1.weight""", f"""{hf_prefix}encoder.layer.{i}.layernorm_before.weight"""))
								rename_keys.append((f"""{prefix}blocks.{i}.norm1.bias""", f"""{hf_prefix}encoder.layer.{i}.layernorm_before.bias"""))
								rename_keys.append(
								    (f"""{prefix}blocks.{i}.attn.proj.weight""", f"""{hf_prefix}encoder.layer.{i}.attention.output.dense.weight"""))
								rename_keys.append(
								    (f"""{prefix}blocks.{i}.attn.proj.bias""", f"""{hf_prefix}encoder.layer.{i}.attention.output.dense.bias"""))
								rename_keys.append(
								    (f"""{prefix}blocks.{i}.norm2.weight""", f"""{hf_prefix}encoder.layer.{i}.layernorm_after.weight"""))
								rename_keys.append((f"""{prefix}blocks.{i}.norm2.bias""", f"""{hf_prefix}encoder.layer.{i}.layernorm_after.bias"""))
								rename_keys.append(
								    (f"""{prefix}blocks.{i}.mlp.fc1.weight""", f"""{hf_prefix}encoder.layer.{i}.intermediate.dense.weight"""))
								rename_keys.append(
								    (f"""{prefix}blocks.{i}.mlp.fc1.bias""", f"""{hf_prefix}encoder.layer.{i}.intermediate.dense.bias"""))
								rename_keys.append((f"""{prefix}blocks.{i}.mlp.fc2.weight""", f"""{hf_prefix}encoder.layer.{i}.output.dense.weight"""))
								rename_keys.append((f"""{prefix}blocks.{i}.mlp.fc2.bias""", f"""{hf_prefix}encoder.layer.{i}.output.dense.bias"""))

				# projection layer + position embeddings
				rename_keys.extend(
				    [
				        (f"""{prefix}cls_token""", f"""{hf_prefix}embeddings.cls_token"""),
				        (f"""{prefix}patch_embed.proj.weight""", f"""{hf_prefix}embeddings.patch_embeddings.projection.weight"""),
				        (f"""{prefix}patch_embed.proj.bias""", f"""{hf_prefix}embeddings.patch_embeddings.projection.bias"""),
				    ])

				if has_lm_head:
								# mask token + shared relative position bias + layernorm
								rename_keys.extend(
								    [
								        ('''mask_token''', f"""{hf_prefix}embeddings.mask_token"""),
								        (
								            '''rel_pos_bias.relative_position_bias_table''',
								            f"""{hf_prefix}encoder.relative_position_bias.relative_position_bias_table""",
								        ),
								        (
								            '''rel_pos_bias.relative_position_index''',
								            f"""{hf_prefix}encoder.relative_position_bias.relative_position_index""",
								        ),
								        ('''norm.weight''', '''layernorm.weight'''),
								        ('''norm.bias''', '''layernorm.bias'''),
								    ])
				elif is_semantic:
								# semantic segmentation classification heads
								rename_keys.extend(
								    [
								        ('''decode_head.conv_seg.weight''', '''decode_head.classifier.weight'''),
								        ('''decode_head.conv_seg.bias''', '''decode_head.classifier.bias'''),
								        ('''auxiliary_head.conv_seg.weight''', '''auxiliary_head.classifier.weight'''),
								        ('''auxiliary_head.conv_seg.bias''', '''auxiliary_head.classifier.bias'''),
								    ])
				else:
								# layernorm + classification head
								rename_keys.extend(
								    [
								        ('''fc_norm.weight''', f"""{hf_prefix}pooler.layernorm.weight"""),
								        ('''fc_norm.bias''', f"""{hf_prefix}pooler.layernorm.bias"""),
								        ('''head.weight''', '''classifier.weight'''),
								        ('''head.bias''', '''classifier.bias'''),
								    ])

				return rename_keys





def      lowercase  (  lowerCAmelCase :						str   , lowerCAmelCase :						Dict   , lowerCAmelCase :						List[str]=False   , lowerCAmelCase :						Any=False   , lowerCAmelCase :						Optional[int]="data2vec_vision."):



				"""simple docstring"""


				for i in range(config.num_hidden_layers):
								_A :				Union[str, Any]				=  '''backbone.''' if is_semantic else ''''''
								# queries, keys and values
								_A :				List[Any]				=  state_dict.pop(f"""{prefix}blocks.{i}.attn.qkv.weight""")
								_A :				Tuple				=  state_dict.pop(f"""{prefix}blocks.{i}.attn.q_bias""")
								_A :				str				=  state_dict.pop(f"""{prefix}blocks.{i}.attn.v_bias""")

								_A :				List[str]				=  in_proj_weight[
								    : config.hidden_size, :
								]
								_A :				Any				=  q_bias
								_A :				Any				=  in_proj_weight[
								    config.hidden_size : config.hidden_size * 2, :
								]
								_A :				Dict				=  in_proj_weight[
								    -config.hidden_size :, :
								]
								_A :				List[Any]				=  v_bias

								# gamma_1 and gamma_2
								# we call them lambda because otherwise they are renamed when using .from_pretrained
								_A :				Dict				=  state_dict.pop(f"""{prefix}blocks.{i}.gamma_1""")
								_A :				int				=  state_dict.pop(f"""{prefix}blocks.{i}.gamma_2""")

								_A :				Any				=  gamma_a
								_A :				List[Any]				=  gamma_a

								# relative_position bias table + index
								if not has_lm_head:
												# each layer has its own relative position bias
												_A :				Optional[int]				=  state_dict.pop(f"""{prefix}blocks.{i}.attn.relative_position_bias_table""")
												_A :				int				=  state_dict.pop(f"""{prefix}blocks.{i}.attn.relative_position_index""")

												_A :				List[str]				=  table
												_A :				str				=  index





def      lowercase  (  ):



				"""simple docstring"""


				_A :				Union[str, Any]				=  argparse.ArgumentParser(
				    '''Convert Data2VecVision to HF for image classification and pretraining'''   , add_help=lowerCAmelCase)
				parser.add_argument('''--hf_checkpoint_name'''   , type=lowerCAmelCase)
				parser.add_argument('''--input_size'''   , default=224   , type=lowerCAmelCase   , help='''images input size''')
				parser.add_argument('''--beit_checkpoint'''   , default=''''''   , help='''beit checkpoint''')

				return parser.parse_args()





def      lowercase  (  lowerCAmelCase :						Optional[int]   , lowerCAmelCase :						str   , lowerCAmelCase :						str):



				"""simple docstring"""


				def load_state_dict(lowerCAmelCase :						str   , lowerCAmelCase :						Dict   , lowerCAmelCase :						Optional[int]=""   , lowerCAmelCase :						Union[str, Any]="relative_position_index"):
								_A :				List[Any]				=  []
								_A :				Optional[Any]				=  []
								_A :				Dict				=  []
								# copy state_dict so _load_from_state_dict can modify it
								_A :				Dict				=  getattr(lowerCAmelCase   , '''_metadata'''   , lowerCAmelCase)
								_A :				List[str]				=  state_dict.copy()
								if metadata is not None:
												_A :				Dict				=  metadata

								def load(lowerCAmelCase :						Dict   , lowerCAmelCase :						Union[str, Any]=""):
												_A :				Union[str, Any]				=  {} if metadata is None else metadata.get(prefix[:-1]   , {})
												module._load_from_state_dict(
												    lowerCAmelCase   , lowerCAmelCase   , lowerCAmelCase   , lowerCAmelCase   , lowerCAmelCase   , lowerCAmelCase   , lowerCAmelCase)
												for name, child in module._modules.items():
																if child is not None:
																				load(lowerCAmelCase   , prefix + name + '''.''')

								load(lowerCAmelCase   , prefix=lowerCAmelCase)

								_A :				List[str]				=  []
								_A :				Dict				=  []
								for key in missing_keys:
												_A :				Optional[int]				=  True
												for ignore_key in ignore_missing.split('''|'''):
																if ignore_key in key:
																				_A :				Dict				=  False
																				break
												if keep_flag:
																warn_missing_keys.append(lowerCAmelCase)
												else:
																ignore_missing_keys.append(lowerCAmelCase)

								_A :				List[str]				=  warn_missing_keys

								if len(lowerCAmelCase) > 0:
												print(
												    '''Weights of {} not initialized from pretrained model: {}'''.format(
												        model.__class__.__name__   , lowerCAmelCase))
								if len(lowerCAmelCase) > 0:
												print('''Weights from pretrained model not used in {}: {}'''.format(model.__class__.__name__   , lowerCAmelCase))
								if len(lowerCAmelCase) > 0:
												print(
												    '''Ignored weights of {} not initialized from pretrained model: {}'''.format(
												        model.__class__.__name__   , lowerCAmelCase))
								if len(lowerCAmelCase) > 0:
												print('''\n'''.join(lowerCAmelCase))

				_A :				str				=  {
				    '''pretrained''': False,
				    '''use_shared_rel_pos_bias''': True,
				    '''use_abs_pos_emb''': False,
				    '''init_values''': 0.1,
				}

				if is_finetuned:
								model_kwargs.update(
								    {
								        '''num_classes''': 1000,
								        '''use_mean_pooling''': True,
								        '''init_scale''': 0.0_0_1,
								        '''use_rel_pos_bias''': True,
								    })

				_A :				Union[str, Any]				=  create_model(
				    '''beit_large_patch16_224''' if is_large else '''beit_base_patch16_224'''   , **lowerCAmelCase   , )
				_A :				Tuple				=  model.patch_embed.patch_size
				_A :				Union[str, Any]				=  (args.input_size // patch_size[0], args.input_size // patch_size[1])
				_A :				Optional[int]				=  torch.load(args.beit_checkpoint   , map_location='''cpu''')

				print(f"""Load ckpt from {args.beit_checkpoint}""")
				_A :				Tuple				=  None
				for model_key in ("model", "module"):
								if model_key in checkpoint:
												_A :				Tuple				=  checkpoint[model_key]
												print(f"""Load state_dict by model_key = {model_key}""")
												break

				_A :				List[str]				=  list(checkpoint_model.keys())
				for key in all_keys:
								if "relative_position_index" in key:
												checkpoint_model.pop(lowerCAmelCase)

								if "relative_position_bias_table" in key:
												_A :				Dict				=  checkpoint_model[key]
												_A	,       _A :				Optional[Any]				=  rel_pos_bias.size()
												_A	,       _A :				Tuple				=  model.state_dict()[key].size()
												_A :				Tuple				=  model.patch_embed.patch_shape
												if dst_patch_shape[0] != dst_patch_shape[1]:
																raise NotImplementedError()

				load_state_dict(lowerCAmelCase   , lowerCAmelCase   , prefix='''''')

				return model





def      lowercase  (  ):



				"""simple docstring"""


				_A :				Tuple				=  get_args()

				_A :				Optional[int]				=  '''ft1k''' in args.hf_checkpoint_name
				_A :				Any				=  '''large''' in args.hf_checkpoint_name

				if is_finetuned:
								# To convert Beit's data2vec_vision to HF you need to copy
								# https://github.com/facebookresearch/data2vec_vision/blob/main/beit/modeling_finetune.py
								# into this folder.
								import modeling_finetune  # noqa: F401
				else:
								# To convert Beit's data2vec_vision to HF you need to copy
								# https://github.com/facebookresearch/data2vec_vision/blob/main/beit/modeling_cyclical.py
								# into this folder
								# IMPORTANT: Note that for now we've only converted the down-stream
								# model and not the full pretrained model. This means for the integration
								# test you need to add a `return x` after the following line:
								# https://github.com/facebookresearch/data2vec_vision/blob/af9a36349aaed59ae66e69b5dabeef2d62fdc5da/beit/modeling_cyclical.py#L197
								# to make the integration test pass.
								import modeling_cyclical  # noqa: F401

				# 1. Create model config
				_A :				Union[str, Any]				=  DataaVecVisionConfig()
				if is_finetuned:
								_A :				List[str]				=  True
								_A :				Union[str, Any]				=  False
								_A :				Tuple				=  True
								_A :				Union[str, Any]				=  1000

								_A :				Optional[int]				=  '''huggingface/label-files'''
								_A :				Optional[int]				=  '''imagenet-1k-id2label.json'''
								_A :				Optional[Any]				=  json.load(open(hf_hub_download(lowerCAmelCase   , lowerCAmelCase   , repo_type='''dataset''')   , '''r'''))
								_A :				Tuple				=  {int(lowerCAmelCase): v for k, v in idalabel.items()}
								_A :				Union[str, Any]				=  idalabel
								_A :				int				=  {v: k for k, v in idalabel.items()}
				else:
								_A :				Any				=  False
								_A :				int				=  True
								_A :				Dict				=  False

				if is_large:
								_A :				Optional[int]				=  1024
								_A :				str				=  4096
								_A :				str				=  24
								_A :				Optional[Any]				=  16

				# 2. Load Beit model
				_A :				Dict				=  load_beit_model(lowerCAmelCase   , lowerCAmelCase   , lowerCAmelCase)
				orig_model.eval()

				# 3. Forward Beit model
				_A :				List[str]				=  BeitImageProcessor(size=config.image_size   , do_center_crop=lowerCAmelCase)
				_A :				Any				=  Image.open('''../../../../tests/fixtures/tests_samples/COCO/000000039769.png''')
				_A :				List[Any]				=  image_processor(images=lowerCAmelCase   , return_tensors='''pt''')
				_A :				str				=  encoding['''pixel_values''']

				_A :				str				=  (pixel_values,) if is_finetuned else (pixel_values, None)
				with torch.no_grad():
								_A :				Tuple				=  orig_model(*lowerCAmelCase)

				# 4. Load HF Data2VecVision model
				if is_finetuned:
								_A :				List[Any]				=  DataaVecVisionForImageClassification(lowerCAmelCase)
								hf_model.eval()
								_A :				int				=  False
								_A :				List[str]				=  '''data2vec_vision.'''
				else:
								_A :				Tuple				=  DataaVecVisionModel(lowerCAmelCase)
								hf_model.eval()
								_A :				Union[str, Any]				=  True
								_A :				List[Any]				=  ''''''

				_A :				List[str]				=  create_rename_keys(lowerCAmelCase   , hf_prefix=lowerCAmelCase   , has_lm_head=lowerCAmelCase)
				_A :				Optional[Any]				=  orig_model.state_dict()
				for src, dest in rename_keys:
								_A :				Dict				=  state_dict.pop(lowerCAmelCase)
								_A :				Any				=  val

				read_in_q_k_v(lowerCAmelCase   , lowerCAmelCase   , hf_prefix=lowerCAmelCase   , has_lm_head=lowerCAmelCase)
				_A	,       _A :				Optional[int]				=  hf_model.load_state_dict(lowerCAmelCase   , strict=lowerCAmelCase)
				print('''HF missing'''   , lowerCAmelCase)
				print('''HF unexpected_keys'''   , lowerCAmelCase)

				# 5. Forward HF Data2VecVision model
				with torch.no_grad():
								_A :				Tuple				=  hf_model(lowerCAmelCase)

				_A :				Dict				=  hf_model_output.logits if is_finetuned else hf_model_output.last_hidden_state

				# 6. Compare
				_A :				Tuple				=  torch.max(torch.abs(hf_output - orig_model_output)).item()

				print(f"""max_absolute_diff = {max_absolute_diff}""")
				_A :				int				=  torch.allclose(lowerCAmelCase   , lowerCAmelCase   , atol=1e-3)
				print('''Do both models output the same tensors?'''   , '''🔥''' if success else '''💩''')
				if not success:
								raise Exception('''Something went wRoNg''')

				# 7. Save
				print(f"""Saving to {args.hf_checkpoint_name}""")
				hf_model.save_pretrained(args.hf_checkpoint_name)
				image_processor.save_pretrained(args.hf_checkpoint_name)


if __name__ == "__main__":
	main()
	# Run the following to convert checkpoints
	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
	#          --beit_checkpoint ./pretrained_base.pt \
	#          --hf_checkpoint_name "./data2vec-vision-base"
	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
	#          --beit_checkpoint ./finetuned_base.pt \
	#          --hf_checkpoint_name "./data2vec-vision-base-ft1k"
	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
	#          --beit_checkpoint ./pretrained_large.pt \
	#          --hf_checkpoint_name "./data2vec-vision-large"
	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
	#          --beit_checkpoint ./finetuned_large.pt \
	#          --hf_checkpoint_name "./data2vec-vision-large-ft1k"







'''simple docstring'''






import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
	import torch

	from transformers import (
	    MODEL_FOR_PRETRAINING_MAPPING,
	    ElectraForCausalLM,
	    ElectraForMaskedLM,
	    ElectraForMultipleChoice,
	    ElectraForPreTraining,
	    ElectraForQuestionAnswering,
	    ElectraForSequenceClassification,
	    ElectraForTokenClassification,
	    ElectraModel,
	)
	from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST



class  lowerCamelCase__				:


	"""simple docstring"""





	def __init__(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__=1_3   ,     UpperCAmelCase__=7   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=9_9   ,     UpperCAmelCase__=3_2   ,     UpperCAmelCase__=5   ,     UpperCAmelCase__=4   ,     UpperCAmelCase__=3_7   ,     UpperCAmelCase__="gelu"   ,     UpperCAmelCase__=0.1   ,     UpperCAmelCase__=0.1   ,     UpperCAmelCase__=5_1_2   ,     UpperCAmelCase__=1_6   ,     UpperCAmelCase__=2   ,     UpperCAmelCase__=0.0_2   ,     UpperCAmelCase__=3   ,     UpperCAmelCase__=4   ,     UpperCAmelCase__=None   ,     )       ->				Optional[Any]:
					_A :				Optional[int]				=  parent
					_A :				List[Any]				=  batch_size
					_A :				Optional[int]				=  seq_length
					_A :				Dict				=  is_training
					_A :				Union[str, Any]				=  use_input_mask
					_A :				Tuple				=  use_token_type_ids
					_A :				str				=  use_labels
					_A :				List[str]				=  vocab_size
					_A :				Tuple				=  hidden_size
					_A :				Tuple				=  num_hidden_layers
					_A :				List[str]				=  num_attention_heads
					_A :				Optional[Any]				=  intermediate_size
					_A :				str				=  hidden_act
					_A :				Optional[int]				=  hidden_dropout_prob
					_A :				int				=  attention_probs_dropout_prob
					_A :				Optional[Any]				=  max_position_embeddings
					_A :				List[Any]				=  type_vocab_size
					_A :				Tuple				=  type_sequence_label_size
					_A :				Any				=  initializer_range
					_A :				Any				=  num_labels
					_A :				str				=  num_choices
					_A :				List[str]				=  scope




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					_A :				List[Any]				=  ids_tensor([self.batch_size, self.seq_length]   ,     self.vocab_size			)

					_A :				Union[str, Any]				=  None
					if self.use_input_mask:
									_A :				Tuple				=  random_attention_mask([self.batch_size, self.seq_length]			)

					_A :				List[str]				=  None
					if self.use_token_type_ids:
									_A :				List[Any]				=  ids_tensor([self.batch_size, self.seq_length]   ,     self.type_vocab_size			)

					_A :				Tuple				=  None
					_A :				Dict				=  None
					_A :				int				=  None
					if self.use_labels:
									_A :				Any				=  ids_tensor([self.batch_size]   ,     self.type_sequence_label_size			)
									_A :				Dict				=  ids_tensor([self.batch_size, self.seq_length]   ,     self.num_labels			)
									_A :				int				=  ids_tensor([self.batch_size]   ,     self.num_choices			)
									_A :				List[str]				=  ids_tensor([self.batch_size, self.seq_length]   ,     1			)

					_A :				List[Any]				=  self.get_config()

					return (
					    config,
					    input_ids,
					    token_type_ids,
					    input_mask,
					    sequence_labels,
					    token_labels,
					    choice_labels,
					    fake_token_labels,
					)




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					return ElectraConfig(
					    vocab_size=self.vocab_size   ,     hidden_size=self.hidden_size   ,     num_hidden_layers=self.num_hidden_layers   ,     num_attention_heads=self.num_attention_heads   ,     intermediate_size=self.intermediate_size   ,     hidden_act=self.hidden_act   ,     hidden_dropout_prob=self.hidden_dropout_prob   ,     attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,     max_position_embeddings=self.max_position_embeddings   ,     type_vocab_size=self.type_vocab_size   ,     is_decoder=UpperCAmelCase__   ,     initializer_range=self.initializer_range   ,     )




	def 				_lowerCamelCase				(					self			)       ->				int:
					(
					    (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       
					) :				List[str]				=  self.prepare_config_and_inputs()

					_A :				Tuple				=  True
					_A :				List[str]				=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
					_A :				str				=  ids_tensor([self.batch_size, self.seq_length]   ,     vocab_size=2			)

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




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Tuple:
					_A :				List[str]				=  ElectraModel(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Tuple				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__			)
					_A :				List[Any]				=  model(UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__			)
					_A :				Tuple				=  model(UpperCAmelCase__			)
					self.parent.assertEqual(result.last_hidden_state.shape   ,     (self.batch_size, self.seq_length, self.hidden_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[str]:
					_A :				Dict				=  True
					_A :				Any				=  ElectraModel(UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				List[str]				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     encoder_hidden_states=UpperCAmelCase__   ,     encoder_attention_mask=UpperCAmelCase__   ,     )
					_A :				Optional[int]				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     encoder_hidden_states=UpperCAmelCase__   ,     )
					_A :				List[Any]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__			)
					self.parent.assertEqual(result.last_hidden_state.shape   ,     (self.batch_size, self.seq_length, self.hidden_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[str]:
					_A :				Optional[int]				=  ElectraForMaskedLM(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				List[str]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length, self.vocab_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Union[str, Any]:
					_A :				List[Any]				=  ElectraForCausalLM(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Optional[int]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length, self.vocab_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[str]:
					_A :				str				=  self.num_labels
					_A :				str				=  ElectraForTokenClassification(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				str				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length, self.num_labels)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[Any]:
					_A :				Optional[int]				=  self.num_labels
					_A :				Optional[int]				=  ElectraForPreTraining(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Dict				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Tuple:
					_A :				Optional[Any]				=  self.num_labels
					_A :				Optional[Any]				=  ElectraForSequenceClassification(UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Any				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.num_labels)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[Any]:
					_A :				Optional[int]				=  ElectraForQuestionAnswering(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Union[str, Any]				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     start_positions=UpperCAmelCase__   ,     end_positions=UpperCAmelCase__   ,     )
					self.parent.assertEqual(result.start_logits.shape   ,     (self.batch_size, self.seq_length)			)
					self.parent.assertEqual(result.end_logits.shape   ,     (self.batch_size, self.seq_length)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Union[str, Any]:
					_A :				Union[str, Any]				=  self.num_choices
					_A :				Union[str, Any]				=  ElectraForMultipleChoice(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Dict				=  input_ids.unsqueeze(1			).expand(-1   ,     self.num_choices   ,     -1			).contiguous()
					_A :				str				=  token_type_ids.unsqueeze(1			).expand(-1   ,     self.num_choices   ,     -1			).contiguous()
					_A :				Dict				=  input_mask.unsqueeze(1			).expand(-1   ,     self.num_choices   ,     -1			).contiguous()
					_A :				str				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__   ,     )
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.num_choices)			)




	def 				_lowerCamelCase				(					self			)       ->				List[str]:
					_A :				Optional[Any]				=  self.prepare_config_and_inputs()
					(
					    (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       
					) :				int				=  config_and_inputs
					_A :				Dict				=  {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
					return config, inputs_dict



@require_torch
class  lowerCamelCase__				(						snake_case_  ,							snake_case_  ,							unittest.TestCase						):


	"""simple docstring"""





	__magic_name__							   =		(
	    (
	        ElectraModel,
	        ElectraForPreTraining,
	        ElectraForMaskedLM,
	        ElectraForCausalLM,
	        ElectraForMultipleChoice,
	        ElectraForTokenClassification,
	        ElectraForSequenceClassification,
	        ElectraForQuestionAnswering,
	    )
	    if is_torch_available()
	    else ()
	)
	__magic_name__							   =		(
	    {
	        """feature-extraction""": ElectraModel,
	        """fill-mask""": ElectraForMaskedLM,
	        """question-answering""": ElectraForQuestionAnswering,
	        """text-classification""": ElectraForSequenceClassification,
	        """text-generation""": ElectraForCausalLM,
	        """token-classification""": ElectraForTokenClassification,
	        """zero-shot""": ElectraForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	__magic_name__							   =		True




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__=False			)       ->				Optional[Any]:
					_A :				List[str]				=  super()._prepare_for_class(UpperCAmelCase__   ,     UpperCAmelCase__   ,     return_labels=UpperCAmelCase__			)

					if return_labels:
									if model_class in get_values(UpperCAmelCase__			):
													_A :				Tuple				=  torch.zeros(
													    (self.model_tester.batch_size, self.model_tester.seq_length)   ,     dtype=torch.long   ,     device=UpperCAmelCase__			)
					return inputs_dict




	def 				_lowerCamelCase				(					self			)       ->				int:
					_A :				Optional[Any]				=  ElectraModelTester(self			)
					_A :				List[str]				=  ConfigTester(self   ,     config_class=UpperCAmelCase__   ,     hidden_size=3_7			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					self.config_tester.run_common_tests()




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					_A :				Union[str, Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_model(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[str]:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_electra_model_as_decoder(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				int				=  self.model_tester.prepare_config_and_inputs()
					for type in ["absolute", "relative_key", "relative_key_query"]:
									_A :				List[str]				=  type
									self.model_tester.create_and_check_electra_model(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_masked_lm(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				Any:
					_A :				Dict				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_token_classification(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				Union[str, Any]:
					_A :				List[str]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_pretraining(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					_A :				Any				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_sequence_classification(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				int:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_question_answering(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_multiple_choice(*UpperCAmelCase__			)




	@slow
	def 				_lowerCamelCase				(					self			)       ->				List[str]:
					for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									_A :				List[str]				=  ElectraModel.from_pretrained(UpperCAmelCase__			)
									self.assertIsNotNone(UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				str				=  self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_electra_for_causal_lm(*UpperCAmelCase__			)



@require_torch
class  lowerCamelCase__				(						unittest.TestCase						):


	"""simple docstring"""





	@slow
	def 				_lowerCamelCase				(					self			)       ->				int:
					_A :				Any				=  ElectraModel.from_pretrained('''google/electra-small-discriminator'''			)
					_A :				Union[str, Any]				=  torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]			)
					_A :				Union[str, Any]				=  torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]			)
					_A :				Optional[int]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__			)[0]
					_A :				Tuple				=  torch.Size((1, 1_1, 2_5_6)			)
					self.assertEqual(output.shape   ,     UpperCAmelCase__			)
					_A :				Optional[Any]				=  torch.tensor(
					    [[[0.4_4_7_1, 0.6_8_2_1, -0.3_2_6_5], [0.4_6_2_7, 0.5_2_5_5, -0.3_6_6_8], [0.4_5_3_2, 0.3_3_1_3, -0.4_3_4_4]]]			)

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