


import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_url
from PIL import Image

from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
A   :					int	=      logging.get_logger(__name__)

def    lowercase_	(	_A    :					Union[str, Any]		):

      """simple docstring"""




      lowerCamelCase__	:       Any					=						DPTConfig(embedding_type="hybrid"		)

      if "large" in checkpoint_url:
            lowerCamelCase__	:       List[str]					=						1024
            lowerCamelCase__	:       Tuple					=						4096
            lowerCamelCase__	:       List[str]					=						24
            lowerCamelCase__	:       int					=						16
            lowerCamelCase__	:       Optional[int]					=						[5, 11, 17, 23]
            lowerCamelCase__	:       Dict					=						[256, 512, 1024, 1024]
            lowerCamelCase__	:       int					=						(1, 384, 384)

      if "nyu" or "midas" in checkpoint_url:
            lowerCamelCase__	:       Dict					=						768
            lowerCamelCase__	:       Tuple					=						[1, 1, 1, 0.5]
            lowerCamelCase__	:       Optional[Any]					=						[256, 512, 768, 768]
            lowerCamelCase__	:       List[Any]					=						150
            lowerCamelCase__	:       Tuple					=						16
            lowerCamelCase__	:       Any					=						(1, 384, 384)
            lowerCamelCase__	:       Tuple					=						False
            lowerCamelCase__	:       Optional[Any]					=						"project"

      if "ade" in checkpoint_url:
            lowerCamelCase__	:       Tuple					=						True
            lowerCamelCase__	:       Any					=						768
            lowerCamelCase__	:       Union[str, Any]					=						[1, 1, 1, 0.5]
            lowerCamelCase__	:       Optional[Any]					=						150
            lowerCamelCase__	:       List[Any]					=						16
            lowerCamelCase__	:       List[Any]					=						"huggingface/label-files"
            lowerCamelCase__	:       Tuple					=						"ade20k-id2label.json"
            lowerCamelCase__	:       Any					=						json.load(open(cached_download(hf_hub_url(_A					,      _A					,      repo_type="dataset"		)		)					,      "r"		)		)
            lowerCamelCase__	:       List[Any]					=						{int(_A		): v for k, v in idalabel.items()}
            lowerCamelCase__	:       Optional[int]					=						idalabel
            lowerCamelCase__	:       List[str]					=						{v: k for k, v in idalabel.items()}
            lowerCamelCase__	:       Optional[int]					=						[1, 150, 480, 480]

      return config, expected_shape

def    lowercase_	(	_A    :					List[str]		):

      """simple docstring"""




      lowerCamelCase__	:       List[Any]					=						["pretrained.model.head.weight", "pretrained.model.head.bias"]
      for k in ignore_keys:
            state_dict.pop(_A					,      _A		)

