



'''simple docstring'''





import argparse
import json

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

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)

def 			a_	(				_lowerCAmelCase			) -> List[str]:
  __lowerCamelCase			:			Union[str, Any]       			=					VideoMAEConfig()

  set_architecture_configs(_lowerCAmelCase			,_lowerCAmelCase			)

  if "finetuned" not in model_name:
    __lowerCamelCase			:			Union[str, Any]       			=					False

  if "finetuned" in model_name:
    __lowerCamelCase			:			List[Any]       			=					'huggingface/label-files'
    if "kinetics" in model_name:
      __lowerCamelCase			:			List[str]       			=					400
      __lowerCamelCase			:			Optional[int]       			=					'kinetics400-id2label.json'
    elif "ssv2" in model_name:
      __lowerCamelCase			:			Dict       			=					174
      __lowerCamelCase			:			List[Any]       			=					'something-something-v2-id2label.json'
    else:
      raise ValueError('Model name should either contain \'kinetics\' or \'ssv2\' in case it\'s fine-tuned.'			)
    __lowerCamelCase			:			Union[str, Any]       			=					json.load(open(hf_hub_download(_lowerCAmelCase			,_lowerCAmelCase			,repo_type='dataset'			)			,'r'			)			)
    __lowerCamelCase			:			Dict       			=					{int(_lowerCAmelCase			): v for k, v in idalabel.items()}
    __lowerCamelCase			:			Any       			=					idalabel
    __lowerCamelCase			:			Optional[Any]       			=					{v: k for k, v in idalabel.items()}

  return config

def 			a_	(				_lowerCAmelCase			,_lowerCAmelCase			) -> Tuple:
  if "small" in model_name:
    __lowerCamelCase			:			int       			=					384
    __lowerCamelCase			:			Union[str, Any]       			=					1536
    __lowerCamelCase			:			List[str]       			=					12
    __lowerCamelCase			:			Tuple       			=					16
    __lowerCamelCase			:			Optional[int]       			=					12
    __lowerCamelCase			:			List[str]       			=					3
    __lowerCamelCase			:			Union[str, Any]       			=					192
    __lowerCamelCase			:			Tuple       			=					768
  elif "large" in model_name:
    __lowerCamelCase			:			Tuple       			=					1024
    __lowerCamelCase			:			Tuple       			=					4096
    __lowerCamelCase			:			Dict       			=					24
    __lowerCamelCase			:			Optional[int]       			=					16
    __lowerCamelCase			:			Union[str, Any]       			=					12
    __lowerCamelCase			:			str       			=					8
    __lowerCamelCase			:			Optional[Any]       			=					512
    __lowerCamelCase			:			int       			=					2048
  elif "huge" in model_name:
    __lowerCamelCase			:			Union[str, Any]       			=					1280
    __lowerCamelCase			:			Tuple       			=					5120
    __lowerCamelCase			:			Dict       			=					32
    __lowerCamelCase			:			Tuple       			=					16
    __lowerCamelCase			:			Dict       			=					12
    __lowerCamelCase			:			str       			=					8
    __lowerCamelCase			:			List[Any]       			=					640
    __lowerCamelCase			:			Optional[int]       			=					2560
  elif "base" not in model_name:
    raise ValueError('Model name should include either "small", "base", "large", or "huge"'			)

