






'''simple docstring'''


import argparse

import requests
import torch
from PIL import Image
from torchvision.transforms import Compose, Normalize, Resize, ToTensor

from transformers import SwinaSRConfig, SwinaSRForImageSuperResolution, SwinaSRImageProcessor







def 	__magic_name__( _A					):
						'''simple docstring'''





						UpperCamelCase__         = SwinaSRConfig()

						if "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
												UpperCamelCase__         = 4
						elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
												UpperCamelCase__         = 4
												UpperCamelCase__         = 48
												UpperCamelCase__         = """pixelshuffle_aux"""
						elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
												UpperCamelCase__         = [6, 6, 6, 6]
												UpperCamelCase__         = 60
												UpperCamelCase__         = [6, 6, 6, 6]
												UpperCamelCase__         = """pixelshuffledirect"""
						elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
												UpperCamelCase__         = 4
												UpperCamelCase__         = """nearest+conv"""
						elif "Swin2SR_Jpeg_dynamic" in checkpoint_url:
												UpperCamelCase__         = 1
												UpperCamelCase__         = 1
												UpperCamelCase__         = 126
												UpperCamelCase__         = 7
												UpperCamelCase__         = 2_5_5.0
												UpperCamelCase__         = """"""

						return config







def 	__magic_name__( _A						,      _A					):
						'''simple docstring'''





						if "patch_embed.proj" in name and "layers" not in name:
												UpperCamelCase__         = name.replace("""patch_embed.proj"""						,      """embeddings.patch_embeddings.projection"""					)
						if "patch_embed.norm" in name:
												UpperCamelCase__         = name.replace("""patch_embed.norm"""						,      """embeddings.patch_embeddings.layernorm"""					)
						if "layers" in name:
												UpperCamelCase__         = name.replace("""layers"""						,      """encoder.stages"""					)
						if "residual_group.blocks" in name:
												UpperCamelCase__         = name.replace("""residual_group.blocks"""						,      """layers"""					)
						if "attn.proj" in name:
												UpperCamelCase__         = name.replace("""attn.proj"""						,      """attention.output.dense"""					)
						if "attn" in name:
												UpperCamelCase__         = name.replace("""attn"""						,      """attention.self"""					)
						if "norm1" in name:
												UpperCamelCase__         = name.replace("""norm1"""						,      """layernorm_before"""					)
						if "norm2" in name:
												UpperCamelCase__         = name.replace("""norm2"""						,      """layernorm_after"""					)
						if "mlp.fc1" in name:
												UpperCamelCase__         = name.replace("""mlp.fc1"""						,      """intermediate.dense"""					)
						if "mlp.fc2" in name:
												UpperCamelCase__         = name.replace("""mlp.fc2"""						,      """output.dense"""					)
						if "q_bias" in name:
												UpperCamelCase__         = name.replace("""q_bias"""						,      """query.bias"""					)
						if "k_bias" in name:
												UpperCamelCase__         = name.replace("""k_bias"""						,      """key.bias"""					)
						if "v_bias" in name:
												UpperCamelCase__         = name.replace("""v_bias"""						,      """value.bias"""					)
						if "cpb_mlp" in name:
												UpperCamelCase__         = name.replace("""cpb_mlp"""						,      """continuous_position_bias_mlp"""					)
						if "patch_embed.proj" in name:
												UpperCamelCase__         = name.replace("""patch_embed.proj"""						,      """patch_embed.projection"""					)

						if name == "norm.weight":
												UpperCamelCase__         = """layernorm.weight"""
						if name == "norm.bias":
												UpperCamelCase__         = """layernorm.bias"""

						if "conv_first" in name:
												UpperCamelCase__         = name.replace("""conv_first"""						,      """first_convolution"""					)

						if (
						    "upsample" in name
						    or "conv_before_upsample" in name
						    or "conv_bicubic" in name
						    or "conv_up" in name
						    or "conv_hr" in name
						    or "conv_last" in name
						    or "aux" in name
						):
												# heads
												if "conv_last" in name:
																		UpperCamelCase__         = name.replace("""conv_last"""						,      """final_convolution"""					)
												if config.upsampler in ["pixelshuffle", "pixelshuffle_aux", "nearest+conv"]:
																		if "conv_before_upsample.0" in name:
																								UpperCamelCase__         = name.replace("""conv_before_upsample.0"""						,      """conv_before_upsample"""					)
																		if "upsample.0" in name:
																								UpperCamelCase__         = name.replace("""upsample.0"""						,      """upsample.convolution_0"""					)
																		if "upsample.2" in name:
																								UpperCamelCase__         = name.replace("""upsample.2"""						,      """upsample.convolution_1"""					)
																		UpperCamelCase__         = """upsample.""" + name
												elif config.upsampler == "pixelshuffledirect":
																		UpperCamelCase__         = name.replace("""upsample.0.weight"""						,      """upsample.conv.weight"""					)
																		UpperCamelCase__         = name.replace("""upsample.0.bias"""						,      """upsample.conv.bias"""					)
												else:
																		pass
						else:
												UpperCamelCase__         = """swin2sr.""" + name

						return name







