





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 timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import (
    BitConfig,
    ViTHybridConfig,
    ViTHybridForImageClassification,
    ViTHybridImageProcessor,
    ViTHybridModel,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
__lowerCamelCase :  Any										=  logging.get_logger(__name__)



def 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase=False  )		->				Union[str, Any]:
 UpperCamelCase			:      Any 			=     []

 # fmt: off
 # stem:
 rename_keys.append(("cls_token", "vit.embeddings.cls_token")  )
 rename_keys.append(("pos_embed", "vit.embeddings.position_embeddings")  )

 rename_keys.append(("patch_embed.proj.weight", "vit.embeddings.patch_embeddings.projection.weight")  )
 rename_keys.append(("patch_embed.proj.bias", "vit.embeddings.patch_embeddings.projection.bias")  )

 # backbone
 rename_keys.append(("patch_embed.backbone.stem.conv.weight", "vit.embeddings.patch_embeddings.backbone.bit.embedder.convolution.weight")  )
 rename_keys.append(("patch_embed.backbone.stem.norm.weight", "vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.weight")  )
 rename_keys.append(("patch_embed.backbone.stem.norm.bias", "vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.bias")  )

 for stage_idx in range(len(config.backbone_config.depths  )  ):
  for layer_idx in range(config.backbone_config.depths[stage_idx]  ):
   rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv1.weight""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv1.weight""")  )
   rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.weight""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.weight""")  )
   rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.bias""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.bias""")  )
   rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv2.weight""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv2.weight""")  )
   rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.weight""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.weight""")  )
   rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.bias""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.bias""")  )
   rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv3.weight""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv3.weight""")  )
   rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.weight""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.weight""")  )
   rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.bias""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.bias""")  )

  rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.conv.weight""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.conv.weight""")  )
  rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.weight""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.weight""")  )
  rename_keys.append((F"""patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.bias""", F"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.bias""")  )

 # transformer encoder
 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""")  )

 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"
  UpperCamelCase			:      str 			=     [(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"),
      ]  )
 # fmt: on

 return rename_keys



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



def 	A_       (  _lowerCAmelCase  )		->				Tuple:
 UpperCamelCase			:      str 			=     ["head.weight", "head.bias"]
 for k in ignore_keys:
  state_dict.pop(_lowerCAmelCase							,     _lowerCAmelCase  )



def 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase  )		->				Any:
 UpperCamelCase			:      Any 			=     dct.pop(_lowerCAmelCase  )
 UpperCamelCase			:      Union[str, Any] 			=     val



def 	A_       (  )		->				Any:
 UpperCamelCase			:      Optional[int] 			=     "http://images.cocodataset.org/val2017/000000039769.jpg"
 UpperCamelCase			:      str 			=     Image.open(requests.get(_lowerCAmelCase							,     stream=_lowerCAmelCase  ).raw  )
 return im



@torch.no_grad()
def 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase=False  )		->				List[str]:
 UpperCamelCase			:      Any 			=     BitConfig(
     global_padding="same"							,     layer_type="bottleneck"							,     depths=(3, 4, 9)							,     out_features=["stage3"]							,     embedding_dynamic_padding=_lowerCAmelCase							,     )
 UpperCamelCase			:      List[Any] 			=     ViTHybridConfig(backbone_config=_lowerCAmelCase							,     image_size=384							,     num_labels=1000  )
 UpperCamelCase			:      Tuple 			=     False

 # load original model from timm
 UpperCamelCase			:      int 			=     timm.create_model(_lowerCAmelCase							,     pretrained=_lowerCAmelCase  )
 timm_model.eval()

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

 UpperCamelCase			:      Union[str, Any] 			=     "huggingface/label-files"
 UpperCamelCase			:      List[str] 			=     "imagenet-1k-id2label.json"
 UpperCamelCase			:      int 			=     json.load(open(hf_hub_download(_lowerCAmelCase							,     _lowerCAmelCase							,     repo_type="dataset"  )							,     "r"  )  )
 UpperCamelCase			:      List[Any] 			=     {int(_lowerCAmelCase  ): v for k, v in idalabel.items()}
 UpperCamelCase			:      Tuple 			=     idalabel
 UpperCamelCase			:      Optional[Any] 			=     {v: k for k, v in idalabel.items()}

 # load HuggingFace model
 if vit_name[-5:] == "in21k":
  UpperCamelCase			:      Optional[int] 			=     ViTHybridModel(_lowerCAmelCase  ).eval()
 else:
  UpperCamelCase			:      Any 			=     ViTHybridForImageClassification(_lowerCAmelCase  ).eval()
 model.load_state_dict(_lowerCAmelCase  )

 # create image processor
 UpperCamelCase			:      Dict 			=     create_transform(**resolve_data_config({}							,     model=_lowerCAmelCase  )  )
 UpperCamelCase			:      Union[str, Any] 			=     transform.transforms

 UpperCamelCase			:      List[Any] 			=     {
     "bilinear": PILImageResampling.BILINEAR,
     "bicubic": PILImageResampling.BICUBIC,
     "nearest": PILImageResampling.NEAREST,
 }

 UpperCamelCase			:      Tuple 			=     ViTHybridImageProcessor(
     do_resize=_lowerCAmelCase							,     size={"shortest_edge": timm_transforms[0].size}							,     resample=pillow_resamplings[timm_transforms[0].interpolation.value]							,     do_center_crop=_lowerCAmelCase							,     crop_size={"height": timm_transforms[1].size[0], "width": timm_transforms[1].size[1]}							,     do_normalize=_lowerCAmelCase							,     image_mean=timm_transforms[-1].mean.tolist()							,     image_std=timm_transforms[-1].std.tolist()							,     )

 UpperCamelCase			:      Tuple 			=     prepare_img()
 UpperCamelCase			:      Dict 			=     transform(_lowerCAmelCase  ).unsqueeze(0  )
 UpperCamelCase			:      Tuple 			=     processor(_lowerCAmelCase							,     return_tensors="pt"  ).pixel_values

 # verify pixel values
 assert torch.allclose(_lowerCAmelCase							,     _lowerCAmelCase  )

 # verify logits
 with torch.no_grad():
  UpperCamelCase			:      Optional[Any] 			=     model(_lowerCAmelCase  )
  UpperCamelCase			:      str 			=     outputs.logits

 print("Predicted class:"							,     logits.argmax(-1  ).item()  )
 if base_model:
  UpperCamelCase			:      Union[str, Any] 			=     timm_model.forward_features(_lowerCAmelCase  )
  assert timm_pooled_output.shape == outputs.pooler_output.shape
  assert torch.allclose(_lowerCAmelCase							,     outputs.pooler_output							,     atol=1e-3  )
 else:
  UpperCamelCase			:      List[str] 			=     timm_model(_lowerCAmelCase  )
  assert timm_logits.shape == outputs.logits.shape
  assert torch.allclose(_lowerCAmelCase							,     outputs.logits							,     atol=1e-3  )
 print("Looks ok!"  )

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

 if push_to_hub:
  print(F"""Pushing model and processor to the hub {vit_name}"""  )
  model.push_to_hub(F"""ybelkada/{vit_name}"""  )
  processor.push_to_hub(F"""ybelkada/{vit_name}"""  )


if __name__ == "__main__":
    __lowerCamelCase :  int										=  argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        """--vit_name""",
        default="""vit_base_r50_s16_384""",
        type=str,
        help="""Name of the hybrid 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."""
    )
    parser.add_argument(
        """--push_to_hub""", action="""store_true""", help="""Whether to upload the model to the HuggingFace hub."""
    )

    __lowerCamelCase :  Union[str, Any]										=  parser.parse_args()
    convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)







from __future__ import annotations

import os
import tempfile
import unittest

from transformers import ConvBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TFConvBertForMaskedLM,
        TFConvBertForMultipleChoice,
        TFConvBertForQuestionAnswering,
        TFConvBertForSequenceClassification,
        TFConvBertForTokenClassification,
        TFConvBertModel,
    )



class       A__     :





      def __init__(				self ,							A_ ,							A_=13 ,							A_=7 ,							A_=True ,							A_=True ,							A_=True ,							A_=True ,							A_=99 ,							A_=32 ,							A_=2 ,							A_=4 ,							A_=37 ,							A_="gelu" ,							A_=0.1 ,							A_=0.1 ,							A_=512 ,							A_=16 ,							A_=2 ,							A_=0.02 ,							A_=3 ,							A_=4 ,							A_=None ,							):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     parent
       UpperCamelCase			:      str 			=     13
       UpperCamelCase			:      int 			=     7
       UpperCamelCase			:      str 			=     True
       UpperCamelCase			:      Dict 			=     True
       UpperCamelCase			:      str 			=     True
       UpperCamelCase			:      Tuple 			=     True
       UpperCamelCase			:      List[str] 			=     99
       UpperCamelCase			:      Optional[Any] 			=     384
       UpperCamelCase			:      Tuple 			=     2
       UpperCamelCase			:      Union[str, Any] 			=     4
       UpperCamelCase			:      Dict 			=     37
       UpperCamelCase			:      Any 			=     "gelu"
       UpperCamelCase			:      List[Any] 			=     0.1
       UpperCamelCase			:      int 			=     0.1
       UpperCamelCase			:      Tuple 			=     512
       UpperCamelCase			:      List[Any] 			=     16
       UpperCamelCase			:      int 			=     2
       UpperCamelCase			:      Dict 			=     0.02
       UpperCamelCase			:      Optional[Any] 			=     3
       UpperCamelCase			:      List[Any] 			=     4
       UpperCamelCase			:      Dict 			=     128
       UpperCamelCase			:      Optional[Any] 			=     2
       UpperCamelCase			:      Optional[int] 			=     9
       UpperCamelCase			:      Optional[int] 			=     1
       UpperCamelCase			:      Union[str, Any] 			=     None





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      int 			=     ids_tensor([self.batch_size, self.seq_length] ,							self.vocab_size     )

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

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

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

       UpperCamelCase			:      Any 			=     ConvBertConfig(
           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 ,							return_dict=A_ ,							)

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





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      str 			=     TFConvBertModel(config=A_     )
       UpperCamelCase			:      Optional[Any] 			=     {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}

       UpperCamelCase			:      Optional[int] 			=     [input_ids, input_mask]
       UpperCamelCase			:      Any 			=     model(A_     )

       UpperCamelCase			:      int 			=     model(A_     )

       self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, self.seq_length, self.hidden_size)     )





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      Tuple 			=     TFConvBertForMaskedLM(config=A_     )
       UpperCamelCase			:      int 			=     {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "token_type_ids": token_type_ids,
       }
       UpperCamelCase			:      Dict 			=     model(A_     )
       self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.vocab_size)     )





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     self.num_labels
       UpperCamelCase			:      int 			=     TFConvBertForSequenceClassification(config=A_     )
       UpperCamelCase			:      List[Any] 			=     {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "token_type_ids": token_type_ids,
       }

       UpperCamelCase			:      Optional[Any] 			=     model(A_     )
       self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.num_labels)     )





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      List[str] 			=     self.num_choices
       UpperCamelCase			:      str 			=     TFConvBertForMultipleChoice(config=A_     )
       UpperCamelCase			:      List[Any] 			=     tf.tile(tf.expand_dims(A_ ,							1     ) ,							(1, self.num_choices, 1)     )
       UpperCamelCase			:      Dict 			=     tf.tile(tf.expand_dims(A_ ,							1     ) ,							(1, self.num_choices, 1)     )
       UpperCamelCase			:      Any 			=     tf.tile(tf.expand_dims(A_ ,							1     ) ,							(1, self.num_choices, 1)     )
       UpperCamelCase			:      List[str] 			=     {
           "input_ids": multiple_choice_inputs_ids,
           "attention_mask": multiple_choice_input_mask,
           "token_type_ids": multiple_choice_token_type_ids,
       }
       UpperCamelCase			:      Optional[Any] 			=     model(A_     )
       self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.num_choices)     )





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     self.num_labels
       UpperCamelCase			:      str 			=     TFConvBertForTokenClassification(config=A_     )
       UpperCamelCase			:      List[Any] 			=     {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "token_type_ids": token_type_ids,
       }
       UpperCamelCase			:      str 			=     model(A_     )
       self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.num_labels)     )





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      List[str] 			=     TFConvBertForQuestionAnswering(config=A_     )
       UpperCamelCase			:      Union[str, Any] 			=     {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "token_type_ids": token_type_ids,
       }

       UpperCamelCase			:      Union[str, Any] 			=     model(A_     )
       self.parent.assertEqual(result.start_logits.shape ,							(self.batch_size, self.seq_length)     )
       self.parent.assertEqual(result.end_logits.shape ,							(self.batch_size, self.seq_length)     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




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



@require_tf
class       A__     ( __snake_case ,     __snake_case ,     unittest.TestCase    ):
      _UpperCAmelCase			:Dict							=     (
          (
              TFConvBertModel,
              TFConvBertForMaskedLM,
              TFConvBertForQuestionAnswering,
              TFConvBertForSequenceClassification,
              TFConvBertForTokenClassification,
              TFConvBertForMultipleChoice,
          )
          if is_tf_available()
          else ()
      )
      _UpperCAmelCase			:Optional[Any]							=     (
          {
              'feature-extraction': TFConvBertModel,
              'fill-mask': TFConvBertForMaskedLM,
              'question-answering': TFConvBertForQuestionAnswering,
              'text-classification': TFConvBertForSequenceClassification,
              'token-classification': TFConvBertForTokenClassification,
              'zero-shot': TFConvBertForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      _UpperCAmelCase			:Any							=     False
      _UpperCAmelCase			:int							=     False
      _UpperCAmelCase			:str							=     False





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     TFConvBertModelTester(self     )
       UpperCamelCase			:      Dict 			=     ConfigTester(self ,							config_class=A_ ,							hidden_size=37     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      List[str] 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      List[str] 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Tuple 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Optional[Any] 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*A_     )





      @slow
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			,					UpperCamelCase			:      Dict 			=     self.model_tester.prepare_config_and_inputs_for_common()
       UpperCamelCase			:      Optional[Any] 			=     True
       UpperCamelCase			:      Any 			=     True

       if hasattr(A_ ,							"use_cache"     ):
        UpperCamelCase			:      List[str] 			=     True

       UpperCamelCase			:      List[Any] 			=     getattr(self.model_tester ,							"encoder_seq_length" ,							self.model_tester.seq_length     )
       UpperCamelCase			:      Any 			=     getattr(self.model_tester ,							"key_length" ,							A_     )

       for model_class in self.all_model_classes:
        UpperCamelCase			:      List[Any] 			=     self._prepare_for_class(A_ ,							A_     )
        UpperCamelCase			:      Dict 			=     model_class(A_     )
        UpperCamelCase			:      Optional[int] 			=     len(model(A_     )     )

        with tempfile.TemporaryDirectory() as tmpdirname:
         model.save_pretrained(A_ ,							saved_model=A_     )
         UpperCamelCase			:      Union[str, Any] 			=     os.path.join(A_ ,							"saved_model" ,							"1"     )
         UpperCamelCase			:      Dict 			=     tf.keras.models.load_model(A_     )
         UpperCamelCase			:      str 			=     model(A_     )

         if self.is_encoder_decoder:
          UpperCamelCase			:      Union[str, Any] 			=     outputs["encoder_hidden_states"]
          UpperCamelCase			:      Any 			=     outputs["encoder_attentions"]
         else:
          UpperCamelCase			:      Any 			=     outputs["hidden_states"]
          UpperCamelCase			:      List[str] 			=     outputs["attentions"]

         self.assertEqual(len(A_     ) ,							A_     )

         UpperCamelCase			:      int 			=     getattr(
             self.model_tester ,							"expected_num_hidden_layers" ,							self.model_tester.num_hidden_layers + 1     )

         self.assertEqual(len(A_     ) ,							A_     )
         self.assertListEqual(
             list(output_hidden_states[0].shape[-2:]     ) ,							[self.model_tester.seq_length, self.model_tester.hidden_size] ,							)

         self.assertEqual(len(A_     ) ,							self.model_tester.num_hidden_layers     )
         self.assertListEqual(
             list(output_attentions[0].shape[-3:]     ) ,							[self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length] ,							)





      @slow
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Union[str, Any] 			=     TFConvBertModel.from_pretrained("YituTech/conv-bert-base"     )
       self.assertIsNotNone(A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			,					UpperCamelCase			:      Tuple 			=     self.model_tester.prepare_config_and_inputs_for_common()
       UpperCamelCase			:      Dict 			=     True
       UpperCamelCase			:      int 			=     getattr(self.model_tester ,							"decoder_seq_length" ,							self.model_tester.seq_length     )
       UpperCamelCase			:      Optional[int] 			=     getattr(self.model_tester ,							"encoder_seq_length" ,							self.model_tester.seq_length     )
       UpperCamelCase			:      Optional[int] 			=     getattr(self.model_tester ,							"key_length" ,							A_     )
       UpperCamelCase			:      Optional[Any] 			=     getattr(self.model_tester ,							"key_length" ,							A_     )

       def check_decoder_attentions_output(A_     ):
        UpperCamelCase			:      Optional[Any] 			=     len(A_     )
        self.assertEqual(out_len % 2 ,							0     )
        UpperCamelCase			:      Any 			=     outputs.decoder_attentions
        self.assertEqual(len(A_     ) ,							self.model_tester.num_hidden_layers     )
        self.assertListEqual(
            list(decoder_attentions[0].shape[-3:]     ) ,							[self.model_tester.num_attention_heads / 2, decoder_seq_length, decoder_key_length] ,							)

       def check_encoder_attentions_output(A_     ):
        UpperCamelCase			:      Dict 			=     [
            t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
        ]
        self.assertEqual(len(A_     ) ,							self.model_tester.num_hidden_layers     )
        self.assertListEqual(
            list(attentions[0].shape[-3:]     ) ,							[self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length] ,							)

       for model_class in self.all_model_classes:
        UpperCamelCase			:      Union[str, Any] 			=     True
        UpperCamelCase			:      List[Any] 			=     False
        UpperCamelCase			:      Dict 			=     model_class(A_     )
        UpperCamelCase			:      Dict 			=     model(self._prepare_for_class(A_ ,							A_     )     )
        UpperCamelCase			:      List[str] 			=     len(A_     )
        self.assertEqual(config.output_hidden_states ,							A_     )
        check_encoder_attentions_output(A_     )

        if self.is_encoder_decoder:
         UpperCamelCase			:      int 			=     model_class(A_     )
         UpperCamelCase			:      Tuple 			=     model(self._prepare_for_class(A_ ,							A_     )     )
         self.assertEqual(config.output_hidden_states ,							A_     )
         check_decoder_attentions_output(A_     )

        # Check that output attentions can also be changed via the config
        del inputs_dict["output_attentions"]
        UpperCamelCase			:      Tuple 			=     True
        UpperCamelCase			:      int 			=     model_class(A_     )
        UpperCamelCase			:      Dict 			=     model(self._prepare_for_class(A_ ,							A_     )     )
        self.assertEqual(config.output_hidden_states ,							A_     )
        check_encoder_attentions_output(A_     )

        # Check attention is always last and order is fine
        UpperCamelCase			:      Optional[int] 			=     True
        UpperCamelCase			:      List[str] 			=     True
        UpperCamelCase			:      Optional[int] 			=     model_class(A_     )
        UpperCamelCase			:      Optional[Any] 			=     model(self._prepare_for_class(A_ ,							A_     )     )

        self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1) ,							len(A_     )     )
        self.assertEqual(model.config.output_hidden_states ,							A_     )
        check_encoder_attentions_output(A_     )



@require_tf
class       A__     ( unittest.TestCase    ):





      @slow
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      str 			=     TFConvBertModel.from_pretrained("YituTech/conv-bert-base"     )
       UpperCamelCase			:      str 			=     tf.constant([[0, 1, 2, 3, 4, 5]]     )
       UpperCamelCase			:      List[str] 			=     model(A_     )[0]

       UpperCamelCase			:      int 			=     [1, 6, 768]
       self.assertEqual(output.shape ,							A_     )

       UpperCamelCase			:      List[str] 			=     tf.constant(
           [
               [
                   [-0.03_47_54_93, -0.4_68_60_34, -0.30_63_88_32],
                   [0.22_63_72_48, -0.26_98_86_46, -0.7_42_34_24],
                   [0.10_32_48_68, -0.45_01_35_08, -0.58_28_07_84],
               ]
           ]     )
       tf.debugging.assert_near(output[:, :3, :3] ,							A_ ,							atol=1e-4     )