def 			a_	(				_lowerCAmelCase			) -> Tuple:
  if "encoder." in name:
    __lowerCamelCase			:			Optional[int]       			=					name.replace('encoder.'			,''			)
  if "cls_token" in name:
    __lowerCamelCase			:			Union[str, Any]       			=					name.replace('cls_token'			,'videomae.embeddings.cls_token'			)
  if "decoder_pos_embed" in name:
    __lowerCamelCase			:			Optional[int]       			=					name.replace('decoder_pos_embed'			,'decoder.decoder_pos_embed'			)
  if "pos_embed" in name and "decoder" not in name:
    __lowerCamelCase			:			Dict       			=					name.replace('pos_embed'			,'videomae.embeddings.position_embeddings'			)
  if "patch_embed.proj" in name:
    __lowerCamelCase			:			Any       			=					name.replace('patch_embed.proj'			,'videomae.embeddings.patch_embeddings.projection'			)
  if "patch_embed.norm" in name:
    __lowerCamelCase			:			Tuple       			=					name.replace('patch_embed.norm'			,'videomae.embeddings.norm'			)
  if "decoder.blocks" in name:
    __lowerCamelCase			:			Optional[Any]       			=					name.replace('decoder.blocks'			,'decoder.decoder_layers'			)
  if "blocks" in name:
    __lowerCamelCase			:			Optional[Any]       			=					name.replace('blocks'			,'videomae.encoder.layer'			)
  if "attn.proj" in name:
    __lowerCamelCase			:			List[str]       			=					name.replace('attn.proj'			,'attention.output.dense'			)
  if "attn" in name and "bias" not in name:
    __lowerCamelCase			:			Tuple       			=					name.replace('attn'			,'attention.self'			)
  if "attn" in name:
    __lowerCamelCase			:			Tuple       			=					name.replace('attn'			,'attention.attention'			)
  if "norm1" in name:
    __lowerCamelCase			:			Optional[int]       			=					name.replace('norm1'			,'layernorm_before'			)
  if "norm2" in name:
    __lowerCamelCase			:			Any       			=					name.replace('norm2'			,'layernorm_after'			)
  if "mlp.fc1" in name:
    __lowerCamelCase			:			Dict       			=					name.replace('mlp.fc1'			,'intermediate.dense'			)
  if "mlp.fc2" in name:
    __lowerCamelCase			:			Dict       			=					name.replace('mlp.fc2'			,'output.dense'			)
  if "decoder_embed" in name:
    __lowerCamelCase			:			Union[str, Any]       			=					name.replace('decoder_embed'			,'decoder.decoder_embed'			)
  if "decoder_norm" in name:
    __lowerCamelCase			:			List[str]       			=					name.replace('decoder_norm'			,'decoder.decoder_norm'			)
  if "decoder_pred" in name:
    __lowerCamelCase			:			List[str]       			=					name.replace('decoder_pred'			,'decoder.decoder_pred'			)
  if "norm.weight" in name and "decoder" not in name and "fc" not in name:
    __lowerCamelCase			:			Optional[Any]       			=					name.replace('norm.weight'			,'videomae.layernorm.weight'			)
  if "norm.bias" in name and "decoder" not in name and "fc" not in name:
    __lowerCamelCase			:			Any       			=					name.replace('norm.bias'			,'videomae.layernorm.bias'			)
  if "head" in name and "decoder" not in name:
    __lowerCamelCase			:			str       			=					name.replace('head'			,'classifier'			)

  return name

def 			a_	(				_lowerCAmelCase			,_lowerCAmelCase			) -> Optional[int]:
  for key in orig_state_dict.copy().keys():
    __lowerCamelCase			:			str       			=					orig_state_dict.pop(_lowerCAmelCase			)

    if key.startswith('encoder.'			):
      __lowerCamelCase			:			Union[str, Any]       			=					key.replace('encoder.'			,''			)

    if "qkv" in key:
      __lowerCamelCase			:			Union[str, Any]       			=					key.split('.'			)
      if key.startswith('decoder.blocks'			):
        __lowerCamelCase			:			Optional[Any]       			=					config.decoder_hidden_size
        __lowerCamelCase			:			Optional[Any]       			=					int(key_split[2]			)
        __lowerCamelCase			:			Any       			=					'decoder.decoder_layers.'
        if "weight" in key:
          __lowerCamelCase			:			List[str]       			=					val[:dim, :]
          __lowerCamelCase			:			Optional[Any]       			=					val[dim : dim * 2, :]
          __lowerCamelCase			:			List[Any]       			=					val[-dim:, :]
      else:
        __lowerCamelCase			:			List[str]       			=					config.hidden_size
        __lowerCamelCase			:			Union[str, Any]       			=					int(key_split[1]			)
        __lowerCamelCase			:			List[Any]       			=					'videomae.encoder.layer.'
        if "weight" in key:
          __lowerCamelCase			:			int       			=					val[:dim, :]
          __lowerCamelCase			:			List[Any]       			=					val[dim : dim * 2, :]
          __lowerCamelCase			:			Tuple       			=					val[-dim:, :]
    else:
      __lowerCamelCase			:			int       			=					val

  return orig_state_dict

def 			a_	(				) -> Dict:
  __lowerCamelCase			:			List[str]       			=					hf_hub_download(
      repo_id='hf-internal-testing/spaghetti-video'			,filename='eating_spaghetti.npy'			,repo_type='dataset'			)
  __lowerCamelCase			:			Union[str, Any]       			=					np.load(_lowerCAmelCase			)
  return list(_lowerCAmelCase			)