def    lowercase_	(	_A    :					Optional[Any]		):

      """simple docstring"""




      if (
          "pretrained.model" in name
          and "cls_token" not in name
          and "pos_embed" not in name
          and "patch_embed" not in name
      ):
            lowerCamelCase__	:       Any					=						name.replace("pretrained.model"					,      "dpt.encoder"		)
      if "pretrained.model" in name:
            lowerCamelCase__	:       int					=						name.replace("pretrained.model"					,      "dpt.embeddings"		)
      if "patch_embed" in name:
            lowerCamelCase__	:       Dict					=						name.replace("patch_embed"					,      ""		)
      if "pos_embed" in name:
            lowerCamelCase__	:       str					=						name.replace("pos_embed"					,      "position_embeddings"		)
      if "attn.proj" in name:
            lowerCamelCase__	:       Union[str, Any]					=						name.replace("attn.proj"					,      "attention.output.dense"		)
      if "proj" in name and "project" not in name:
            lowerCamelCase__	:       str					=						name.replace("proj"					,      "projection"		)
      if "blocks" in name:
            lowerCamelCase__	:       List[Any]					=						name.replace("blocks"					,      "layer"		)
      if "mlp.fc1" in name:
            lowerCamelCase__	:       Union[str, Any]					=						name.replace("mlp.fc1"					,      "intermediate.dense"		)
      if "mlp.fc2" in name:
            lowerCamelCase__	:       Optional[int]					=						name.replace("mlp.fc2"					,      "output.dense"		)
      if "norm1" in name and "backbone" not in name:
            lowerCamelCase__	:       Optional[int]					=						name.replace("norm1"					,      "layernorm_before"		)
      if "norm2" in name and "backbone" not in name:
            lowerCamelCase__	:       Optional[int]					=						name.replace("norm2"					,      "layernorm_after"		)
      if "scratch.output_conv" in name:
            lowerCamelCase__	:       Tuple					=						name.replace("scratch.output_conv"					,      "head"		)
      if "scratch" in name:
            lowerCamelCase__	:       List[str]					=						name.replace("scratch"					,      "neck"		)
      if "layer1_rn" in name:
            lowerCamelCase__	:       Union[str, Any]					=						name.replace("layer1_rn"					,      "convs.0"		)
      if "layer2_rn" in name:
            lowerCamelCase__	:       Tuple					=						name.replace("layer2_rn"					,      "convs.1"		)
      if "layer3_rn" in name:
            lowerCamelCase__	:       Any					=						name.replace("layer3_rn"					,      "convs.2"		)
      if "layer4_rn" in name:
            lowerCamelCase__	:       Any					=						name.replace("layer4_rn"					,      "convs.3"		)
      if "refinenet" in name:
            lowerCamelCase__	:       str					=						int(name[len("neck.refinenet"		) : len("neck.refinenet"		) + 1]		)
            # tricky here: we need to map 4 to 0, 3 to 1, 2 to 2 and 1 to 3
            lowerCamelCase__	:       List[Any]					=						name.replace(F"refinenet{layer_idx}"					,      F"fusion_stage.layers.{abs(layer_idx-4		)}"		)
      if "out_conv" in name:
            lowerCamelCase__	:       Dict					=						name.replace("out_conv"					,      "projection"		)
      if "resConfUnit1" in name:
            lowerCamelCase__	:       Optional[Any]					=						name.replace("resConfUnit1"					,      "residual_layer1"		)
      if "resConfUnit2" in name:
            lowerCamelCase__	:       List[str]					=						name.replace("resConfUnit2"					,      "residual_layer2"		)
      if "conv1" in name:
            lowerCamelCase__	:       Union[str, Any]					=						name.replace("conv1"					,      "convolution1"		)
      if "conv2" in name:
            lowerCamelCase__	:       Optional[Any]					=						name.replace("conv2"					,      "convolution2"		)
      # readout blocks
      if "pretrained.act_postprocess1.0.project.0" in name:
            lowerCamelCase__	:       List[Any]					=						name.replace("pretrained.act_postprocess1.0.project.0"					,      "neck.reassemble_stage.readout_projects.0.0"		)
      if "pretrained.act_postprocess2.0.project.0" in name:
            lowerCamelCase__	:       List[str]					=						name.replace("pretrained.act_postprocess2.0.project.0"					,      "neck.reassemble_stage.readout_projects.1.0"		)
      if "pretrained.act_postprocess3.0.project.0" in name:
            lowerCamelCase__	:       Any					=						name.replace("pretrained.act_postprocess3.0.project.0"					,      "neck.reassemble_stage.readout_projects.2.0"		)
      if "pretrained.act_postprocess4.0.project.0" in name:
            lowerCamelCase__	:       Optional[Any]					=						name.replace("pretrained.act_postprocess4.0.project.0"					,      "neck.reassemble_stage.readout_projects.3.0"		)

      # resize blocks
      if "pretrained.act_postprocess1.3" in name:
            lowerCamelCase__	:       str					=						name.replace("pretrained.act_postprocess1.3"					,      "neck.reassemble_stage.layers.0.projection"		)
      if "pretrained.act_postprocess1.4" in name:
            lowerCamelCase__	:       Any					=						name.replace("pretrained.act_postprocess1.4"					,      "neck.reassemble_stage.layers.0.resize"		)
      if "pretrained.act_postprocess2.3" in name:
            lowerCamelCase__	:       int					=						name.replace("pretrained.act_postprocess2.3"					,      "neck.reassemble_stage.layers.1.projection"		)
      if "pretrained.act_postprocess2.4" in name:
            lowerCamelCase__	:       Tuple					=						name.replace("pretrained.act_postprocess2.4"					,      "neck.reassemble_stage.layers.1.resize"		)
      if "pretrained.act_postprocess3.3" in name:
            lowerCamelCase__	:       Optional[Any]					=						name.replace("pretrained.act_postprocess3.3"					,      "neck.reassemble_stage.layers.2.projection"		)
      if "pretrained.act_postprocess4.3" in name:
            lowerCamelCase__	:       str					=						name.replace("pretrained.act_postprocess4.3"					,      "neck.reassemble_stage.layers.3.projection"		)
      if "pretrained.act_postprocess4.4" in name:
            lowerCamelCase__	:       Any					=						name.replace("pretrained.act_postprocess4.4"					,      "neck.reassemble_stage.layers.3.resize"		)
      if "pretrained" in name:
            lowerCamelCase__	:       int					=						name.replace("pretrained"					,      "dpt"		)
      if "bn" in name:
            lowerCamelCase__	:       List[str]					=						name.replace("bn"					,      "batch_norm"		)
      if "head" in name:
            lowerCamelCase__	:       Tuple					=						name.replace("head"					,      "head.head"		)
      if "encoder.norm" in name:
            lowerCamelCase__	:       Union[str, Any]					=						name.replace("encoder.norm"					,      "layernorm"		)
      if "auxlayer" in name:
            lowerCamelCase__	:       List[Any]					=						name.replace("auxlayer"					,      "auxiliary_head.head"		)
      if "backbone" in name:
            lowerCamelCase__	:       int					=						name.replace("backbone"					,      "backbone.bit.encoder"		)

      if ".." in name:
            lowerCamelCase__	:       int					=						name.replace(".."					,      "."		)

      if "stem.conv" in name:
            lowerCamelCase__	:       Any					=						name.replace("stem.conv"					,      "bit.embedder.convolution"		)
      if "blocks" in name:
            lowerCamelCase__	:       Any					=						name.replace("blocks"					,      "layers"		)
      if "convolution" in name and "backbone" in name:
            lowerCamelCase__	:       Dict					=						name.replace("convolution"					,      "conv"		)
      if "layer" in name and "backbone" in name:
            lowerCamelCase__	:       List[str]					=						name.replace("layer"					,      "layers"		)
      if "backbone.bit.encoder.bit" in name:
            lowerCamelCase__	:       List[str]					=						name.replace("backbone.bit.encoder.bit"					,      "backbone.bit"		)
      if "embedder.conv" in name:
            lowerCamelCase__	:       Optional[Any]					=						name.replace("embedder.conv"					,      "embedder.convolution"		)
      if "backbone.bit.encoder.stem.norm" in name:
            lowerCamelCase__	:       str					=						name.replace("backbone.bit.encoder.stem.norm"					,      "backbone.bit.embedder.norm"		)
      return name

