
"""simple docstring"""





import argparse

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)



def 				lowerCamelCase__		(      __snake_case,					__snake_case     )							->		str:


					"""simple docstring"""





					_UpperCamelCase										=  XCLIPTextConfig()

					# derive patch size from model name
					_UpperCamelCase										=  model_name.find('''patch'''     )
					_UpperCamelCase										=  int(model_name[start_idx + len('''patch'''     ) : start_idx + len('''patch'''     ) + 2]     )
					_UpperCamelCase										=  XCLIPVisionConfig(patch_size=__snake_case,					num_frames=__snake_case     )

					if "large" in model_name:
										_UpperCamelCase										=  7_68
										_UpperCamelCase										=  30_72
										_UpperCamelCase										=  12

										_UpperCamelCase										=  10_24
										_UpperCamelCase										=  40_96
										_UpperCamelCase										=  16
										_UpperCamelCase										=  24
										_UpperCamelCase										=  7_68
										_UpperCamelCase										=  30_72

					if model_name == "xclip-large-patch14-16-frames":
										_UpperCamelCase										=  3_36

					_UpperCamelCase										=  XCLIPConfig.from_text_vision_configs(__snake_case,					__snake_case     )

					if "large" in model_name:
										_UpperCamelCase										=  7_68

					return config



def 				lowerCamelCase__		(      __snake_case     )							->		List[Any]:


					"""simple docstring"""





					if name == "token_embedding.weight":
										_UpperCamelCase										=  name.replace('''token_embedding.weight''',					'''text_model.embeddings.token_embedding.weight'''     )
					if name == "positional_embedding":
										_UpperCamelCase										=  name.replace('''positional_embedding''',					'''text_model.embeddings.position_embedding.weight'''     )
					if "ln_1" in name:
										_UpperCamelCase										=  name.replace('''ln_1''',					'''layer_norm1'''     )
					if "ln_2" in name:
										_UpperCamelCase										=  name.replace('''ln_2''',					'''layer_norm2'''     )
					if "c_fc" in name:
										_UpperCamelCase										=  name.replace('''c_fc''',					'''fc1'''     )
					if "c_proj" in name:
										_UpperCamelCase										=  name.replace('''c_proj''',					'''fc2'''     )
					if name.startswith('''transformer.resblocks'''     ):
										_UpperCamelCase										=  name.replace('''transformer.resblocks''',					'''text_model.encoder.layers'''     )
					if "attn.out_proj" in name and "message" not in name:
										_UpperCamelCase										=  name.replace('''attn.out_proj''',					'''self_attn.out_proj'''     )
					if "ln_final" in name:
										_UpperCamelCase										=  name.replace('''ln_final''',					'''text_model.final_layer_norm'''     )
					# visual encoder
					if name == "visual.class_embedding":
										_UpperCamelCase										=  name.replace('''visual.class_embedding''',					'''vision_model.embeddings.class_embedding'''     )
					if name == "visual.positional_embedding":
										_UpperCamelCase										=  name.replace('''visual.positional_embedding''',					'''vision_model.embeddings.position_embedding.weight'''     )
					if name.startswith('''visual.transformer.resblocks'''     ):
										_UpperCamelCase										=  name.replace('''visual.transformer.resblocks''',					'''vision_model.encoder.layers'''     )
					if "visual.conv1" in name:
										_UpperCamelCase										=  name.replace('''visual.conv1''',					'''vision_model.embeddings.patch_embedding'''     )
					if "visual.ln_pre" in name:
										_UpperCamelCase										=  name.replace('''visual.ln_pre''',					'''vision_model.pre_layernorm'''     )
					if "visual.ln_post" in name:
										_UpperCamelCase										=  name.replace('''visual.ln_post''',					'''vision_model.post_layernorm'''     )
					if "visual.proj" in name:
										_UpperCamelCase										=  name.replace('''visual.proj''',					'''visual_projection.weight'''     )
					if "text_projection" in name:
										_UpperCamelCase										=  name.replace('''text_projection''',					'''text_projection.weight'''     )
					# things on top
					if "prompts_visual_proj" in name:
										_UpperCamelCase										=  name.replace('''prompts_visual_proj''',					'''prompts_visual_projection'''     )
					if "prompts_visual_ln" in name:
										_UpperCamelCase										=  name.replace('''prompts_visual_ln''',					'''prompts_visual_layernorm'''     )
					# mit
					if name == "mit.positional_embedding":
										_UpperCamelCase										=  name.replace('''positional''',					'''position'''     )
					if name.startswith('''mit.resblocks'''     ):
										_UpperCamelCase										=  name.replace('''mit.resblocks''',					'''mit.encoder.layers'''     )
					# prompts generator
					if name.startswith('''prompts_generator.norm'''     ):
										_UpperCamelCase										=  name.replace('''prompts_generator.norm''',					'''prompts_generator.layernorm'''     )

					return name