def 			a_	(				_lowerCAmelCase			,_lowerCAmelCase			,_lowerCAmelCase			,_lowerCAmelCase			) -> Any:
  __lowerCamelCase			:			List[str]       			=					get_videomae_config(_lowerCAmelCase			)

  if "finetuned" in model_name:
    __lowerCamelCase			:			str       			=					VideoMAEForVideoClassification(_lowerCAmelCase			)
  else:
    __lowerCamelCase			:			Any       			=					VideoMAEForPreTraining(_lowerCAmelCase			)

  # download original checkpoint, hosted on Google Drive
  __lowerCamelCase			:			Optional[int]       			=					'pytorch_model.bin'
  gdown.cached_download(_lowerCAmelCase			,_lowerCAmelCase			,quiet=_lowerCAmelCase			)
  __lowerCamelCase			:			List[str]       			=					torch.load(_lowerCAmelCase			,map_location='cpu'			)
  if "model" in files:
    __lowerCamelCase			:			Optional[int]       			=					files['model']
  else:
    __lowerCamelCase			:			str       			=					files['module']
  __lowerCamelCase			:			Optional[Any]       			=					convert_state_dict(_lowerCAmelCase			,_lowerCAmelCase			)

  model.load_state_dict(_lowerCAmelCase			)
  model.eval()

  # verify model on basic input
  __lowerCamelCase			:			Union[str, Any]       			=					VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]			,image_std=[0.5, 0.5, 0.5]			)
  __lowerCamelCase			:			Any       			=					prepare_video()
  __lowerCamelCase			:			int       			=					image_processor(_lowerCAmelCase			,return_tensors='pt'			)

  if "finetuned" not in model_name:
    __lowerCamelCase			:			Union[str, Any]       			=					hf_hub_download(repo_id='hf-internal-testing/bool-masked-pos'			,filename='bool_masked_pos.pt'			)
    __lowerCamelCase			:			Tuple       			=					torch.load(_lowerCAmelCase			)

  __lowerCamelCase			:			Union[str, Any]       			=					model(**_lowerCAmelCase			)
  __lowerCamelCase			:			List[str]       			=					outputs.logits

  __lowerCamelCase			:			Dict       			=					[
      'videomae-small-finetuned-kinetics',
      'videomae-small-finetuned-ssv2',
      # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
      'videomae-base-short',
      'videomae-base-short-finetuned-kinetics',
      'videomae-base',
      'videomae-base-finetuned-kinetics',
      'videomae-large',
      'videomae-large-finetuned-kinetics',
      'videomae-huge-finetuned-kinetics',
      # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
      'videomae-base-short-ssv2',
      'videomae-base-short-finetuned-ssv2',
      'videomae-base-ssv2',
      'videomae-base-finetuned-ssv2',
  ]

  # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
  if model_name == "videomae-small-finetuned-kinetics":
    __lowerCamelCase			:			Optional[Any]       			=					torch.Size([1, 400]			)
    __lowerCamelCase			:			str       			=					torch.tensor([-0.9291, -0.4061, -0.9307]			)
  elif model_name == "videomae-small-finetuned-ssv2":
    __lowerCamelCase			:			List[Any]       			=					torch.Size([1, 174]			)
    __lowerCamelCase			:			Any       			=					torch.tensor([0.2671, -0.4689, -0.8235]			)
  elif model_name == "videomae-base":
    __lowerCamelCase			:			Optional[int]       			=					torch.Size([1, 1408, 1536]			)
    __lowerCamelCase			:			Optional[int]       			=					torch.tensor([[0.7739, 0.7968, 0.7089], [0.6701, 0.7487, 0.6209], [0.4287, 0.5158, 0.4773]]			)
  elif model_name == "videomae-base-short":
    __lowerCamelCase			:			Any       			=					torch.Size([1, 1408, 1536]			)
    __lowerCamelCase			:			Any       			=					torch.tensor([[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]]			)
    # we verified the loss both for normalized and unnormalized targets for this one
    __lowerCamelCase			:			Optional[Any]       			=					torch.tensor([0.5142]			) if config.norm_pix_loss else torch.tensor([0.6469]			)
  elif model_name == "videomae-large":
    __lowerCamelCase			:			Optional[int]       			=					torch.Size([1, 1408, 1536]			)
    __lowerCamelCase			:			List[Any]       			=					torch.tensor([[0.7149, 0.7997, 0.6966], [0.6768, 0.7869, 0.6948], [0.5139, 0.6221, 0.5605]]			)
  elif model_name == "videomae-large-finetuned-kinetics":
    __lowerCamelCase			:			str       			=					torch.Size([1, 400]			)
    __lowerCamelCase			:			Union[str, Any]       			=					torch.tensor([0.0771, 0.0011, -0.3625]			)
  elif model_name == "videomae-huge-finetuned-kinetics":
    __lowerCamelCase			:			List[str]       			=					torch.Size([1, 400]			)
    __lowerCamelCase			:			Dict       			=					torch.tensor([0.2433, 0.1632, -0.4894]			)
  elif model_name == "videomae-base-short-finetuned-kinetics":
    __lowerCamelCase			:			Optional[int]       			=					torch.Size([1, 400]			)
    __lowerCamelCase			:			Any       			=					torch.tensor([0.6588, 0.0990, -0.2493]			)
  elif model_name == "videomae-base-finetuned-kinetics":
    __lowerCamelCase			:			Optional[int]       			=					torch.Size([1, 400]			)
    __lowerCamelCase			:			Tuple       			=					torch.tensor([0.3669, -0.0688, -0.2421]			)
  elif model_name == "videomae-base-short-ssv2":
    __lowerCamelCase			:			Optional[Any]       			=					torch.Size([1, 1408, 1536]			)
    __lowerCamelCase			:			List[Any]       			=					torch.tensor([[0.4712, 0.5296, 0.5786], [0.2278, 0.2729, 0.4026], [0.0352, 0.0730, 0.2506]]			)
  elif model_name == "videomae-base-short-finetuned-ssv2":
    __lowerCamelCase			:			Tuple       			=					torch.Size([1, 174]			)
    __lowerCamelCase			:			str       			=					torch.tensor([-0.0537, -0.1539, -0.3266]			)
  elif model_name == "videomae-base-ssv2":
    __lowerCamelCase			:			Tuple       			=					torch.Size([1, 1408, 1536]			)
    __lowerCamelCase			:			Any       			=					torch.tensor([[0.8131, 0.8727, 0.8546], [0.7366, 0.9377, 0.8870], [0.5935, 0.8874, 0.8564]]			)
  elif model_name == "videomae-base-finetuned-ssv2":
    __lowerCamelCase			:			Dict       			=					torch.Size([1, 174]			)
    __lowerCamelCase			:			Any       			=					torch.tensor([0.1961, -0.8337, -0.6389]			)
  else:
    raise ValueError(F'Model name not supported. Should be one of {model_names}'			)

  # verify logits
  assert logits.shape == expected_shape
  if "finetuned" in model_name:
    assert torch.allclose(logits[0, :3]			,_lowerCAmelCase			,atol=1E-4			)
  else:
    print('Logits:'			,logits[0, :3, :3]			)
    assert torch.allclose(logits[0, :3, :3]			,_lowerCAmelCase			,atol=1E-4			)
  print('Logits ok!'			)

  # verify loss, if applicable
  if model_name == "videomae-base-short":
    __lowerCamelCase			:			Optional[Any]       			=					outputs.loss
    assert torch.allclose(_lowerCAmelCase			,_lowerCAmelCase			,atol=1E-4			)
    print('Loss ok!'			)

  if pytorch_dump_folder_path is not None:
    print(F'Saving model and image processor to {pytorch_dump_folder_path}'			)
    image_processor.save_pretrained(_lowerCAmelCase			)
    model.save_pretrained(_lowerCAmelCase			)

  if push_to_hub:
    print('Pushing to the hub...'			)
    model.push_to_hub(_lowerCAmelCase			,organization='nielsr'			)


