






'''simple docstring'''
import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
a_					= logging.get_logger(__name__)





def  _a(  UpperCamelCase__				:						Any,	UpperCamelCase__				:						Dict=False							):




      '''simple docstring'''






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

      # projection layer + position embeddings
      rename_keys.extend(
          [
              ('''cls_token''', '''vit.embeddings.cls_token'''),
              ('''patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
              ('''patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
              ('''pos_embed''', '''vit.embeddings.position_embeddings'''),
          ]							)

      if base_model:
            # layernorm + pooler
            rename_keys.extend(
                [
                    ('''norm.weight''', '''layernorm.weight'''),
                    ('''norm.bias''', '''layernorm.bias'''),
                    ('''pre_logits.fc.weight''', '''pooler.dense.weight'''),
                    ('''pre_logits.fc.bias''', '''pooler.dense.bias'''),
                ]							)

            # if just the base model, we should remove "vit" from all keys that start with "vit"
            SCREAMING_SNAKE_CASE__     :				Tuple									=[(pair[0], pair[1][4:]) if pair[1].startswith('''vit'''							) else pair for pair in rename_keys]
      else:
            # layernorm + classification head
            rename_keys.extend(
                [
                    ('''norm.weight''', '''vit.layernorm.weight'''),
                    ('''norm.bias''', '''vit.layernorm.bias'''),
                    ('''head.weight''', '''classifier.weight'''),
                    ('''head.bias''', '''classifier.bias'''),
                ]							)

      return rename_keys





def  _a(  UpperCamelCase__				:						List[str],	UpperCamelCase__				:						List[Any],	UpperCamelCase__				:						Optional[int]=False							):




      '''simple docstring'''






      for i in range(config.num_hidden_layers							):
            if base_model:
                  SCREAMING_SNAKE_CASE__     :				List[Any]									=''''''
            else:
                  SCREAMING_SNAKE_CASE__     :				Optional[int]									='''vit.'''
            # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
            SCREAMING_SNAKE_CASE__     :				List[Any]									=state_dict.pop(f"blocks.{i}.attn.qkv.weight"							)
            SCREAMING_SNAKE_CASE__     :				Optional[int]									=state_dict.pop(f"blocks.{i}.attn.qkv.bias"							)
            # next, add query, keys and values (in that order) to the state dict
            SCREAMING_SNAKE_CASE__     :				int									=in_proj_weight[
                : config.hidden_size, :
            ]
            SCREAMING_SNAKE_CASE__     :				str									=in_proj_bias[: config.hidden_size]
            SCREAMING_SNAKE_CASE__     :				str									=in_proj_weight[
                config.hidden_size : config.hidden_size * 2, :
            ]
            SCREAMING_SNAKE_CASE__     :				str									=in_proj_bias[
                config.hidden_size : config.hidden_size * 2
            ]
            SCREAMING_SNAKE_CASE__     :				List[str]									=in_proj_weight[
                -config.hidden_size :, :
            ]
            SCREAMING_SNAKE_CASE__     :				Dict									=in_proj_bias[-config.hidden_size :]





def  _a(  UpperCamelCase__				:						int							):




      '''simple docstring'''






      SCREAMING_SNAKE_CASE__     :				int									=['''head.weight''', '''head.bias''']
      for k in ignore_keys:
            state_dict.pop(UpperCamelCase__,	UpperCamelCase__							)





def  _a(  UpperCamelCase__				:						Optional[Any],	UpperCamelCase__				:						str,	UpperCamelCase__				:						Optional[Any]							):




      '''simple docstring'''






      SCREAMING_SNAKE_CASE__     :				Optional[int]									=dct.pop(UpperCamelCase__							)
      SCREAMING_SNAKE_CASE__     :				Optional[int]									=val





def  _a(  ):




      '''simple docstring'''






      SCREAMING_SNAKE_CASE__     :				Union[str, Any]									='''http://images.cocodataset.org/val2017/000000039769.jpg'''
      SCREAMING_SNAKE_CASE__     :				List[Any]									=Image.open(requests.get(UpperCamelCase__,	stream=UpperCamelCase__							).raw							)
      return im





@torch.no_grad()
def  _a(  UpperCamelCase__				:						int,	UpperCamelCase__				:						List[str]							):




      '''simple docstring'''






      SCREAMING_SNAKE_CASE__     :				List[str]									=ViTConfig()
      SCREAMING_SNAKE_CASE__     :				int									=False
      # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
      if vit_name[-5:] == "in21k":
            SCREAMING_SNAKE_CASE__     :				Tuple									=True
            SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=int(vit_name[-1_2:-1_0]							)
            SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=int(vit_name[-9:-6]							)
      else:
            SCREAMING_SNAKE_CASE__     :				Any									=1_0_0_0
            SCREAMING_SNAKE_CASE__     :				List[Any]									='''huggingface/label-files'''
            SCREAMING_SNAKE_CASE__     :				int									='''imagenet-1k-id2label.json'''
            SCREAMING_SNAKE_CASE__     :				List[Any]									=json.load(open(hf_hub_download(UpperCamelCase__,	UpperCamelCase__,	repo_type='''dataset'''							),	'''r'''							)							)
            SCREAMING_SNAKE_CASE__     :				List[str]									={int(UpperCamelCase__							): v for k, v in idalabel.items()}
            SCREAMING_SNAKE_CASE__     :				int									=idalabel
            SCREAMING_SNAKE_CASE__     :				str									={v: k for k, v in idalabel.items()}
            SCREAMING_SNAKE_CASE__     :				List[Any]									=int(vit_name[-6:-4]							)
            SCREAMING_SNAKE_CASE__     :				int									=int(vit_name[-3:]							)
      # size of the architecture
      if "deit" in vit_name:
            if vit_name[9:].startswith('''tiny'''							):
                  SCREAMING_SNAKE_CASE__     :				Optional[int]									=1_9_2
                  SCREAMING_SNAKE_CASE__     :				Dict									=7_6_8
                  SCREAMING_SNAKE_CASE__     :				Optional[int]									=1_2
                  SCREAMING_SNAKE_CASE__     :				Tuple									=3
            elif vit_name[9:].startswith('''small'''							):
                  SCREAMING_SNAKE_CASE__     :				Any									=3_8_4
                  SCREAMING_SNAKE_CASE__     :				Any									=1_5_3_6
                  SCREAMING_SNAKE_CASE__     :				Any									=1_2
                  SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=6
            else:
                  pass
      else:
            if vit_name[4:].startswith('''small'''							):
                  SCREAMING_SNAKE_CASE__     :				str									=7_6_8
                  SCREAMING_SNAKE_CASE__     :				Optional[int]									=2_3_0_4
                  SCREAMING_SNAKE_CASE__     :				List[str]									=8
                  SCREAMING_SNAKE_CASE__     :				List[Any]									=8
            elif vit_name[4:].startswith('''base'''							):
                  pass
            elif vit_name[4:].startswith('''large'''							):
                  SCREAMING_SNAKE_CASE__     :				str									=1_0_2_4
                  SCREAMING_SNAKE_CASE__     :				Optional[Any]									=4_0_9_6
                  SCREAMING_SNAKE_CASE__     :				Optional[Any]									=2_4
                  SCREAMING_SNAKE_CASE__     :				Dict									=1_6
            elif vit_name[4:].startswith('''huge'''							):
                  SCREAMING_SNAKE_CASE__     :				List[Any]									=1_2_8_0
                  SCREAMING_SNAKE_CASE__     :				int									=5_1_2_0
                  SCREAMING_SNAKE_CASE__     :				Tuple									=3_2
                  SCREAMING_SNAKE_CASE__     :				int									=1_6

    # load original model from timm
      SCREAMING_SNAKE_CASE__     :				str									=timm.create_model(UpperCamelCase__,	pretrained=UpperCamelCase__							)
      timm_model.eval()

      # load state_dict of original model, remove and rename some keys
      SCREAMING_SNAKE_CASE__     :				Optional[int]									=timm_model.state_dict()
      if base_model:
            remove_classification_head_(UpperCamelCase__							)
      SCREAMING_SNAKE_CASE__     :				List[str]									=create_rename_keys(UpperCamelCase__,	UpperCamelCase__							)
      for src, dest in rename_keys:
            rename_key(UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__							)
      read_in_q_k_v(UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__							)

      # load HuggingFace model
      if vit_name[-5:] == "in21k":
            SCREAMING_SNAKE_CASE__     :				List[str]									=ViTModel(UpperCamelCase__							).eval()
      else:
            SCREAMING_SNAKE_CASE__     :				Dict									=ViTForImageClassification(UpperCamelCase__							).eval()
      model.load_state_dict(UpperCamelCase__							)

      # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
      if "deit" in vit_name:
            SCREAMING_SNAKE_CASE__     :				Optional[Any]									=DeiTImageProcessor(size=config.image_size							)
      else:
            SCREAMING_SNAKE_CASE__     :				Any									=ViTImageProcessor(size=config.image_size							)
      SCREAMING_SNAKE_CASE__     :				Optional[Any]									=image_processor(images=prepare_img(),	return_tensors='''pt'''							)
      SCREAMING_SNAKE_CASE__     :				Dict									=encoding['''pixel_values''']
      SCREAMING_SNAKE_CASE__     :				Any									=model(UpperCamelCase__							)

      if base_model:
            SCREAMING_SNAKE_CASE__     :				str									=timm_model.forward_features(UpperCamelCase__							)
            assert timm_pooled_output.shape == outputs.pooler_output.shape
            assert torch.allclose(UpperCamelCase__,	outputs.pooler_output,	atol=1e-3							)
      else:
            SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=timm_model(UpperCamelCase__							)
            assert timm_logits.shape == outputs.logits.shape
            assert torch.allclose(UpperCamelCase__,	outputs.logits,	atol=1e-3							)

      Path(UpperCamelCase__							).mkdir(exist_ok=UpperCamelCase__							)
      print(f"Saving model {vit_name} to {pytorch_dump_folder_path}"							)
      model.save_pretrained(UpperCamelCase__							)
      print(f"Saving image processor to {pytorch_dump_folder_path}"							)
      image_processor.save_pretrained(UpperCamelCase__							)


if __name__ == "__main__":
  a_					= argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--vit_name',
      default='vit_base_patch16_224',
      type=str,
      help='Name of the ViT timm model you\'d like to convert.',
  )
  parser.add_argument(
      '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
  )

  a_					= parser.parse_args()
  convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)