def 	__magic_name__( _A						,      _A					):
						'''simple docstring'''





						for key in orig_state_dict.copy().keys():
												UpperCamelCase__         = orig_state_dict.pop(_A					)

												if "qkv" in key:
																		UpperCamelCase__         = key.split("""."""					)
																		UpperCamelCase__         = int(key_split[1]					)
																		UpperCamelCase__         = int(key_split[4]					)
																		UpperCamelCase__         = config.embed_dim

																		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:]
																		pass
												else:
																		UpperCamelCase__         = val

						return orig_state_dict







def 	__magic_name__( _A						,      _A						,      _A					):
						'''simple docstring'''





						UpperCamelCase__         = get_config(_A					)
						UpperCamelCase__         = SwinaSRForImageSuperResolution(_A					)
						model.eval()

						UpperCamelCase__         = torch.hub.load_state_dict_from_url(_A						,      map_location="""cpu"""					)
						UpperCamelCase__         = convert_state_dict(_A						,      _A					)
						UpperCamelCase__    ,   UpperCamelCase__         = model.load_state_dict(_A						,      strict=_A					)

						if len(_A					) > 0:
												raise ValueError("""Missing keys when converting: {}""".format(_A					)					)
						for key in unexpected_keys:
												if not ("relative_position_index" in key or "relative_coords_table" in key or "self_mask" in key):
																		raise ValueError(f"Unexpected key {key} in state_dict"					)

    # verify values
						UpperCamelCase__         = """https://github.com/mv-lab/swin2sr/blob/main/testsets/real-inputs/shanghai.jpg?raw=true"""
						UpperCamelCase__         = Image.open(requests.get(_A						,      stream=_A					).raw					).convert("""RGB"""					)
						UpperCamelCase__         = SwinaSRImageProcessor()
						# pixel_values = processor(image, return_tensors="pt").pixel_values

						UpperCamelCase__         = 126 if """Jpeg""" in checkpoint_url else 256
						UpperCamelCase__         = Compose(
						    [
						        Resize((image_size, image_size)					),
						        ToTensor(),
						        Normalize(mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]						,      std=[0.2_2_9, 0.2_2_4, 0.2_2_5]					),
						    ]					)
						UpperCamelCase__         = transforms(_A					).unsqueeze(0					)

						if config.num_channels == 1:
												UpperCamelCase__         = pixel_values[:, 0, :, :].unsqueeze(1					)

						UpperCamelCase__         = model(_A					)

						# assert values
						if "Swin2SR_ClassicalSR_X2_64" in checkpoint_url:
												UpperCamelCase__         = torch.Size([1, 3, 512, 512]					)
												UpperCamelCase__         = torch.tensor(
												    [[-0.7_0_8_7, -0.7_1_3_8, -0.6_7_2_1], [-0.8_3_4_0, -0.8_0_9_5, -0.7_2_9_8], [-0.9_1_4_9, -0.8_4_1_4, -0.7_9_4_0]]					)
						elif "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
												UpperCamelCase__         = torch.Size([1, 3, 1024, 1024]					)
												UpperCamelCase__         = torch.tensor(
												    [[-0.7_7_7_5, -0.8_1_0_5, -0.8_9_3_3], [-0.7_7_6_4, -0.8_3_5_6, -0.9_2_2_5], [-0.7_9_7_6, -0.8_6_8_6, -0.9_5_7_9]]					)
						elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
												# TODO values didn't match exactly here
												UpperCamelCase__         = torch.Size([1, 3, 1024, 1024]					)
												UpperCamelCase__         = torch.tensor(
												    [[-0.8_0_3_5, -0.7_5_0_4, -0.7_4_9_1], [-0.8_5_3_8, -0.8_1_2_4, -0.7_7_8_2], [-0.8_8_0_4, -0.8_6_5_1, -0.8_4_9_3]]					)
						elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
												UpperCamelCase__         = torch.Size([1, 3, 512, 512]					)
												UpperCamelCase__         = torch.tensor(
												    [[-0.7_6_6_9, -0.8_6_6_2, -0.8_7_6_7], [-0.8_8_1_0, -0.9_9_6_2, -0.9_8_2_0], [-0.9_3_4_0, -1.0_3_2_2, -1.1_1_4_9]]					)
						elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
												UpperCamelCase__         = torch.Size([1, 3, 1024, 1024]					)
												UpperCamelCase__         = torch.tensor(
												    [[-0.5_2_3_8, -0.5_5_5_7, -0.6_3_2_1], [-0.6_0_1_6, -0.5_9_0_3, -0.6_3_9_1], [-0.6_2_4_4, -0.6_3_3_4, -0.6_8_8_9]]					)

						assert (
						    outputs.reconstruction.shape == expected_shape
						), f"Shape of reconstruction should be {expected_shape}, but is {outputs.reconstruction.shape}"
						assert torch.allclose(outputs.reconstruction[0, 0, :3, :3]						,      _A						,      atol=1e-3					)
						print("""Looks ok!"""					)

						UpperCamelCase__         = {
						    """https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X2_64.pth""": (
						        """swin2SR-classical-sr-x2-64"""
						    ),
						    """https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X4_64.pth""": (
						        """swin2SR-classical-sr-x4-64"""
						    ),
						    """https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_CompressedSR_X4_48.pth""": (
						        """swin2SR-compressed-sr-x4-48"""
						    ),
						    """https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_Lightweight_X2_64.pth""": (
						        """swin2SR-lightweight-x2-64"""
						    ),
						    """https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR.pth""": (
						        """swin2SR-realworld-sr-x4-64-bsrgan-psnr"""
						    ),
						}
						UpperCamelCase__         = url_to_name[checkpoint_url]

						if pytorch_dump_folder_path is not None:
												print(f"Saving model {model_name} to {pytorch_dump_folder_path}"					)
												model.save_pretrained(_A					)
												print(f"Saving image processor to {pytorch_dump_folder_path}"					)
												processor.save_pretrained(_A					)

						if push_to_hub:
												model.push_to_hub(f"caidas/{model_name}"					)
												processor.push_to_hub(f"caidas/{model_name}"					)