def 				lowerCamelCase__		(      __snake_case,					__snake_case     )							->		Optional[int]:


					"""simple docstring"""





					for key in orig_state_dict.copy().keys():
										_UpperCamelCase										=  orig_state_dict.pop(__snake_case     )

										if "attn.in_proj" in key:
															_UpperCamelCase										=  key.split('''.'''     )
															if key.startswith('''visual'''     ):
																				_UpperCamelCase										=  key_split[3]
																				_UpperCamelCase										=  config.vision_config.hidden_size
																				if "message_attn" in key:
																									if "weight" in key:
																														_UpperCamelCase										=  val[
																														    :dim, :
																														]
																														_UpperCamelCase										=  val[
																														    dim : dim * 2, :
																														]
																														_UpperCamelCase										=  val[
																														    -dim:, :
																														]
																									else:
																														_UpperCamelCase										=  val[
																														    :dim
																														]
																														_UpperCamelCase										=  val[
																														    dim : dim * 2
																														]
																														_UpperCamelCase										=  val[
																														    -dim:
																														]
																				else:
																									if "weight" in key:
																														_UpperCamelCase										=  val[
																														    :dim, :
																														]
																														_UpperCamelCase										=  val[
																														    dim : dim * 2, :
																														]
																														_UpperCamelCase										=  val[
																														    -dim:, :
																														]
																									else:
																														_UpperCamelCase										=  val[:dim]
																														_UpperCamelCase										=  val[
																														    dim : dim * 2
																														]
																														_UpperCamelCase										=  val[-dim:]
															elif key.startswith('''mit'''     ):
																				_UpperCamelCase										=  key_split[2]
																				_UpperCamelCase										=  config.vision_config.mit_hidden_size
																				if "weight" in key:
																									_UpperCamelCase										=  val[:dim, :]
																									_UpperCamelCase										=  val[dim : dim * 2, :]
																									_UpperCamelCase										=  val[-dim:, :]
																				else:
																									_UpperCamelCase										=  val[:dim]
																									_UpperCamelCase										=  val[dim : dim * 2]
																									_UpperCamelCase										=  val[-dim:]
															else:
																				_UpperCamelCase										=  key_split[2]
																				_UpperCamelCase										=  config.text_config.hidden_size
																				if "weight" in key:
																									_UpperCamelCase										=  val[:dim, :]
																									_UpperCamelCase										=  val[
																									    dim : dim * 2, :
																									]
																									_UpperCamelCase										=  val[-dim:, :]
																				else:
																									_UpperCamelCase										=  val[:dim]
																									_UpperCamelCase										=  val[
																									    dim : dim * 2
																									]
																									_UpperCamelCase										=  val[-dim:]
										else:
															_UpperCamelCase										=  rename_key(__snake_case     )
															if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
																				_UpperCamelCase										=  val.T
															_UpperCamelCase										=  val

					return orig_state_dict



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


					"""simple docstring"""





					if num_frames == 8:
										_UpperCamelCase										=  '''eating_spaghetti_8_frames.npy'''
					elif num_frames == 16:
										_UpperCamelCase										=  '''eating_spaghetti.npy'''
					elif num_frames == 32:
										_UpperCamelCase										=  '''eating_spaghetti_32_frames.npy'''
					_UpperCamelCase										=  hf_hub_download(
					    repo_id='''hf-internal-testing/spaghetti-video''',					filename=__snake_case,					repo_type='''dataset''',					)
					_UpperCamelCase										=  np.load(__snake_case     )
					return list(__snake_case     )