if __name__ == "__main__":
   _UpperCamelCase											=		argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--checkpoint_url',
       default='https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4',
       type=str,
       help=(
           'URL of the original PyTorch checkpoint (on Google Drive) you\'d like to convert. Should be a direct'
           ' download link.'
       ),
   )
   parser.add_argument(
       '--pytorch_dump_folder_path',
       default='/Users/nielsrogge/Documents/VideoMAE/Test',
       type=str,
       help='Path to the output PyTorch model directory.',
   )
   parser.add_argument('--model_name', default='videomae-base', type=str, help='Name of the model.')
   parser.add_argument(
       '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
   )

   _UpperCamelCase											=		parser.parse_args()
   convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)











'''simple docstring'''





import os
import tempfile
import unittest

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

from ...generation.test_utils import GenerationTesterMixin
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,
       NezhaForMaskedLM,
       NezhaForMultipleChoice,
       NezhaForNextSentencePrediction,
       NezhaForPreTraining,
       NezhaForQuestionAnswering,
       NezhaForSequenceClassification,
       NezhaForTokenClassification,
       NezhaModel,
   )
   from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST






class    lowerCamelCase_       :
  """simple docstring"""




  def __init__(					self :      str     ,      _a :      Dict     ,      _a :      List[str]=13     ,      _a :      List[str]=7     ,      _a :      Union[str, Any]=True     ,      _a :      List[Any]=True     ,      _a :      Optional[Any]=True     ,      _a :      Any=True     ,      _a :      Optional[Any]=99     ,      _a :      List[str]=32     ,      _a :      Optional[Any]=5     ,      _a :      str=4     ,      _a :      str=37     ,      _a :      List[Any]="gelu"     ,      _a :      List[Any]=0.1     ,      _a :      Optional[int]=0.1     ,      _a :      Optional[Any]=128     ,      _a :      Tuple=32     ,      _a :      List[Any]=16     ,      _a :      Optional[int]=2     ,      _a :      List[str]=0.02     ,      _a :      List[str]=3     ,      _a :      Any=4     ,      _a :      List[str]=None     ,      )   ->   Any:
    __lowerCamelCase			:			Optional[Any]       			=					parent
    __lowerCamelCase			:			Optional[int]       			=					batch_size
    __lowerCamelCase			:			List[str]       			=					seq_length
    __lowerCamelCase			:			List[str]       			=					is_training
    __lowerCamelCase			:			Dict       			=					use_input_mask
    __lowerCamelCase			:			Optional[int]       			=					use_token_type_ids
    __lowerCamelCase			:			Union[str, Any]       			=					use_labels
    __lowerCamelCase			:			Tuple       			=					vocab_size
    __lowerCamelCase			:			Any       			=					hidden_size
    __lowerCamelCase			:			Tuple       			=					num_hidden_layers
    __lowerCamelCase			:			List[str]       			=					num_attention_heads
    __lowerCamelCase			:			str       			=					intermediate_size
    __lowerCamelCase			:			Tuple       			=					hidden_act
    __lowerCamelCase			:			List[str]       			=					hidden_dropout_prob
    __lowerCamelCase			:			List[Any]       			=					attention_probs_dropout_prob
    __lowerCamelCase			:			List[str]       			=					max_position_embeddings
    __lowerCamelCase			:			str       			=					type_vocab_size
    __lowerCamelCase			:			Optional[Any]       			=					type_sequence_label_size
    __lowerCamelCase			:			Optional[Any]       			=					initializer_range
    __lowerCamelCase			:			Tuple       			=					num_labels
    __lowerCamelCase			:			Tuple       			=					num_choices
    __lowerCamelCase			:			Optional[int]       			=					scope




  def 		_lowercase	(					self :      Optional[int]  )   ->   Dict:
    __lowerCamelCase			:			Tuple       			=					ids_tensor([self.batch_size, self.seq_length]     ,      self.vocab_size  )

    __lowerCamelCase			:			Optional[int]       			=					None
    if self.use_input_mask:
      __lowerCamelCase			:			Optional[int]       			=					random_attention_mask([self.batch_size, self.seq_length]  )

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

    __lowerCamelCase			:			Optional[int]       			=					None
    __lowerCamelCase			:			Union[str, Any]       			=					None
    __lowerCamelCase			:			List[str]       			=					None
    if self.use_labels:
      __lowerCamelCase			:			Union[str, Any]       			=					ids_tensor([self.batch_size]     ,      self.type_sequence_label_size  )
      __lowerCamelCase			:			Optional[int]       			=					ids_tensor([self.batch_size, self.seq_length]     ,      self.num_labels  )
      __lowerCamelCase			:			Any       			=					ids_tensor([self.batch_size]     ,      self.num_choices  )

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

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




  def 		_lowercase	(					self :      List[str]  )   ->   int:
    return NezhaConfig(
        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=_a     ,      initializer_range=self.initializer_range     ,      )




  def 		_lowercase	(					self :      Tuple  )   ->   Optional[Any]:
    (
        (
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,
    )			:			List[Any]       			=					self.prepare_config_and_inputs()

    __lowerCamelCase			:			Optional[Any]       			=					True
    __lowerCamelCase			:			Optional[int]       			=					floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
    __lowerCamelCase			:			Optional[Any]       			=					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 		_lowercase	(					self :      Optional[int]     ,      _a :      List[Any]     ,      _a :      Dict     ,      _a :      Union[str, Any]     ,      _a :      Tuple     ,      _a :      Tuple     ,      _a :      Dict     ,      _a :      Any  )   ->   Tuple:
    __lowerCamelCase			:			List[Any]       			=					NezhaModel(config=_a  )
    model.to(_a  )
    model.eval()
    __lowerCamelCase			:			Any       			=					model(_a     ,      attention_mask=_a     ,      token_type_ids=_a  )
    __lowerCamelCase			:			int       			=					model(_a     ,      token_type_ids=_a  )
    __lowerCamelCase			:			Optional[Any]       			=					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 		_lowercase	(					self :      Optional[Any]     ,      _a :      List[str]     ,      _a :      Dict     ,      _a :      Optional[Any]     ,      _a :      int     ,      _a :      List[str]     ,      _a :      Optional[int]     ,      _a :      List[str]     ,      _a :      Optional[int]     ,      _a :      Any     ,      )   ->   Dict:
    __lowerCamelCase			:			Optional[Any]       			=					True
    __lowerCamelCase			:			Union[str, Any]       			=					NezhaModel(_a  )
    model.to(_a  )
    model.eval()
    __lowerCamelCase			:			Any       			=					model(
        _a     ,      attention_mask=_a     ,      token_type_ids=_a     ,      encoder_hidden_states=_a     ,      encoder_attention_mask=_a     ,      )
    __lowerCamelCase			:			Tuple       			=					model(
        _a     ,      attention_mask=_a     ,      token_type_ids=_a     ,      encoder_hidden_states=_a     ,      )
    __lowerCamelCase			:			str       			=					model(_a     ,      attention_mask=_a     ,      token_type_ids=_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 		_lowercase	(					self :      Union[str, Any]     ,      _a :      Optional[int]     ,      _a :      int     ,      _a :      Optional[Any]     ,      _a :      Any     ,      _a :      Tuple     ,      _a :      Optional[int]     ,      _a :      int  )   ->   List[Any]:
    __lowerCamelCase			:			int       			=					NezhaForMaskedLM(config=_a  )
    model.to(_a  )
    model.eval()
    __lowerCamelCase			:			Optional[Any]       			=					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 		_lowercase	(					self :      int     ,      _a :      Tuple     ,      _a :      List[Any]     ,      _a :      Any     ,      _a :      Optional[Any]     ,      _a :      Dict     ,      _a :      Dict     ,      _a :      List[Any]  )   ->   str:
    __lowerCamelCase			:			str       			=					NezhaForNextSentencePrediction(config=_a  )
    model.to(_a  )
    model.eval()
    __lowerCamelCase			:			str       			=					model(
        _a     ,      attention_mask=_a     ,      token_type_ids=_a     ,      labels=_a     ,      )
    self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, 2)  )




  def 		_lowercase	(					self :      Any     ,      _a :      List[str]     ,      _a :      str     ,      _a :      List[Any]     ,      _a :      str     ,      _a :      Union[str, Any]     ,      _a :      int     ,      _a :      Tuple  )   ->   Dict:
    __lowerCamelCase			:			List[str]       			=					NezhaForPreTraining(config=_a  )
    model.to(_a  )
    model.eval()
    __lowerCamelCase			:			Optional[int]       			=					model(
        _a     ,      attention_mask=_a     ,      token_type_ids=_a     ,      labels=_a     ,      next_sentence_label=_a     ,      )
    self.parent.assertEqual(result.prediction_logits.shape     ,      (self.batch_size, self.seq_length, self.vocab_size)  )
    self.parent.assertEqual(result.seq_relationship_logits.shape     ,      (self.batch_size, 2)  )




  def 		_lowercase	(					self :      int     ,      _a :      Dict     ,      _a :      Any     ,      _a :      Any     ,      _a :      Tuple     ,      _a :      List[str]     ,      _a :      Any     ,      _a :      List[Any]  )   ->   List[Any]:
    __lowerCamelCase			:			Any       			=					NezhaForQuestionAnswering(config=_a  )
    model.to(_a  )
    model.eval()
    __lowerCamelCase			:			List[Any]       			=					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 		_lowercase	(					self :      Optional[int]     ,      _a :      str     ,      _a :      Tuple     ,      _a :      List[str]     ,      _a :      List[str]     ,      _a :      Any     ,      _a :      str     ,      _a :      Union[str, Any]  )   ->   int:
    __lowerCamelCase			:			Optional[int]       			=					self.num_labels
    __lowerCamelCase			:			List[str]       			=					NezhaForSequenceClassification(_a  )
    model.to(_a  )
    model.eval()
    __lowerCamelCase			:			Any       			=					model(_a     ,      attention_mask=_a     ,      token_type_ids=_a     ,      labels=_a  )
    self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.num_labels)  )




  def 		_lowercase	(					self :      List[Any]     ,      _a :      List[Any]     ,      _a :      str     ,      _a :      Tuple     ,      _a :      Dict     ,      _a :      Dict     ,      _a :      Union[str, Any]     ,      _a :      List[Any]  )   ->   int:
    __lowerCamelCase			:			int       			=					self.num_labels
    __lowerCamelCase			:			int       			=					NezhaForTokenClassification(config=_a  )
    model.to(_a  )
    model.eval()
    __lowerCamelCase			:			List[str]       			=					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 		_lowercase	(					self :      List[str]     ,      _a :      Any     ,      _a :      List[str]     ,      _a :      Tuple     ,      _a :      Optional[Any]     ,      _a :      Optional[Any]     ,      _a :      str     ,      _a :      Optional[Any]  )   ->   int:
    __lowerCamelCase			:			List[Any]       			=					self.num_choices
    __lowerCamelCase			:			Optional[Any]       			=					NezhaForMultipleChoice(config=_a  )
    model.to(_a  )
    model.eval()
    __lowerCamelCase			:			Optional[Any]       			=					input_ids.unsqueeze(1  ).expand(-1     ,      self.num_choices     ,      -1  ).contiguous()
    __lowerCamelCase			:			Any       			=					token_type_ids.unsqueeze(1  ).expand(-1     ,      self.num_choices     ,      -1  ).contiguous()
    __lowerCamelCase			:			List[str]       			=					input_mask.unsqueeze(1  ).expand(-1     ,      self.num_choices     ,      -1  ).contiguous()
    __lowerCamelCase			:			str       			=					model(
        _a     ,      attention_mask=_a     ,      token_type_ids=_a     ,      labels=_a     ,      )
    self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.num_choices)  )




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






