




import argparse
import json
from collections import OrderedDict
from pathlib import Path

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

from transformers import DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCAmelCase    : Dict		= logging.get_logger(__name__)

# here we list all keys to be renamed (original name on the left, our name on the right)
_lowerCAmelCase    : Optional[int]		= []
for i in range(6):
  # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
  rename_keys.append(
      (F'''transformer.encoder.layers.{i}.self_attn.out_proj.weight''', F'''encoder.layers.{i}.self_attn.out_proj.weight''')
  )
  rename_keys.append(
      (F'''transformer.encoder.layers.{i}.self_attn.out_proj.bias''', F'''encoder.layers.{i}.self_attn.out_proj.bias''')
  )
  rename_keys.append((F'''transformer.encoder.layers.{i}.linear1.weight''', F'''encoder.layers.{i}.fc1.weight'''))
  rename_keys.append((F'''transformer.encoder.layers.{i}.linear1.bias''', F'''encoder.layers.{i}.fc1.bias'''))
  rename_keys.append((F'''transformer.encoder.layers.{i}.linear2.weight''', F'''encoder.layers.{i}.fc2.weight'''))
  rename_keys.append((F'''transformer.encoder.layers.{i}.linear2.bias''', F'''encoder.layers.{i}.fc2.bias'''))
  rename_keys.append(
      (F'''transformer.encoder.layers.{i}.norm1.weight''', F'''encoder.layers.{i}.self_attn_layer_norm.weight''')
  )
  rename_keys.append((F'''transformer.encoder.layers.{i}.norm1.bias''', F'''encoder.layers.{i}.self_attn_layer_norm.bias'''))
  rename_keys.append((F'''transformer.encoder.layers.{i}.norm2.weight''', F'''encoder.layers.{i}.final_layer_norm.weight'''))
  rename_keys.append((F'''transformer.encoder.layers.{i}.norm2.bias''', F'''encoder.layers.{i}.final_layer_norm.bias'''))
  # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
  rename_keys.append(
      (F'''transformer.decoder.layers.{i}.self_attn.out_proj.weight''', F'''decoder.layers.{i}.self_attn.out_proj.weight''')
  )
  rename_keys.append(
      (F'''transformer.decoder.layers.{i}.self_attn.out_proj.bias''', F'''decoder.layers.{i}.self_attn.out_proj.bias''')
  )
  rename_keys.append(
      (
          F'''transformer.decoder.layers.{i}.multihead_attn.out_proj.weight''',
          F'''decoder.layers.{i}.encoder_attn.out_proj.weight''',
      )
  )
  rename_keys.append(
      (
          F'''transformer.decoder.layers.{i}.multihead_attn.out_proj.bias''',
          F'''decoder.layers.{i}.encoder_attn.out_proj.bias''',
      )
  )
  rename_keys.append((F'''transformer.decoder.layers.{i}.linear1.weight''', F'''decoder.layers.{i}.fc1.weight'''))
  rename_keys.append((F'''transformer.decoder.layers.{i}.linear1.bias''', F'''decoder.layers.{i}.fc1.bias'''))
  rename_keys.append((F'''transformer.decoder.layers.{i}.linear2.weight''', F'''decoder.layers.{i}.fc2.weight'''))
  rename_keys.append((F'''transformer.decoder.layers.{i}.linear2.bias''', F'''decoder.layers.{i}.fc2.bias'''))
  rename_keys.append(
      (F'''transformer.decoder.layers.{i}.norm1.weight''', F'''decoder.layers.{i}.self_attn_layer_norm.weight''')
  )
  rename_keys.append((F'''transformer.decoder.layers.{i}.norm1.bias''', F'''decoder.layers.{i}.self_attn_layer_norm.bias'''))
  rename_keys.append(
      (F'''transformer.decoder.layers.{i}.norm2.weight''', F'''decoder.layers.{i}.encoder_attn_layer_norm.weight''')
  )
  rename_keys.append(
      (F'''transformer.decoder.layers.{i}.norm2.bias''', F'''decoder.layers.{i}.encoder_attn_layer_norm.bias''')
  )
  rename_keys.append((F'''transformer.decoder.layers.{i}.norm3.weight''', F'''decoder.layers.{i}.final_layer_norm.weight'''))
  rename_keys.append((F'''transformer.decoder.layers.{i}.norm3.bias''', F'''decoder.layers.{i}.final_layer_norm.bias'''))

# convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
rename_keys.extend(
    [
        ("input_proj.weight", "input_projection.weight"),
        ("input_proj.bias", "input_projection.bias"),
        ("query_embed.weight", "query_position_embeddings.weight"),
        ("transformer.decoder.norm.weight", "decoder.layernorm.weight"),
        ("transformer.decoder.norm.bias", "decoder.layernorm.bias"),
        ("class_embed.weight", "class_labels_classifier.weight"),
        ("class_embed.bias", "class_labels_classifier.bias"),
        ("bbox_embed.layers.0.weight", "bbox_predictor.layers.0.weight"),
        ("bbox_embed.layers.0.bias", "bbox_predictor.layers.0.bias"),
        ("bbox_embed.layers.1.weight", "bbox_predictor.layers.1.weight"),
        ("bbox_embed.layers.1.bias", "bbox_predictor.layers.1.bias"),
        ("bbox_embed.layers.2.weight", "bbox_predictor.layers.2.weight"),
        ("bbox_embed.layers.2.bias", "bbox_predictor.layers.2.bias"),
    ]
)



def     lowerCAmelCase		(		_lowerCAmelCase				: str    ,				_lowerCAmelCase				: Optional[int]    ,				_lowerCAmelCase				: List[str]						):




 """simple docstring"""







 UpperCAmelCase__							=    state_dict.pop(_lowerCAmelCase						)
 UpperCAmelCase__							=    val



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




 """simple docstring"""







 UpperCAmelCase__							=    OrderedDict()
 for key, value in state_dict.items():
  if "backbone.0.body" in key:
   UpperCAmelCase__							=    key.replace("backbone.0.body"    ,				"backbone.conv_encoder.model"						)
   UpperCAmelCase__							=    value
  else:
   UpperCAmelCase__							=    value

 return new_state_dict



def     lowerCAmelCase		(		_lowerCAmelCase				: str    ,				_lowerCAmelCase				: List[str]=False						):




 """simple docstring"""







 UpperCAmelCase__							=    ""
 if is_panoptic:
  UpperCAmelCase__							=    "detr."

 # first: transformer encoder
 for i in range(6						):
  # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
  UpperCAmelCase__							=    state_dict.pop(F'''{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight'''						)
  UpperCAmelCase__							=    state_dict.pop(F'''{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias'''						)
  # next, add query, keys and values (in that order) to the state dict
  UpperCAmelCase__							=    in_proj_weight[:256, :]
  UpperCAmelCase__							=    in_proj_bias[:256]
  UpperCAmelCase__							=    in_proj_weight[256:512, :]
  UpperCAmelCase__							=    in_proj_bias[256:512]
  UpperCAmelCase__							=    in_proj_weight[-256:, :]
  UpperCAmelCase__							=    in_proj_bias[-256:]
 # next: transformer decoder (which is a bit more complex because it also includes cross-attention)
 for i in range(6						):
  # read in weights + bias of input projection layer of self-attention
  UpperCAmelCase__							=    state_dict.pop(F'''{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight'''						)
  UpperCAmelCase__							=    state_dict.pop(F'''{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias'''						)
  # next, add query, keys and values (in that order) to the state dict
  UpperCAmelCase__							=    in_proj_weight[:256, :]
  UpperCAmelCase__							=    in_proj_bias[:256]
  UpperCAmelCase__							=    in_proj_weight[256:512, :]
  UpperCAmelCase__							=    in_proj_bias[256:512]
  UpperCAmelCase__							=    in_proj_weight[-256:, :]
  UpperCAmelCase__							=    in_proj_bias[-256:]
  # read in weights + bias of input projection layer of cross-attention
  UpperCAmelCase__							=    state_dict.pop(
      F'''{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight'''						)
  UpperCAmelCase__							=    state_dict.pop(F'''{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias'''						)
  # next, add query, keys and values (in that order) of cross-attention to the state dict
  UpperCAmelCase__							=    in_proj_weight_cross_attn[:256, :]
  UpperCAmelCase__							=    in_proj_bias_cross_attn[:256]
  UpperCAmelCase__							=    in_proj_weight_cross_attn[256:512, :]
  UpperCAmelCase__							=    in_proj_bias_cross_attn[256:512]
  UpperCAmelCase__							=    in_proj_weight_cross_attn[-256:, :]
  UpperCAmelCase__							=    in_proj_bias_cross_attn[-256:]