def    lowercase_	(	_A    :					List[Any]					,      _A    :					Dict		):

      """simple docstring"""




      for i in range(config.num_hidden_layers		):
            # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
            lowerCamelCase__	:       str					=						state_dict.pop(F"dpt.encoder.layer.{i}.attn.qkv.weight"		)
            lowerCamelCase__	:       Any					=						state_dict.pop(F"dpt.encoder.layer.{i}.attn.qkv.bias"		)
            # next, add query, keys and values (in that order) to the state dict
            lowerCamelCase__	:       Union[str, Any]					=						in_proj_weight[: config.hidden_size, :]
            lowerCamelCase__	:       Tuple					=						in_proj_bias[: config.hidden_size]
            lowerCamelCase__	:       Dict					=						in_proj_weight[
                config.hidden_size : config.hidden_size * 2, :
            ]
            lowerCamelCase__	:       str					=						in_proj_bias[
                config.hidden_size : config.hidden_size * 2
            ]
            lowerCamelCase__	:       Dict					=						in_proj_weight[
                -config.hidden_size :, :
            ]
            lowerCamelCase__	:       str					=						in_proj_bias[-config.hidden_size :]

def    lowercase_	(	):

      """simple docstring"""




      lowerCamelCase__	:       int					=						"http://images.cocodataset.org/val2017/000000039769.jpg"
      lowerCamelCase__	:       str					=						Image.open(requests.get(_A					,      stream=_A		).raw		)
      return im