@require_torch
class    lowerCamelCase_       (							SCREAMING_SNAKE_CASE_  ,  SCREAMING_SNAKE_CASE_  ,  SCREAMING_SNAKE_CASE_  ,  unittest.TestCase	):
  """simple docstring"""




  a_				   =(
      (
          NezhaModel,
          NezhaForMaskedLM,
          NezhaForMultipleChoice,
          NezhaForNextSentencePrediction,
          NezhaForPreTraining,
          NezhaForQuestionAnswering,
          NezhaForSequenceClassification,
          NezhaForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  a_				   =(
      {
          """feature-extraction""": NezhaModel,
          """fill-mask""": NezhaForMaskedLM,
          """question-answering""": NezhaForQuestionAnswering,
          """text-classification""": NezhaForSequenceClassification,
          """token-classification""": NezhaForTokenClassification,
          """zero-shot""": NezhaForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  a_				   =True




  def 		_lowercase	(					self :      List[str]     ,      _a :      Union[str, Any]     ,      _a :      Any     ,      _a :      List[Any]=False  )   ->   Optional[Any]:
    __lowerCamelCase			:			List[str]       			=					super()._prepare_for_class(_a     ,      _a     ,      return_labels=_a  )

    if return_labels:
      if model_class in get_values(_a  ):
        __lowerCamelCase			:			Tuple       			=					torch.zeros(
            (self.model_tester.batch_size, self.model_tester.seq_length)     ,      dtype=torch.long     ,      device=_a  )
        __lowerCamelCase			:			List[Any]       			=					torch.zeros(
            self.model_tester.batch_size     ,      dtype=torch.long     ,      device=_a  )
    return inputs_dict




  def 		_lowercase	(					self :      Optional[int]  )   ->   Any:
    __lowerCamelCase			:			Dict       			=					NezhaModelTester(self  )
    __lowerCamelCase			:			int       			=					ConfigTester(self     ,      config_class=_a     ,      hidden_size=37  )




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




  def 		_lowercase	(					self :      int  )   ->   List[str]:
    __lowerCamelCase			:			str       			=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*_a  )




  def 		_lowercase	(					self :      Optional[Any]  )   ->   Union[str, Any]:
    __lowerCamelCase			:			List[str]       			=					self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_model_as_decoder(*_a  )




  def 		_lowercase	(					self :      str  )   ->   Any:
    # This regression test was failing with PyTorch < 1.3
    (
        (
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,(
        __lowerCamelCase
    )			,
    )			:			str       			=					self.model_tester.prepare_config_and_inputs_for_decoder()

    __lowerCamelCase			:			Dict       			=					None

    self.model_tester.create_and_check_model_as_decoder(
        _a     ,      _a     ,      _a     ,      _a     ,      _a     ,      _a     ,      _a     ,      _a     ,      _a     ,      )




  def 		_lowercase	(					self :      int  )   ->   Tuple:
    __lowerCamelCase			:			Any       			=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*_a  )




  def 		_lowercase	(					self :      str  )   ->   str:
    __lowerCamelCase			:			str       			=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_multiple_choice(*_a  )




  def 		_lowercase	(					self :      List[str]  )   ->   int:
    __lowerCamelCase			:			int       			=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_next_sequence_prediction(*_a  )




  def 		_lowercase	(					self :      Any  )   ->   Any:
    __lowerCamelCase			:			Union[str, Any]       			=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_pretraining(*_a  )




  def 		_lowercase	(					self :      Union[str, Any]  )   ->   Optional[int]:
    __lowerCamelCase			:			Dict       			=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_question_answering(*_a  )




  def 		_lowercase	(					self :      Optional[int]  )   ->   str:
    __lowerCamelCase			:			Optional[Any]       			=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_sequence_classification(*_a  )




  def 		_lowercase	(					self :      Dict  )   ->   Any:
    __lowerCamelCase			:			str       			=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*_a  )




  @slow
  def 		_lowercase	(					self :      Optional[Any]  )   ->   Union[str, Any]:
    for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      __lowerCamelCase			:			List[str]       			=					NezhaModel.from_pretrained(_a  )
      self.assertIsNotNone(_a  )




  @slow
  @require_torch_gpu
  def 		_lowercase	(					self :      List[Any]  )   ->   str:
    __lowerCamelCase			,__lowerCamelCase			:			Optional[int]       			=					self.model_tester.prepare_config_and_inputs_for_common()
    for model_class in self.all_model_classes:
      # NezhaForMultipleChoice behaves incorrectly in JIT environments.
      if model_class == NezhaForMultipleChoice:
        return

      __lowerCamelCase			:			List[str]       			=					True
      __lowerCamelCase			:			Dict       			=					model_class(config=_a  )

      __lowerCamelCase			:			List[str]       			=					self._prepare_for_class(_a     ,      _a  )
      __lowerCamelCase			:			Dict       			=					torch.jit.trace(
          _a     ,      (inputs_dict['input_ids'].to('cpu'  ), inputs_dict['attention_mask'].to('cpu'  ))  )

      with tempfile.TemporaryDirectory() as tmp:
        torch.jit.save(_a     ,      os.path.join(_a     ,      'bert.pt'  )  )
        __lowerCamelCase			:			int       			=					torch.jit.load(os.path.join(_a     ,      'bert.pt'  )     ,      map_location=_a  )
        loaded(inputs_dict['input_ids'].to(_a  )     ,      inputs_dict['attention_mask'].to(_a  )  )






@require_torch
class    lowerCamelCase_       (							unittest.TestCase	):
  """simple docstring"""




  @slow
  def 		_lowercase	(					self :      Dict  )   ->   Optional[int]:
    __lowerCamelCase			:			Dict       			=					NezhaModel.from_pretrained('sijunhe/nezha-cn-base'  )
    __lowerCamelCase			:			int       			=					torch.tensor([[0, 1, 2, 3, 4, 5]]  )
    __lowerCamelCase			:			Dict       			=					torch.tensor([[0, 1, 1, 1, 1, 1]]  )
    with torch.no_grad():
      __lowerCamelCase			:			Tuple       			=					model(_a     ,      attention_mask=_a  )[0]
    __lowerCamelCase			:			Optional[int]       			=					torch.Size((1, 6, 768)  )
    self.assertEqual(output.shape     ,      _a  )
    __lowerCamelCase			:			Union[str, Any]       			=					torch.tensor([[[0.0685, 0.2441, 0.1102], [0.0600, 0.1906, 0.1349], [0.0221, 0.0819, 0.0586]]]  )

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




  @slow
  def 		_lowercase	(					self :      Dict  )   ->   Dict:
    __lowerCamelCase			:			Dict       			=					NezhaForMaskedLM.from_pretrained('sijunhe/nezha-cn-base'  )
    __lowerCamelCase			:			Optional[Any]       			=					torch.tensor([[0, 1, 2, 3, 4, 5]]  )
    __lowerCamelCase			:			List[str]       			=					torch.tensor([[1, 1, 1, 1, 1, 1]]  )
    with torch.no_grad():
      __lowerCamelCase			:			Union[str, Any]       			=					model(_a     ,      attention_mask=_a  )[0]
    __lowerCamelCase			:			Optional[Any]       			=					torch.Size((1, 6, 2_1128)  )
    self.assertEqual(output.shape     ,      _a  )
    __lowerCamelCase			:			Optional[Any]       			=					torch.tensor(
        [[-2.7939, -1.7902, -2.2189], [-2.8585, -1.8908, -2.3723], [-2.6499, -1.7750, -2.2558]]  )

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