def     lowerCAmelCase		(		):




 """simple docstring"""







 UpperCAmelCase__							=    "http://images.cocodataset.org/val2017/000000039769.jpg"
 UpperCAmelCase__							=    Image.open(requests.get(_lowerCAmelCase    ,				stream=_lowerCAmelCase						).raw						)

 return im



@torch.no_grad()
def     lowerCAmelCase		(		_lowerCAmelCase				: Any    ,				_lowerCAmelCase				: Optional[int]						):




 """simple docstring"""







 UpperCAmelCase__							=    DetrConfig()
 # set backbone and dilation attributes
 if "resnet101" in model_name:
  UpperCAmelCase__							=    "resnet101"
 if "dc5" in model_name:
  UpperCAmelCase__							=    True
 UpperCAmelCase__							=    "panoptic" in model_name
 if is_panoptic:
  UpperCAmelCase__							=    250
 else:
  UpperCAmelCase__							=    91
  UpperCAmelCase__							=    "huggingface/label-files"
  UpperCAmelCase__							=    "coco-detection-id2label.json"
  UpperCAmelCase__							=    json.load(open(hf_hub_download(_lowerCAmelCase    ,				_lowerCAmelCase    ,				repo_type="dataset"						)    ,				"r"						)						)
  UpperCAmelCase__							=    {int(_lowerCAmelCase						): v for k, v in idalabel.items()}
  UpperCAmelCase__							=    idalabel
  UpperCAmelCase__							=    {v: k for k, v in idalabel.items()}

 # load image processor
 UpperCAmelCase__							=    "coco_panoptic" if is_panoptic else "coco_detection"
 UpperCAmelCase__							=    DetrImageProcessor(format=_lowerCAmelCase						)

 # prepare image
 UpperCAmelCase__							=    prepare_img()
 UpperCAmelCase__							=    image_processor(images=_lowerCAmelCase    ,				return_tensors="pt"						)
 UpperCAmelCase__							=    encoding["pixel_values"]

 logger.info(F'''Converting model {model_name}...'''						)

 # load original model from torch hub
 UpperCAmelCase__							=    torch.hub.load("facebookresearch/detr"    ,				_lowerCAmelCase    ,				pretrained=_lowerCAmelCase						).eval()
 UpperCAmelCase__							=    detr.state_dict()
 # rename keys
 for src, dest in rename_keys:
  if is_panoptic:
   UpperCAmelCase__							=    "detr." + src
  rename_key(_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase						)
 UpperCAmelCase__							=    rename_backbone_keys(_lowerCAmelCase						)
 # query, key and value matrices need special treatment
 read_in_q_k_v(_lowerCAmelCase    ,				is_panoptic=_lowerCAmelCase						)
 # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
 UpperCAmelCase__							=    "detr.model." if is_panoptic else "model."
 for key in state_dict.copy().keys():
  if is_panoptic:
   if (
       key.startswith("detr"						)
       and not key.startswith("class_labels_classifier"						)
       and not key.startswith("bbox_predictor"						)
   ):
    UpperCAmelCase__							=    state_dict.pop(_lowerCAmelCase						)
    UpperCAmelCase__							=    val
   elif "class_labels_classifier" in key or "bbox_predictor" in key:
    UpperCAmelCase__							=    state_dict.pop(_lowerCAmelCase						)
    UpperCAmelCase__							=    val
   elif key.startswith("bbox_attention"						) or key.startswith("mask_head"						):
    continue
   else:
    UpperCAmelCase__							=    state_dict.pop(_lowerCAmelCase						)
    UpperCAmelCase__							=    val
  else:
   if not key.startswith("class_labels_classifier"						) and not key.startswith("bbox_predictor"						):
    UpperCAmelCase__							=    state_dict.pop(_lowerCAmelCase						)
    UpperCAmelCase__							=    val
    # finally, create HuggingFace model and load state dict
 UpperCAmelCase__							=    DetrForSegmentation(_lowerCAmelCase						) if is_panoptic else DetrForObjectDetection(_lowerCAmelCase						)
 model.load_state_dict(_lowerCAmelCase						)
 model.eval()
 # verify our conversion
 UpperCAmelCase__							=    detr(_lowerCAmelCase						)
 UpperCAmelCase__							=    model(_lowerCAmelCase						)
 assert torch.allclose(outputs.logits    ,				original_outputs["pred_logits"]    ,				atol=1E-4						)
 assert torch.allclose(outputs.pred_boxes    ,				original_outputs["pred_boxes"]    ,				atol=1E-4						)
 if is_panoptic:
  assert torch.allclose(outputs.pred_masks    ,				original_outputs["pred_masks"]    ,				atol=1E-4						)

 # Save model and image processor
 logger.info(F'''Saving PyTorch model and image processor to {pytorch_dump_folder_path}...'''						)
 Path(_lowerCAmelCase						).mkdir(exist_ok=_lowerCAmelCase						)
 model.save_pretrained(_lowerCAmelCase						)
 image_processor.save_pretrained(_lowerCAmelCase						)