def 				lowerCamelCase__		(      __snake_case,					__snake_case=None,					__snake_case=False     )							->		Union[str, Any]:


					"""simple docstring"""





					_UpperCamelCase										=  {
					    # fully supervised kinetics-400 checkpoints
					    '''xclip-base-patch32''': '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth''',
					    '''xclip-base-patch32-16-frames''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth'''
					    ),
					    '''xclip-base-patch16''': '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth''',
					    '''xclip-base-patch16-16-frames''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth'''
					    ),
					    '''xclip-large-patch14''': '''https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb''',
					    '''xclip-large-patch14-16-frames''': '''https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f''',
					    # fully supervised kinetics-600 checkpoints
					    '''xclip-base-patch16-kinetics-600''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth'''
					    ),
					    '''xclip-base-patch16-kinetics-600-16-frames''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth'''
					    ),
					    '''xclip-large-patch14-kinetics-600''': '''https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be''',
					    # few shot
					    '''xclip-base-patch16-hmdb-2-shot''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth'''
					    ),
					    '''xclip-base-patch16-hmdb-4-shot''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth'''
					    ),
					    '''xclip-base-patch16-hmdb-8-shot''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth'''
					    ),
					    '''xclip-base-patch16-hmdb-16-shot''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth'''
					    ),
					    '''xclip-base-patch16-ucf-2-shot''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth'''
					    ),
					    '''xclip-base-patch16-ucf-4-shot''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth'''
					    ),
					    '''xclip-base-patch16-ucf-8-shot''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth'''
					    ),
					    '''xclip-base-patch16-ucf-16-shot''': (
					        '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth'''
					    ),
					    # zero shot
					    '''xclip-base-patch16-zero-shot''': '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth''',
					}

					_UpperCamelCase										=  model_to_url[model_name]
					_UpperCamelCase										=  8
					if "16-frames" in model_name:
										_UpperCamelCase										=  16
					elif "shot" in model_name:
										_UpperCamelCase										=  32

					_UpperCamelCase										=  get_xclip_config(__snake_case,					__snake_case     )
					_UpperCamelCase										=  XCLIPModel(__snake_case     )
					model.eval()

					if "drive" in checkpoint_url:
										_UpperCamelCase										=  '''pytorch_model.bin'''
										gdown.cached_download(__snake_case,					__snake_case,					quiet=__snake_case     )
										_UpperCamelCase										=  torch.load(__snake_case,					map_location='''cpu'''     )['''model''']
					else:
										_UpperCamelCase										=  torch.hub.load_state_dict_from_url(__snake_case     )['''model''']

					_UpperCamelCase										=  convert_state_dict(__snake_case,					__snake_case     )

					_UpperCamelCase										=  XCLIPModel(__snake_case     )
					_UpperCamelCase			,  _UpperCamelCase										=  model.load_state_dict(__snake_case,					strict=__snake_case     )
					assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
					model.eval()

					_UpperCamelCase										=  3_36 if model_name == '''xclip-large-patch14-16-frames''' else 2_24
					_UpperCamelCase										=  VideoMAEImageProcessor(size=__snake_case     )
					_UpperCamelCase										=  CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32'''     )
					_UpperCamelCase										=  CLIPTokenizerFast.from_pretrained('''openai/clip-vit-base-patch32'''     )
					_UpperCamelCase										=  XCLIPProcessor(image_processor=__snake_case,					tokenizer=__snake_case     )

					_UpperCamelCase										=  prepare_video(__snake_case     )
					_UpperCamelCase										=  processor(
					    text=['''playing sports''', '''eating spaghetti''', '''go shopping'''],					videos=__snake_case,					return_tensors='''pt''',					padding=__snake_case     )

					print('''Shape of pixel values:''',					inputs.pixel_values.shape     )

					with torch.no_grad():
										_UpperCamelCase										=  model(**__snake_case     )

					# Verify outputs
					_UpperCamelCase										=  outputs.logits_per_video
					_UpperCamelCase										=  logits_per_video.softmax(dim=1     )
					print('''Probs:''',					__snake_case     )
					# kinetics-400
					if model_name == "xclip-base-patch32":
										_UpperCamelCase										=  torch.tensor([[0.0019, 0.9951, 0.0030]]     )
					elif model_name == "xclip-base-patch32-16-frames":
										_UpperCamelCase										=  torch.tensor([[7.0_999e-04, 9.9_883e-01, 4.5_580e-04]]     )
					elif model_name == "xclip-base-patch16":
										_UpperCamelCase										=  torch.tensor([[0.0083, 0.9681, 0.0236]]     )
					elif model_name == "xclip-base-patch16-16-frames":
										_UpperCamelCase										=  torch.tensor([[7.6_937e-04, 9.9_728e-01, 1.9_473e-03]]     )
					elif model_name == "xclip-large-patch14":
										_UpperCamelCase										=  torch.tensor([[0.0062, 0.9864, 0.0075]]     )
					elif model_name == "xclip-large-patch14-16-frames":
										_UpperCamelCase										=  torch.tensor([[3.3_877e-04, 9.9_937e-01, 2.8_888e-04]]     )
					# kinetics-600
					elif model_name == "xclip-base-patch16-kinetics-600":
										_UpperCamelCase										=  torch.tensor([[0.0555, 0.8914, 0.0531]]     )
					elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
										_UpperCamelCase										=  torch.tensor([[3.8_554e-04, 9.9_929e-01, 3.2_754e-04]]     )
					elif model_name == "xclip-large-patch14-kinetics-600":
										_UpperCamelCase										=  torch.tensor([[0.0036, 0.9920, 0.0045]]     )
					# few shot
					elif model_name == "xclip-base-patch16-hmdb-2-shot":
										_UpperCamelCase										=  torch.tensor([[7.1_890e-06, 9.9_994e-01, 5.6_559e-05]]     )
					elif model_name == "xclip-base-patch16-hmdb-4-shot":
										_UpperCamelCase										=  torch.tensor([[1.0_320e-05, 9.9_993e-01, 6.2_435e-05]]     )
					elif model_name == "xclip-base-patch16-hmdb-8-shot":
										_UpperCamelCase										=  torch.tensor([[4.1_377e-06, 9.9_990e-01, 9.8_386e-05]]     )
					elif model_name == "xclip-base-patch16-hmdb-16-shot":
										_UpperCamelCase										=  torch.tensor([[4.1_347e-05, 9.9_962e-01, 3.3_411e-04]]     )
					elif model_name == "xclip-base-patch16-ucf-2-shot":
										_UpperCamelCase										=  torch.tensor([[8.5_857e-05, 9.9_928e-01, 6.3_291e-04]]     )
					elif model_name == "xclip-base-patch16-ucf-4-shot":
										_UpperCamelCase										=  torch.tensor([[8.5_857e-05, 9.9_928e-01, 6.3_291e-04]]     )
					elif model_name == "xclip-base-patch16-ucf-8-shot":
										_UpperCamelCase										=  torch.tensor([[0.0027, 0.9904, 0.0070]]     )
					elif model_name == "xclip-base-patch16-ucf-16-shot":
										_UpperCamelCase										=  torch.tensor([[9.8_219e-04, 9.9_593e-01, 3.0_863e-03]]     )
					# zero shot
					elif model_name == "xclip-base-patch16-zero-shot":
										_UpperCamelCase										=  torch.tensor([[3.5_082e-04, 9.9_785e-01, 1.7_966e-03]]     )
					else:
										raise ValueError(F'''Model name {model_name} not supported'''     )
					assert torch.allclose(__snake_case,					__snake_case,					atol=1e-3     )
					print('''Looks ok!'''     )

					if pytorch_dump_folder_path is not None:
										print(F'''Saving model {model_name} to {pytorch_dump_folder_path}'''     )
										model.save_pretrained(__snake_case     )

					if push_to_hub:
										print('''Pushing model, processor and slow tokenizer files to the hub...'''     )
										model.push_to_hub(__snake_case,					organization='''nielsr'''     )
										processor.push_to_hub(__snake_case,					organization='''nielsr'''     )
										slow_tokenizer.push_to_hub(__snake_case,					organization='''nielsr'''     )


if __name__ == "__main__":
				_a														=		argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    """--model_name""",
				    default="""xclip-base-patch32""",
				    type=str,
				    help="""Name of the model.""",
				)
				parser.add_argument(
				    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
				)
				parser.add_argument(
				    """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
				)

				_a														=		parser.parse_args()
				convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








"""simple docstring"""





import unittest

from transformers import AlbertConfig, 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, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import (
				    MODEL_FOR_PRETRAINING_MAPPING,
				    AlbertForMaskedLM,
				    AlbertForMultipleChoice,
				    AlbertForPreTraining,
				    AlbertForQuestionAnswering,
				    AlbertForSequenceClassification,
				    AlbertForTokenClassification,
				    AlbertModel,
				)
				from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST






class    _UpperCAmelCase:



		def __init__( self    ,		__a    ,		__a=13    ,		__a=7    ,		__a=True    ,		__a=True    ,		__a=True    ,		__a=True    ,		__a=99    ,		__a=16    ,		__a=36    ,		__a=6    ,		__a=6    ,		__a=6    ,		__a=37    ,		__a="gelu"    ,		__a=0.1    ,		__a=0.1    ,		__a=5_12    ,		__a=16    ,		__a=2    ,		__a=0.02    ,		__a=3    ,		__a=4    ,		__a=None    ,		)	->	List[str]:






							'''simple docstring'''
							_UpperCamelCase										=  parent
							_UpperCamelCase										=  batch_size
							_UpperCamelCase										=  seq_length
							_UpperCamelCase										=  is_training
							_UpperCamelCase										=  use_input_mask
							_UpperCamelCase										=  use_token_type_ids
							_UpperCamelCase										=  use_labels
							_UpperCamelCase										=  vocab_size
							_UpperCamelCase										=  embedding_size
							_UpperCamelCase										=  hidden_size
							_UpperCamelCase										=  num_hidden_layers
							_UpperCamelCase										=  num_hidden_groups
							_UpperCamelCase										=  num_attention_heads
							_UpperCamelCase										=  intermediate_size
							_UpperCamelCase										=  hidden_act
							_UpperCamelCase										=  hidden_dropout_prob
							_UpperCamelCase										=  attention_probs_dropout_prob
							_UpperCamelCase										=  max_position_embeddings
							_UpperCamelCase										=  type_vocab_size
							_UpperCamelCase										=  type_sequence_label_size
							_UpperCamelCase										=  initializer_range
							_UpperCamelCase										=  num_labels
							_UpperCamelCase										=  num_choices
							_UpperCamelCase										=  scope



		def 						UpperCAmelCase							( self)	->	List[str]:






							'''simple docstring'''
							_UpperCamelCase										=  ids_tensor([self.batch_size, self.seq_length]    ,		self.vocab_size)

							_UpperCamelCase										=  None
							if self.use_input_mask:
												_UpperCamelCase										=  random_attention_mask([self.batch_size, self.seq_length])

							_UpperCamelCase										=  None
							if self.use_token_type_ids:
												_UpperCamelCase										=  ids_tensor([self.batch_size, self.seq_length]    ,		self.type_vocab_size)

							_UpperCamelCase										=  None
							_UpperCamelCase										=  None
							_UpperCamelCase										=  None
							if self.use_labels:
												_UpperCamelCase										=  ids_tensor([self.batch_size]    ,		self.type_sequence_label_size)
												_UpperCamelCase										=  ids_tensor([self.batch_size, self.seq_length]    ,		self.num_labels)
												_UpperCamelCase										=  ids_tensor([self.batch_size]    ,		self.num_choices)

							_UpperCamelCase										=  self.get_config()

							return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



		def 						UpperCAmelCase							( self)	->	List[str]:






							'''simple docstring'''
							return AlbertConfig(
							    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    ,		initializer_range=self.initializer_range    ,		num_hidden_groups=self.num_hidden_groups    ,		)



		def 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a)	->	Optional[Any]:






							'''simple docstring'''
							_UpperCamelCase										=  AlbertModel(config=__a)
							model.to(__a)
							model.eval()
							_UpperCamelCase										=  model(__a    ,		attention_mask=__a    ,		token_type_ids=__a)
							_UpperCamelCase										=  model(__a    ,		token_type_ids=__a)
							_UpperCamelCase										=  model(__a)
							self.parent.assertEqual(result.last_hidden_state.shape    ,		(self.batch_size, self.seq_length, self.hidden_size))
							self.parent.assertEqual(result.pooler_output.shape    ,		(self.batch_size, self.hidden_size))



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






							'''simple docstring'''
							_UpperCamelCase										=  AlbertForPreTraining(config=__a)
							model.to(__a)
							model.eval()
							_UpperCamelCase										=  model(
							    __a    ,		attention_mask=__a    ,		token_type_ids=__a    ,		labels=__a    ,		sentence_order_label=__a    ,		)
							self.parent.assertEqual(result.prediction_logits.shape    ,		(self.batch_size, self.seq_length, self.vocab_size))
							self.parent.assertEqual(result.sop_logits.shape    ,		(self.batch_size, config.num_labels))



		def 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a)	->	List[Any]:






							'''simple docstring'''
							_UpperCamelCase										=  AlbertForMaskedLM(config=__a)
							model.to(__a)
							model.eval()
							_UpperCamelCase										=  model(__a    ,		attention_mask=__a    ,		token_type_ids=__a    ,		labels=__a)
							self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.seq_length, self.vocab_size))



		def 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a)	->	Any:






							'''simple docstring'''
							_UpperCamelCase										=  AlbertForQuestionAnswering(config=__a)
							model.to(__a)
							model.eval()
							_UpperCamelCase										=  model(
							    __a    ,		attention_mask=__a    ,		token_type_ids=__a    ,		start_positions=__a    ,		end_positions=__a    ,		)
							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 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a)	->	List[str]:






							'''simple docstring'''
							_UpperCamelCase										=  self.num_labels
							_UpperCamelCase										=  AlbertForSequenceClassification(__a)
							model.to(__a)
							model.eval()
							_UpperCamelCase										=  model(__a    ,		attention_mask=__a    ,		token_type_ids=__a    ,		labels=__a)
							self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.num_labels))



		def 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a)	->	Any:






							'''simple docstring'''
							_UpperCamelCase										=  self.num_labels
							_UpperCamelCase										=  AlbertForTokenClassification(config=__a)
							model.to(__a)
							model.eval()
							_UpperCamelCase										=  model(__a    ,		attention_mask=__a    ,		token_type_ids=__a    ,		labels=__a)
							self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.seq_length, self.num_labels))



		def 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a)	->	Any:






							'''simple docstring'''
							_UpperCamelCase										=  self.num_choices
							_UpperCamelCase										=  AlbertForMultipleChoice(config=__a)
							model.to(__a)
							model.eval()
							_UpperCamelCase										=  input_ids.unsqueeze(1).expand(-1    ,		self.num_choices    ,		-1).contiguous()
							_UpperCamelCase										=  token_type_ids.unsqueeze(1).expand(-1    ,		self.num_choices    ,		-1).contiguous()
							_UpperCamelCase										=  input_mask.unsqueeze(1).expand(-1    ,		self.num_choices    ,		-1).contiguous()
							_UpperCamelCase										=  model(
							    __a    ,		attention_mask=__a    ,		token_type_ids=__a    ,		labels=__a    ,		)
							self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.num_choices))




		def 						UpperCAmelCase							( self)	->	List[str]:






							'''simple docstring'''
							_UpperCamelCase										=  self.prepare_config_and_inputs()
							(
							    (
							    _UpperCamelCase
							)			,  (
							    _UpperCamelCase
							)			,  (
							    _UpperCamelCase
							)			,  (
							    _UpperCamelCase
							)			,  (
							    _UpperCamelCase
							)			,  (
							    _UpperCamelCase
							)			,  (
							    _UpperCamelCase
							)			,  
							)										=  config_and_inputs
							_UpperCamelCase										=  {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
							return config, inputs_dict






@require_torch
class    _UpperCAmelCase(   lowerCamelCase				,					lowerCamelCase				,					unittest.TestCase     ):
		lowercase__      =				(
		    (
		        AlbertModel,
		        AlbertForPreTraining,
		        AlbertForMaskedLM,
		        AlbertForMultipleChoice,
		        AlbertForSequenceClassification,
		        AlbertForTokenClassification,
		        AlbertForQuestionAnswering,
		    )
		    if is_torch_available()
		    else ()
		)
		lowercase__      =				(
		    {
		        'feature-extraction': AlbertModel,
		        'fill-mask': AlbertForMaskedLM,
		        'question-answering': AlbertForQuestionAnswering,
		        'text-classification': AlbertForSequenceClassification,
		        'token-classification': AlbertForTokenClassification,
		        'zero-shot': AlbertForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		lowercase__      =				True



		def 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a=False)	->	Tuple:






							'''simple docstring'''
							_UpperCamelCase										=  super()._prepare_for_class(__a    ,		__a    ,		return_labels=__a)

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



		def 						UpperCAmelCase							( self)	->	Dict:






							'''simple docstring'''
							_UpperCamelCase										=  AlbertModelTester(self)
							_UpperCamelCase										=  ConfigTester(self    ,		config_class=__a    ,		hidden_size=37)



		def 						UpperCAmelCase							( self)	->	Optional[int]:






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



		def 						UpperCAmelCase							( self)	->	List[Any]:






							'''simple docstring'''
							_UpperCamelCase										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__a)



		def 						UpperCAmelCase							( self)	->	int:






							'''simple docstring'''
							_UpperCamelCase										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_pretraining(*__a)



		def 						UpperCAmelCase							( self)	->	Optional[int]:






							'''simple docstring'''
							_UpperCamelCase										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*__a)



		def 						UpperCAmelCase							( self)	->	Any:






							'''simple docstring'''
							_UpperCamelCase										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*__a)



		def 						UpperCAmelCase							( self)	->	Any:






							'''simple docstring'''
							_UpperCamelCase										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*__a)



		def 						UpperCAmelCase							( self)	->	List[Any]:






							'''simple docstring'''
							_UpperCamelCase										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*__a)



		def 						UpperCAmelCase							( self)	->	Dict:






							'''simple docstring'''
							_UpperCamelCase										=  self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
												_UpperCamelCase										=  type
												self.model_tester.create_and_check_model(*__a)




		@slow
		def 						UpperCAmelCase							( self)	->	int:






							'''simple docstring'''
							for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												_UpperCamelCase										=  AlbertModel.from_pretrained(__a)
												self.assertIsNotNone(__a)






@require_torch
class    _UpperCAmelCase(   unittest.TestCase     ):



		@slow
		def 						UpperCAmelCase							( self)	->	Tuple:






							'''simple docstring'''
							_UpperCamelCase										=  AlbertModel.from_pretrained('''albert-base-v2''')
							_UpperCamelCase										=  torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]])
							_UpperCamelCase										=  torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
							with torch.no_grad():
												_UpperCamelCase										=  model(__a    ,		attention_mask=__a)[0]
							_UpperCamelCase										=  torch.Size((1, 11, 7_68))
							self.assertEqual(output.shape    ,		__a)
							_UpperCamelCase										=  torch.tensor(
							    [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]])

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