@torch.no_grad()
def    lowercase_	(	_A    :					List[str]					,      _A    :					str					,      _A    :					List[str]					,      _A    :					Optional[Any]					,      _A    :					List[str]		):

      """simple docstring"""




      lowerCamelCase__					,     lowerCamelCase__	:       Tuple					=						get_dpt_config(_A		)
      # load original state_dict from URL
      # state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")
      lowerCamelCase__	:       Dict					=						torch.load(_A					,      map_location="cpu"		)
      # remove certain keys
      remove_ignore_keys_(_A		)
      # rename keys
      for key in state_dict.copy().keys():
            lowerCamelCase__	:       Optional[Any]					=						state_dict.pop(_A		)
            lowerCamelCase__	:       Optional[Any]					=						val
      # read in qkv matrices
      read_in_q_k_v(_A					,      _A		)

      # load HuggingFace model
      lowerCamelCase__	:       Optional[int]					=						DPTForSemanticSegmentation(_A		) if "ade" in checkpoint_url else DPTForDepthEstimation(_A		)
      model.load_state_dict(_A		)
      model.eval()

      # Check outputs on an image
      lowerCamelCase__	:       List[Any]					=						480 if "ade" in checkpoint_url else 384
      lowerCamelCase__	:       int					=						DPTImageProcessor(size=_A		)

      lowerCamelCase__	:       Optional[Any]					=						prepare_img()
      lowerCamelCase__	:       List[Any]					=						image_processor(_A					,      return_tensors="pt"		)

      # forward pass
      lowerCamelCase__	:       str					=						model(**_A		).logits if "ade" in checkpoint_url else model(**_A		).predicted_depth

      if show_prediction:
            lowerCamelCase__	:       Optional[int]					=						(
                torch.nn.functional.interpolate(
                    outputs.unsqueeze(1		)					,      size=(image.size[1], image.size[0])					,      mode="bicubic"					,      align_corners=_A					,      )
                .squeeze()
                .cpu()
                .numpy()
            )

            Image.fromarray((prediction / prediction.max()) * 255		).show()

      if pytorch_dump_folder_path is not None:
            Path(_A		).mkdir(exist_ok=_A		)
            print(F"Saving model to {pytorch_dump_folder_path}"		)
            model.save_pretrained(_A		)
            print(F"Saving image processor to {pytorch_dump_folder_path}"		)
            image_processor.save_pretrained(_A		)

      if push_to_hub:
            model.push_to_hub("ybelkada/dpt-hybrid-midas"		)
            image_processor.push_to_hub("ybelkada/dpt-hybrid-midas"		)