if __name__ == "__main__":
  _lowerCAmelCase    : Tuple		= argparse.ArgumentParser()

  parser.add_argument(
      "--model_name", default="detr_resnet50", type=str, help="Name of the DETR model you'd like to convert."
  )
  parser.add_argument(
      "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
  )
  _lowerCAmelCase    : List[Any]		= parser.parse_args()
  convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path)











from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
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 TF_MODEL_FOR_PRETRAINING_MAPPING
  from transformers.models.albert.modeling_tf_albert import (
      TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFAlbertForMaskedLM,
      TFAlbertForMultipleChoice,
      TFAlbertForPreTraining,
      TFAlbertForQuestionAnswering,
      TFAlbertForSequenceClassification,
      TFAlbertForTokenClassification,
      TFAlbertModel,
  )
class       _UpperCamelCase   :
      def __init__(      self :Optional[int]			,   lowerCamelCase :Tuple			,   lowerCamelCase :Union[str, Any]=13			,   lowerCamelCase :List[str]=7			,   lowerCamelCase :List[Any]=True			,   lowerCamelCase :Any=True			,   lowerCamelCase :Any=True			,   lowerCamelCase :Optional[Any]=True			,   lowerCamelCase :Optional[int]=99			,   lowerCamelCase :Tuple=16			,   lowerCamelCase :Optional[Any]=32			,   lowerCamelCase :Optional[Any]=2			,   lowerCamelCase :List[Any]=4			,   lowerCamelCase :Dict=37			,   lowerCamelCase :Optional[Any]="gelu"			,   lowerCamelCase :Dict=0.1			,   lowerCamelCase :Optional[Any]=0.1			,   lowerCamelCase :Tuple=512			,   lowerCamelCase :Optional[int]=16			,   lowerCamelCase :Any=2			,   lowerCamelCase :Tuple=0.02			,   lowerCamelCase :Union[str, Any]=3			,   lowerCamelCase :str=4			,   lowerCamelCase :Dict=None			,   )  ->      List[Any]:
       UpperCAmelCase__							=    parent
       UpperCAmelCase__							=    13
       UpperCAmelCase__							=    7
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    99
       UpperCAmelCase__							=    16
       UpperCAmelCase__							=    32
       UpperCAmelCase__							=    2
       UpperCAmelCase__							=    4
       UpperCAmelCase__							=    37
       UpperCAmelCase__							=    "gelu"
       UpperCAmelCase__							=    0.1
       UpperCAmelCase__							=    0.1
       UpperCAmelCase__							=    512
       UpperCAmelCase__							=    16
       UpperCAmelCase__							=    2
       UpperCAmelCase__							=    0.02
       UpperCAmelCase__							=    3
       UpperCAmelCase__							=    4
       UpperCAmelCase__							=    None
      def 						UpperCAmelCase_    (      self :Optional[int]   )  ->      Dict:
       UpperCAmelCase__							=    ids_tensor([self.batch_size, self.seq_length]			,   self.vocab_size   )

       UpperCAmelCase__							=    None
       if self.use_input_mask:
        UpperCAmelCase__							=    random_attention_mask([self.batch_size, self.seq_length]   )

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

       UpperCAmelCase__							=    None
       UpperCAmelCase__							=    None
       UpperCAmelCase__							=    None
       if self.use_labels:
        UpperCAmelCase__							=    ids_tensor([self.batch_size]			,   self.type_sequence_label_size   )
        UpperCAmelCase__							=    ids_tensor([self.batch_size, self.seq_length]			,   self.num_labels   )
        UpperCAmelCase__							=    ids_tensor([self.batch_size]			,   self.num_choices   )

       UpperCAmelCase__							=    AlbertConfig(
           vocab_size=self.vocab_size			,   hidden_size=self.hidden_size			,   embedding_size=self.embedding_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 config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
      def 						UpperCAmelCase_    (      self :Optional[int]			,   lowerCamelCase :str			,   lowerCamelCase :Tuple			,   lowerCamelCase :Tuple			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :int			,   lowerCamelCase :str   )  ->      Optional[Any]:
       UpperCAmelCase__							=    TFAlbertModel(config=lowerCamelCase   )
       # inputs = {'input_ids': input_ids,
       #           'attention_mask': input_mask,
       #           'token_type_ids': token_type_ids}
       # sequence_output, pooled_output = model(**inputs)
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )

       UpperCAmelCase__							=    [input_ids, input_mask]
       UpperCAmelCase__							=    model(lowerCamelCase   )

       UpperCAmelCase__							=    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 						UpperCAmelCase_    (      self :Optional[Any]			,   lowerCamelCase :Any			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Any			,   lowerCamelCase :List[str]			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Tuple   )  ->      Optional[int]:
       UpperCAmelCase__							=    self.num_labels
       UpperCAmelCase__							=    TFAlbertForPreTraining(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(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, self.num_labels)   )
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :Dict			,   lowerCamelCase :List[Any]			,   lowerCamelCase :List[str]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :Tuple			,   lowerCamelCase :Any   )  ->      Optional[Any]:
       UpperCAmelCase__							=    TFAlbertForMaskedLM(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.seq_length, self.vocab_size)   )
      def 						UpperCAmelCase_    (      self :Tuple			,   lowerCamelCase :str			,   lowerCamelCase :Dict			,   lowerCamelCase :Dict			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Optional[Any]   )  ->      str:
       UpperCAmelCase__							=    self.num_labels
       UpperCAmelCase__							=    TFAlbertForSequenceClassification(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.num_labels)   )
      def 						UpperCAmelCase_    (      self :Dict			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :str			,   lowerCamelCase :Dict			,   lowerCamelCase :List[Any]			,   lowerCamelCase :List[str]			,   lowerCamelCase :int			,   lowerCamelCase :Optional[Any]   )  ->      Dict:
       UpperCAmelCase__							=    TFAlbertForQuestionAnswering(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(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 						UpperCAmelCase_    (      self :Optional[int]			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :Dict			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Tuple			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Union[str, Any]   )  ->      int:
       UpperCAmelCase__							=    self.num_choices
       UpperCAmelCase__							=    TFAlbertForMultipleChoice(config=lowerCamelCase   )
       UpperCAmelCase__							=    tf.tile(tf.expand_dims(lowerCamelCase			,   1   )			,   (1, self.num_choices, 1)   )
       UpperCAmelCase__							=    tf.tile(tf.expand_dims(lowerCamelCase			,   1   )			,   (1, self.num_choices, 1)   )
       UpperCAmelCase__							=    tf.tile(tf.expand_dims(lowerCamelCase			,   1   )			,   (1, self.num_choices, 1)   )
       UpperCAmelCase__							=    {
           "input_ids": multiple_choice_inputs_ids,
           "attention_mask": multiple_choice_input_mask,
           "token_type_ids": multiple_choice_token_type_ids,
       }
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertListEqual(list(result["logits"].shape   )			,   [self.batch_size, self.num_choices]   )
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :List[str]			,   lowerCamelCase :int			,   lowerCamelCase :int			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Dict			,   lowerCamelCase :Dict   )  ->      List[str]:
       UpperCAmelCase__							=    self.num_labels
       UpperCAmelCase__							=    TFAlbertForTokenClassification(config=lowerCamelCase   )
       UpperCAmelCase__							=    {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "token_type_ids": token_type_ids,
       }
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertListEqual(list(result["logits"].shape   )			,   [self.batch_size, self.seq_length, self.num_labels]   )







      def 						UpperCAmelCase_    (      self :str   )  ->      Any:
       UpperCAmelCase__							=    self.prepare_config_and_inputs()
       (
           (
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			
       )							=    config_and_inputs
       UpperCAmelCase__							=    {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
       return config, inputs_dict



@require_tf
class       _UpperCamelCase   (					lowerCAmelCase						,							lowerCAmelCase						,							unittest.TestCase		):
      UpperCAmelCase_          =		(
          (
              TFAlbertModel,
              TFAlbertForPreTraining,
              TFAlbertForMaskedLM,
              TFAlbertForSequenceClassification,
              TFAlbertForQuestionAnswering,
              TFAlbertForTokenClassification,
              TFAlbertForMultipleChoice,
          )
          if is_tf_available()
          else ()
      )
      UpperCAmelCase_          =		(
          {
              """feature-extraction""": TFAlbertModel,
              """fill-mask""": TFAlbertForMaskedLM,
              """question-answering""": TFAlbertForQuestionAnswering,
              """text-classification""": TFAlbertForSequenceClassification,
              """token-classification""": TFAlbertForTokenClassification,
              """zero-shot""": TFAlbertForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      UpperCAmelCase_          =		False
      UpperCAmelCase_          =		False
      def 						UpperCAmelCase_    (      self :str			,   lowerCamelCase :Any			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :Optional[int]=False   )  ->      str:
       UpperCAmelCase__							=    super()._prepare_for_class(lowerCamelCase			,   lowerCamelCase			,   return_labels=lowerCamelCase   )

       if return_labels:
        if model_class in get_values(lowerCamelCase   ):
         UpperCAmelCase__							=    tf.zeros(self.model_tester.batch_size			,   dtype=tf.intaa   )

       return inputs_dict
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      List[Any]:
       UpperCAmelCase__							=    TFAlbertModelTester(self   )
       UpperCAmelCase__							=    ConfigTester(self			,   config_class=lowerCamelCase			,   hidden_size=37   )
      def 						UpperCAmelCase_    (      self :Dict   )  ->      Union[str, Any]:
       self.config_tester.run_common_tests()
      def 						UpperCAmelCase_    (      self :Optional[Any]   )  ->      Tuple:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_model(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      Dict:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_pretraining(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      str:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_masked_lm(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Dict   )  ->      int:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_multiple_choice(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Union[str, Any]   )  ->      Dict:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_sequence_classification(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :str   )  ->      int:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_question_answering(*lowerCamelCase   )







      @slow
      def 						UpperCAmelCase_    (      self :int   )  ->      str:
       for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase__							=    TFAlbertModel.from_pretrained(lowerCamelCase   )
        self.assertIsNotNone(lowerCamelCase   )



@require_tf
class       _UpperCamelCase   (					unittest.TestCase		):
      @slow
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      List[str]:
       UpperCAmelCase__							=    TFAlbertForPreTraining.from_pretrained("albert-base-v2"   )
       UpperCAmelCase__							=    tf.constant([[0, 1, 2, 3, 4, 5]]   )
       UpperCAmelCase__							=    model(lowerCamelCase   )[0]

       UpperCAmelCase__							=    [1, 6, 3_0000]
       self.assertEqual(output.shape			,   lowerCamelCase   )

       UpperCAmelCase__							=    tf.constant(
           [
               [
                   [4.59_56_68, 0.74_46_27_54, -1.81_81_47],
                   [4.5_95_43_47, 0.7_45_41_84, -1.8_18_82_58],
                   [4.5_95_49_05, 0.7_44_82_35, -1.8_18_23_16],
               ]
           ]   )
       tf.debugging.assert_near(output[:, :3, :3]			,   lowerCamelCase			,   atol=1e-4   )