if __name__ == "__main__":
					lowerCamelCase_				: Union[str, Any]    	=							argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '''--checkpoint_url''',
					    default='''https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X2_64.pth''',
					    type=str,
					    help='''URL of the original Swin2SR checkpoint you\'d like to convert.''',
					)
					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 to push the converted model to the hub.''')

					lowerCamelCase_				: Union[str, Any]    	=							parser.parse_args()
					convert_swinasr_checkpoint(args.checkpoint_url, 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    _SCREAMING_SNAKE_CASE							:


			'''simple docstring'''





			def __init__(   self							:						Optional[int] ,      lowercase							:						List[str] ,      lowercase							:						List[Any]=1_3 ,      lowercase							:						Union[str, Any]=7 ,      lowercase							:						Dict=True ,      lowercase							:						Optional[int]=True ,      lowercase							:						List[Any]=True ,      lowercase							:						Dict=True ,      lowercase							:						List[str]=9_9 ,      lowercase							:						Dict=1_6 ,      lowercase							:						Dict=3_6 ,      lowercase							:						str=6 ,      lowercase							:						List[Any]=6 ,      lowercase							:						int=6 ,      lowercase							:						Union[str, Any]=3_7 ,      lowercase							:						Union[str, Any]="gelu" ,      lowercase							:						List[Any]=0.1 ,      lowercase							:						List[str]=0.1 ,      lowercase							:						str=5_1_2 ,      lowercase							:						Any=1_6 ,      lowercase							:						str=2 ,      lowercase							:						List[Any]=0.0_2 ,      lowercase							:						Tuple=3 ,      lowercase							:						Dict=4 ,      lowercase							:						Dict=None ,      )							->    List[Any]:

									'''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 				A      (   self							:						List[str]			)							->    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 				A      (   self							:						Optional[int]			)							->    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 				A      (   self							:						List[str] ,      lowercase							:						Union[str, Any] ,      lowercase							:						Optional[Any] ,      lowercase							:						Dict ,      lowercase							:						str ,      lowercase							:						int ,      lowercase							:						List[str] ,      lowercase							:						Optional[Any]			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = AlbertModel(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = model(lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase			)
									UpperCamelCase__         = model(lowercase ,      token_type_ids=lowercase			)
									UpperCamelCase__         = model(lowercase			)
									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 				A      (   self							:						int ,      lowercase							:						List[str] ,      lowercase							:						Optional[int] ,      lowercase							:						Optional[Any] ,      lowercase							:						Optional[Any] ,      lowercase							:						Optional[int] ,      lowercase							:						List[Any] ,      lowercase							:						int			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = AlbertForPreTraining(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = model(
									    lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase ,      labels=lowercase ,      sentence_order_label=lowercase ,      )
									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 				A      (   self							:						Optional[int] ,      lowercase							:						Dict ,      lowercase							:						List[Any] ,      lowercase							:						Any ,      lowercase							:						List[str] ,      lowercase							:						int ,      lowercase							:						Any ,      lowercase							:						Dict			)							->    Any:

									'''simple docstring'''







									UpperCamelCase__         = AlbertForMaskedLM(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = model(lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase ,      labels=lowercase			)
									self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.seq_length, self.vocab_size)			)






			def 				A      (   self							:						Optional[Any] ,      lowercase							:						Any ,      lowercase							:						Union[str, Any] ,      lowercase							:						List[Any] ,      lowercase							:						List[Any] ,      lowercase							:						List[str] ,      lowercase							:						Optional[Any] ,      lowercase							:						Optional[Any]			)							->    List[Any]:

									'''simple docstring'''







									UpperCamelCase__         = AlbertForQuestionAnswering(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = model(
									    lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase ,      start_positions=lowercase ,      end_positions=lowercase ,      )
									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 				A      (   self							:						Any ,      lowercase							:						Optional[Any] ,      lowercase							:						Any ,      lowercase							:						Dict ,      lowercase							:						Any ,      lowercase							:						Optional[int] ,      lowercase							:						str ,      lowercase							:						Dict			)							->    List[Any]:

									'''simple docstring'''







									UpperCamelCase__         = self.num_labels
									UpperCamelCase__         = AlbertForSequenceClassification(lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = model(lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase ,      labels=lowercase			)
									self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.num_labels)			)






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

									'''simple docstring'''







									UpperCamelCase__         = self.num_labels
									UpperCamelCase__         = AlbertForTokenClassification(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = model(lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase ,      labels=lowercase			)
									self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.seq_length, self.num_labels)			)






			def 				A      (   self							:						List[Any] ,      lowercase							:						Union[str, Any] ,      lowercase							:						List[Any] ,      lowercase							:						Any ,      lowercase							:						List[str] ,      lowercase							:						Optional[Any] ,      lowercase							:						List[str] ,      lowercase							:						List[Any]			)							->    str:

									'''simple docstring'''







									UpperCamelCase__         = self.num_choices
									UpperCamelCase__         = AlbertForMultipleChoice(config=lowercase			)
									model.to(lowercase			)
									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(
									    lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase ,      labels=lowercase ,      )
									self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.num_choices)			)







			def 				A      (   self							:						Optional[int]			)							->    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    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       ,unittest.TestCase	):


			'''simple docstring'''





			__a							:							Any					     =					(
			    (
			        AlbertModel,
			        AlbertForPreTraining,
			        AlbertForMaskedLM,
			        AlbertForMultipleChoice,
			        AlbertForSequenceClassification,
			        AlbertForTokenClassification,
			        AlbertForQuestionAnswering,
			    )
			    if is_torch_available()
			    else ()
			)
			__a							:							List[Any]					     =					(
			    {
			        "feature-extraction": AlbertModel,
			        "fill-mask": AlbertForMaskedLM,
			        "question-answering": AlbertForQuestionAnswering,
			        "text-classification": AlbertForSequenceClassification,
			        "token-classification": AlbertForTokenClassification,
			        "zero-shot": AlbertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			__a							:							List[str]					     =					True






			def 				A      (   self							:						Any ,      lowercase							:						int ,      lowercase							:						Dict ,      lowercase							:						Optional[Any]=False			)							->    Tuple:

									'''simple docstring'''







									UpperCamelCase__         = super()._prepare_for_class(lowercase ,      lowercase ,      return_labels=lowercase			)

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






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

									'''simple docstring'''







									UpperCamelCase__         = AlbertModelTester(self			)
									UpperCamelCase__         = ConfigTester(self ,      config_class=lowercase ,      hidden_size=3_7			)






			def 				A      (   self							:						int			)							->    List[str]:

									'''simple docstring'''







									self.config_tester.run_common_tests()






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

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowercase			)






			def 				A      (   self							:						int			)							->    Dict:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*lowercase			)






			def 				A      (   self							:						Tuple			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*lowercase			)






			def 				A      (   self							:						Any			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*lowercase			)






			def 				A      (   self							:						List[str]			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*lowercase			)






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

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*lowercase			)






			def 				A      (   self							:						Tuple			)							->    List[str]:

									'''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(*lowercase			)







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

									'''simple docstring'''







									for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															UpperCamelCase__         = AlbertModel.from_pretrained(lowercase			)
															self.assertIsNotNone(lowercase			)
@require_torch
class    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):


			'''simple docstring'''





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

									'''simple docstring'''







									UpperCamelCase__         = AlbertModel.from_pretrained("""albert-base-v2"""			)
									UpperCamelCase__         = 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]]			)
									UpperCamelCase__         = torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]			)
									with torch.no_grad():
															UpperCamelCase__         = model(lowercase ,      attention_mask=lowercase			)[0]
									UpperCamelCase__         = torch.Size((1, 1_1, 7_6_8)			)
									self.assertEqual(output.shape ,      lowercase			)
									UpperCamelCase__         = torch.tensor(
									    [[[-0.6_5_1_3, 1.5_0_3_5, -0.2_7_6_6], [-0.6_5_1_5, 1.5_0_4_6, -0.2_7_8_0], [-0.6_5_1_2, 1.5_0_4_9, -0.2_7_8_4]]]			)

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