if __name__ == "__main__":
    A   :					Optional[Any]	=      argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--checkpoint_url",
        default="https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt",
        type=str,
        help="URL of the original DPT checkpoint you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        type=str,
        required=False,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument(
        "--push_to_hub",
        action="store_true",
    )
    parser.add_argument(
        "--model_name",
        default="dpt-large",
        type=str,
        help="Name of the model, in case you're pushing to the hub.",
    )
    parser.add_argument(
        "--show_prediction",
        action="store_true",
    )

    A   :					List[Any]	=      parser.parse_args()
    convert_dpt_checkpoint(
        args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name, args.show_prediction
    )






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     _lowercase				:


       """simple docstring"""




       def __init__(     self							:    List[str]      ,   __lowerCamelCase							:    List[str]      ,   __lowerCamelCase							:    List[str]=13      ,   __lowerCamelCase							:    Dict=7      ,   __lowerCamelCase							:    List[Any]=True      ,   __lowerCamelCase							:    List[Any]=True      ,   __lowerCamelCase							:    Union[str, Any]=True      ,   __lowerCamelCase							:    Optional[Any]=True      ,   __lowerCamelCase							:    Optional[int]=99      ,   __lowerCamelCase							:    List[Any]=16      ,   __lowerCamelCase							:    Optional[Any]=36      ,   __lowerCamelCase							:    Optional[int]=6      ,   __lowerCamelCase							:    Union[str, Any]=6      ,   __lowerCamelCase							:    Optional[int]=6      ,   __lowerCamelCase							:    Dict=37      ,   __lowerCamelCase							:    List[Any]="gelu"      ,   __lowerCamelCase							:    Tuple=0.1      ,   __lowerCamelCase							:    Optional[Any]=0.1      ,   __lowerCamelCase							:    List[Any]=512      ,   __lowerCamelCase							:    Dict=16      ,   __lowerCamelCase							:    Union[str, Any]=2      ,   __lowerCamelCase							:    Dict=0.0_2      ,   __lowerCamelCase							:    Optional[Any]=3      ,   __lowerCamelCase							:    Dict=4      ,   __lowerCamelCase							:    Dict=None      ,   ):


             '''simple docstring'''



             lowerCamelCase__	:       Dict					=						parent
             lowerCamelCase__	:       List[Any]					=						batch_size
             lowerCamelCase__	:       Any					=						seq_length
             lowerCamelCase__	:       List[str]					=						is_training
             lowerCamelCase__	:       int					=						use_input_mask
             lowerCamelCase__	:       List[str]					=						use_token_type_ids
             lowerCamelCase__	:       int					=						use_labels
             lowerCamelCase__	:       Dict					=						vocab_size
             lowerCamelCase__	:       List[Any]					=						embedding_size
             lowerCamelCase__	:       Dict					=						hidden_size
             lowerCamelCase__	:       Any					=						num_hidden_layers
             lowerCamelCase__	:       Optional[Any]					=						num_hidden_groups
             lowerCamelCase__	:       Optional[int]					=						num_attention_heads
             lowerCamelCase__	:       List[str]					=						intermediate_size
             lowerCamelCase__	:       Optional[Any]					=						hidden_act
             lowerCamelCase__	:       str					=						hidden_dropout_prob
             lowerCamelCase__	:       Union[str, Any]					=						attention_probs_dropout_prob
             lowerCamelCase__	:       Optional[int]					=						max_position_embeddings
             lowerCamelCase__	:       List[Any]					=						type_vocab_size
             lowerCamelCase__	:       Optional[Any]					=						type_sequence_label_size
             lowerCamelCase__	:       Optional[int]					=						initializer_range
             lowerCamelCase__	:       str					=						num_labels
             lowerCamelCase__	:       List[Any]					=						num_choices
             lowerCamelCase__	:       Any					=						scope





       def 		lowerCAmelCase  (     self							:    Dict ):


             '''simple docstring'''



             lowerCamelCase__	:       Optional[Any]					=						ids_tensor([self.batch_size, self.seq_length]      ,   self.vocab_size )

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

             lowerCamelCase__	:       Optional[Any]					=						None
             if self.use_token_type_ids:
                   lowerCamelCase__	:       Union[str, Any]					=						ids_tensor([self.batch_size, self.seq_length]      ,   self.type_vocab_size )

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

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

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





       def 		lowerCAmelCase  (     self							:    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 		lowerCAmelCase  (     self							:    Dict      ,   __lowerCamelCase							:    Optional[Any]      ,   __lowerCamelCase							:    Dict      ,   __lowerCamelCase							:    int      ,   __lowerCamelCase							:    List[str]      ,   __lowerCamelCase							:    List[str]      ,   __lowerCamelCase							:    Any      ,   __lowerCamelCase							:    List[Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       int					=						AlbertModel(config=__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()
             lowerCamelCase__	:       Any					=						model(__lowerCamelCase      ,   attention_mask=__lowerCamelCase      ,   token_type_ids=__lowerCamelCase )
             lowerCamelCase__	:       Any					=						model(__lowerCamelCase      ,   token_type_ids=__lowerCamelCase )
             lowerCamelCase__	:       Optional[int]					=						model(__lowerCamelCase )
             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 		lowerCAmelCase  (     self							:    Union[str, Any]      ,   __lowerCamelCase							:    Union[str, Any]      ,   __lowerCamelCase							:    Optional[Any]      ,   __lowerCamelCase							:    Dict      ,   __lowerCamelCase							:    str      ,   __lowerCamelCase							:    Union[str, Any]      ,   __lowerCamelCase							:    int      ,   __lowerCamelCase							:    Tuple ):


             '''simple docstring'''



             lowerCamelCase__	:       Any					=						AlbertForPreTraining(config=__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()
             lowerCamelCase__	:       Union[str, Any]					=						model(
                 __lowerCamelCase      ,   attention_mask=__lowerCamelCase      ,   token_type_ids=__lowerCamelCase      ,   labels=__lowerCamelCase      ,   sentence_order_label=__lowerCamelCase      ,   )
             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 		lowerCAmelCase  (     self							:    str      ,   __lowerCamelCase							:    str      ,   __lowerCamelCase							:    List[Any]      ,   __lowerCamelCase							:    Any      ,   __lowerCamelCase							:    str      ,   __lowerCamelCase							:    List[Any]      ,   __lowerCamelCase							:    Optional[int]      ,   __lowerCamelCase							:    Tuple ):


             '''simple docstring'''



             lowerCamelCase__	:       Dict					=						AlbertForMaskedLM(config=__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()
             lowerCamelCase__	:       Tuple					=						model(__lowerCamelCase      ,   attention_mask=__lowerCamelCase      ,   token_type_ids=__lowerCamelCase      ,   labels=__lowerCamelCase )
             self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.seq_length, self.vocab_size) )





       def 		lowerCAmelCase  (     self							:    Union[str, Any]      ,   __lowerCamelCase							:    Optional[int]      ,   __lowerCamelCase							:    List[Any]      ,   __lowerCamelCase							:    Optional[Any]      ,   __lowerCamelCase							:    Optional[Any]      ,   __lowerCamelCase							:    Optional[Any]      ,   __lowerCamelCase							:    Any      ,   __lowerCamelCase							:    int ):


             '''simple docstring'''



             lowerCamelCase__	:       str					=						AlbertForQuestionAnswering(config=__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()
             lowerCamelCase__	:       str					=						model(
                 __lowerCamelCase      ,   attention_mask=__lowerCamelCase      ,   token_type_ids=__lowerCamelCase      ,   start_positions=__lowerCamelCase      ,   end_positions=__lowerCamelCase      ,   )
             self.parent.assertEqual(result.start_logits.shape      ,   (self.batch_size, self.seq_length) )
             self.parent.assertEqual(result.end_logits.shape      ,   (self.batch_size, self.seq_length) )





       def 		lowerCAmelCase  (     self							:    Optional[int]      ,   __lowerCamelCase							:    Tuple      ,   __lowerCamelCase							:    Optional[Any]      ,   __lowerCamelCase							:    Optional[Any]      ,   __lowerCamelCase							:    Optional[int]      ,   __lowerCamelCase							:    Optional[int]      ,   __lowerCamelCase							:    Union[str, Any]      ,   __lowerCamelCase							:    Dict ):


             '''simple docstring'''



             lowerCamelCase__	:       int					=						self.num_labels
             lowerCamelCase__	:       Optional[int]					=						AlbertForSequenceClassification(__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()
             lowerCamelCase__	:       Any					=						model(__lowerCamelCase      ,   attention_mask=__lowerCamelCase      ,   token_type_ids=__lowerCamelCase      ,   labels=__lowerCamelCase )
             self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.num_labels) )





       def 		lowerCAmelCase  (     self							:    Dict      ,   __lowerCamelCase							:    Dict      ,   __lowerCamelCase							:    List[str]      ,   __lowerCamelCase							:    List[str]      ,   __lowerCamelCase							:    Any      ,   __lowerCamelCase							:    Dict      ,   __lowerCamelCase							:    Any      ,   __lowerCamelCase							:    List[Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       Optional[int]					=						self.num_labels
             lowerCamelCase__	:       List[str]					=						AlbertForTokenClassification(config=__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()
             lowerCamelCase__	:       Tuple					=						model(__lowerCamelCase      ,   attention_mask=__lowerCamelCase      ,   token_type_ids=__lowerCamelCase      ,   labels=__lowerCamelCase )
             self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.seq_length, self.num_labels) )





       def 		lowerCAmelCase  (     self							:    Tuple      ,   __lowerCamelCase							:    Any      ,   __lowerCamelCase							:    Optional[Any]      ,   __lowerCamelCase							:    List[Any]      ,   __lowerCamelCase							:    List[Any]      ,   __lowerCamelCase							:    Any      ,   __lowerCamelCase							:    Any      ,   __lowerCamelCase							:    Union[str, Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       Union[str, Any]					=						self.num_choices
             lowerCamelCase__	:       Optional[int]					=						AlbertForMultipleChoice(config=__lowerCamelCase )
             model.to(__lowerCamelCase )
             model.eval()
             lowerCamelCase__	:       Any					=						input_ids.unsqueeze(1 ).expand(-1      ,   self.num_choices      ,   -1 ).contiguous()
             lowerCamelCase__	:       Optional[int]					=						token_type_ids.unsqueeze(1 ).expand(-1      ,   self.num_choices      ,   -1 ).contiguous()
             lowerCamelCase__	:       Dict					=						input_mask.unsqueeze(1 ).expand(-1      ,   self.num_choices      ,   -1 ).contiguous()
             lowerCamelCase__	:       int					=						model(
                 __lowerCamelCase      ,   attention_mask=__lowerCamelCase      ,   token_type_ids=__lowerCamelCase      ,   labels=__lowerCamelCase      ,   )
             self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.num_choices) )





       def 		lowerCAmelCase  (     self							:    str ):


             '''simple docstring'''



             lowerCamelCase__	:       int					=						self.prepare_config_and_inputs()
             (
                 (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     (
                 lowerCamelCase__
             )					,     
             )	:       Union[str, Any]					=						config_and_inputs
             lowerCamelCase__	:       str					=						{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
             return config, inputs_dict







@require_torch
class     _lowercase				(							lowercase__							,						lowercase__							,						unittest.TestCase):


       """simple docstring"""




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





       def 		lowerCAmelCase  (     self							:    List[str]      ,   __lowerCamelCase							:    int      ,   __lowerCamelCase							:    Tuple      ,   __lowerCamelCase							:    Dict=False ):


             '''simple docstring'''



             lowerCamelCase__	:       Any					=						super()._prepare_for_class(__lowerCamelCase      ,   __lowerCamelCase      ,   return_labels=__lowerCamelCase )

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





       def 		lowerCAmelCase  (     self							:    Optional[Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       Optional[int]					=						AlbertModelTester(self )
             lowerCamelCase__	:       Optional[Any]					=						ConfigTester(self      ,   config_class=__lowerCamelCase      ,   hidden_size=37 )





       def 		lowerCAmelCase  (     self							:    Tuple ):


             '''simple docstring'''



             self.config_tester.run_common_tests()





       def 		lowerCAmelCase  (     self							:    Dict ):


             '''simple docstring'''



             lowerCamelCase__	:       Optional[Any]					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*__lowerCamelCase )





       def 		lowerCAmelCase  (     self							:    Union[str, Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       Union[str, Any]					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_pretraining(*__lowerCamelCase )





       def 		lowerCAmelCase  (     self							:    Any ):


             '''simple docstring'''



             lowerCamelCase__	:       int					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase )





       def 		lowerCAmelCase  (     self							:    Any ):


             '''simple docstring'''



             lowerCamelCase__	:       List[Any]					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_multiple_choice(*__lowerCamelCase )





       def 		lowerCAmelCase  (     self							:    Tuple ):


             '''simple docstring'''



             lowerCamelCase__	:       Optional[Any]					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_question_answering(*__lowerCamelCase )





       def 		lowerCAmelCase  (     self							:    Union[str, Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       List[str]					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_sequence_classification(*__lowerCamelCase )





       def 		lowerCAmelCase  (     self							:    Optional[Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       List[str]					=						self.model_tester.prepare_config_and_inputs()
             for type in ["absolute", "relative_key", "relative_key_query"]:
                   lowerCamelCase__	:       Dict					=						type
                   self.model_tester.create_and_check_model(*__lowerCamelCase )





       @slow
       def 		lowerCAmelCase  (     self							:    Optional[Any] ):


             '''simple docstring'''



             for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   lowerCamelCase__	:       List[str]					=						AlbertModel.from_pretrained(__lowerCamelCase )
                   self.assertIsNotNone(__lowerCamelCase )







@require_torch
class     _lowercase				(							unittest.TestCase):


       """simple docstring"""




       @slow
       def 		lowerCAmelCase  (     self							:    Union[str, Any] ):


             '''simple docstring'''



             lowerCamelCase__	:       List[Any]					=						AlbertModel.from_pretrained("albert-base-v2" )
             lowerCamelCase__	:       Any					=						torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]] )
             lowerCamelCase__	:       int					=						torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
             with torch.no_grad():
                   lowerCamelCase__	:       List[Any]					=						model(__lowerCamelCase      ,   attention_mask=__lowerCamelCase )[0]
             lowerCamelCase__	:       Tuple					=						torch.Size((1, 11, 768) )
             self.assertEqual(output.shape      ,   __lowerCamelCase )
             lowerCamelCase__	:       Dict					=						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]      ,   __lowerCamelCase      ,   atol=1E-4 ) )