'''simple docstring'''
import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, 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,
      MegatronBertForCausalLM,
      MegatronBertForMaskedLM,
      MegatronBertForMultipleChoice,
      MegatronBertForNextSentencePrediction,
      MegatronBertForPreTraining,
      MegatronBertForQuestionAnswering,
      MegatronBertForSequenceClassification,
      MegatronBertForTokenClassification,
      MegatronBertModel,
  )






class      __SCREAMING_SNAKE_CASE      :




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




 def      __magic_name__	(      self	:							str      )	->						Optional[Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=ids_tensor([self.batch_size, self.seq_length]					,      self.vocab_size      )

       SCREAMING_SNAKE_CASE__     :				int									=None
       if self.use_input_mask:
             SCREAMING_SNAKE_CASE__     :				List[str]									=random_attention_mask([self.batch_size, self.seq_length]      )

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

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

       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.get_config()

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




 def      __magic_name__	(      self	:							List[str]      )	->						Any:
       return MegatronBertConfig(
           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					,      embedding_size=self.embedding_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=__lowercase					,      initializer_range=self.initializer_range					,      )




 def      __magic_name__	(      self	:							Any					,      __lowercase	:							Tuple					,      __lowercase	:							Any					,      __lowercase	:							List[Any]					,      __lowercase	:							Optional[int]					,      __lowercase	:							List[str]					,      __lowercase	:							Dict					,      __lowercase	:							Union[str, Any]      )	->						int:
       SCREAMING_SNAKE_CASE__     :				List[str]									=MegatronBertModel(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[Any]									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase      )
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=model(__lowercase					,      token_type_ids=__lowercase      )
       SCREAMING_SNAKE_CASE__     :				str									=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      __magic_name__	(      self	:							Dict					,      __lowercase	:							Optional[int]					,      __lowercase	:							List[str]					,      __lowercase	:							str					,      __lowercase	:							Tuple					,      __lowercase	:							Any					,      __lowercase	:							Optional[int]					,      __lowercase	:							Any      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				str									=MegatronBertForMaskedLM(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Tuple									=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      __magic_name__	(      self	:							List[str]					,      __lowercase	:							Tuple					,      __lowercase	:							Tuple					,      __lowercase	:							Optional[Any]					,      __lowercase	:							Any					,      __lowercase	:							str					,      __lowercase	:							Any					,      __lowercase	:							str      )	->						Tuple:
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=MegatronBertForCausalLM(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Dict									=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      __magic_name__	(      self	:							Union[str, Any]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Dict					,      __lowercase	:							Union[str, Any]					,      __lowercase	:							Dict					,      __lowercase	:							Union[str, Any]					,      __lowercase	:							int					,      __lowercase	:							Any      )	->						Dict:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=MegatronBertForNextSentencePrediction(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[str]									=model(
           __lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase					,      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, 2)      )




 def      __magic_name__	(      self	:							List[str]					,      __lowercase	:							Optional[int]					,      __lowercase	:							str					,      __lowercase	:							Optional[int]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Optional[Any]					,      __lowercase	:							str					,      __lowercase	:							Optional[Any]      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=MegatronBertForPreTraining(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[str]									=model(
           __lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase					,      next_sentence_label=__lowercase					,      )
       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      __magic_name__	(      self	:							Optional[int]					,      __lowercase	:							Tuple					,      __lowercase	:							Optional[int]					,      __lowercase	:							Union[str, Any]					,      __lowercase	:							int					,      __lowercase	:							Any					,      __lowercase	:							Any					,      __lowercase	:							List[str]      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				List[Any]									=MegatronBertForQuestionAnswering(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Dict									=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      __magic_name__	(      self	:							Optional[int]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Optional[int]					,      __lowercase	:							List[Any]					,      __lowercase	:							List[Any]					,      __lowercase	:							Optional[int]					,      __lowercase	:							int					,      __lowercase	:							List[str]      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=self.num_labels
       SCREAMING_SNAKE_CASE__     :				Dict									=MegatronBertForSequenceClassification(__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.num_labels)      )




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




 def      __magic_name__	(      self	:							str      )	->						Any:
       SCREAMING_SNAKE_CASE__     :				Tuple									=self.prepare_config_and_inputs()
       (
           (
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	
       )     :				str									=config_and_inputs
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									={'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict






@require_torch
class      __SCREAMING_SNAKE_CASE      (			lowerCamelCase	,				lowerCamelCase	,				unittest.TestCase				):
 snake_case_  							=							(
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 snake_case_  							=							(
     {
         """feature-extraction""": MegatronBertModel,
         """fill-mask""": MegatronBertForMaskedLM,
         """question-answering""": MegatronBertForQuestionAnswering,
         """text-classification""": MegatronBertForSequenceClassification,
         """text-generation""": MegatronBertForCausalLM,
         """token-classification""": MegatronBertForTokenClassification,
         """zero-shot""": MegatronBertForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 snake_case_  							=							True
 # test_resize_embeddings = False
 snake_case_  							=							False




 def      __magic_name__	(      self	:							List[Any]					,      __lowercase	:							List[str]					,      __lowercase	:							Tuple					,      __lowercase	:							Tuple=False      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=super()._prepare_for_class(__lowercase					,      __lowercase					,      return_labels=__lowercase      )

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




 def      __magic_name__	(      self	:							Tuple      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=MegatronBertModelTester(self      )
       SCREAMING_SNAKE_CASE__     :				Tuple									=ConfigTester(self					,      config_class=__lowercase					,      hidden_size=37      )




 def      __magic_name__	(      self	:							str      )	->						Dict:
       self.config_tester.run_common_tests()




 def      __magic_name__	(      self	:							Tuple      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				str									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_model(*__lowercase      )




 def      __magic_name__	(      self	:							List[str]      )	->						Any:
       SCREAMING_SNAKE_CASE__     :				int									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__lowercase      )




 def      __magic_name__	(      self	:							Optional[Any]      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				Dict									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__lowercase      )




 def      __magic_name__	(      self	:							Optional[int]      )	->						Dict:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__lowercase      )




 def      __magic_name__	(      self	:							Dict      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_pretraining(*__lowercase      )




 def      __magic_name__	(      self	:							int      )	->						Dict:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_question_answering(*__lowercase      )




 def      __magic_name__	(      self	:							List[Any]      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__lowercase      )




 def      __magic_name__	(      self	:							Union[str, Any]      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				List[Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_token_classification(*__lowercase      )





def  _a(  UpperCamelCase__				:						List[str]							):




      '''simple docstring'''






      return torch.tensor(
          UpperCamelCase__,	dtype=torch.long,	device=UpperCamelCase__,	)


a_					= 1E-4







@require_torch
@require_sentencepiece
@require_tokenizers
class      __SCREAMING_SNAKE_CASE      (			unittest.TestCase				):




 @slow
 @unittest.skip('''Model is not available.'''      )
 def      __magic_name__	(      self	:							Any      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				Any									='''nvidia/megatron-bert-uncased-345m'''
       if "MYDIR" in os.environ:
             SCREAMING_SNAKE_CASE__     :				Optional[int]									=os.path.join(os.environ['''MYDIR''']					,      __lowercase      )
       SCREAMING_SNAKE_CASE__     :				Any									=MegatronBertModel.from_pretrained(__lowercase      )
       model.to(__lowercase      )
       model.half()
       SCREAMING_SNAKE_CASE__     :				Dict									=_long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]      )
       with torch.no_grad():
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=model(__lowercase      )[0]
       SCREAMING_SNAKE_CASE__     :				Dict									=torch.Size((1, 9, 10_24)      )
       self.assertEqual(output.shape					,      __lowercase      )

       SCREAMING_SNAKE_CASE__     :				str									=[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
       for ii in range(3      ):
             for jj in range(3      ):
                   SCREAMING_SNAKE_CASE__     :				List[Any]									=output[0, ii, jj]
                   SCREAMING_SNAKE_CASE__     :				Tuple									=expected[3 * ii + jj]
                   SCREAMING_SNAKE_CASE__     :				List[str]									='''ii={} jj={} a={} b={}'''.format(__lowercase					,      __lowercase					,      __lowercase					,      __lowercase      )
                   self.assertTrue(math.isclose(__lowercase					,      __lowercase					,      rel_tol=__lowercase					,      abs_tol=__lowercase      )					,      msg=__lowercase      )