





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 DeformableDetrConfig, DeformableDetrForObjectDetection, DeformableDetrImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCAmelCase  :				List[str]						=				logging.get_logger(__name__)



def   lowerCAmelCase						(						_lowerCAmelCase							:				Tuple		):




 """simple docstring"""







 if "backbone.0.body" in orig_key:
  UpperCAmelCase__								=  orig_key.replace("backbone.0.body"  ,							"backbone.conv_encoder.model"		)
 if "transformer" in orig_key:
  UpperCAmelCase__								=  orig_key.replace("transformer."  ,							""		)
 if "norm1" in orig_key:
  if "encoder" in orig_key:
   UpperCAmelCase__								=  orig_key.replace("norm1"  ,							"self_attn_layer_norm"		)
  else:
   UpperCAmelCase__								=  orig_key.replace("norm1"  ,							"encoder_attn_layer_norm"		)
 if "norm2" in orig_key:
  if "encoder" in orig_key:
   UpperCAmelCase__								=  orig_key.replace("norm2"  ,							"final_layer_norm"		)
  else:
   UpperCAmelCase__								=  orig_key.replace("norm2"  ,							"self_attn_layer_norm"		)
 if "norm3" in orig_key:
  UpperCAmelCase__								=  orig_key.replace("norm3"  ,							"final_layer_norm"		)
 if "linear1" in orig_key:
  UpperCAmelCase__								=  orig_key.replace("linear1"  ,							"fc1"		)
 if "linear2" in orig_key:
  UpperCAmelCase__								=  orig_key.replace("linear2"  ,							"fc2"		)
 if "query_embed" in orig_key:
  UpperCAmelCase__								=  orig_key.replace("query_embed"  ,							"query_position_embeddings"		)
 if "cross_attn" in orig_key:
  UpperCAmelCase__								=  orig_key.replace("cross_attn"  ,							"encoder_attn"		)

 return orig_key



def   lowerCAmelCase						(						_lowerCAmelCase							:				Any		):




 """simple docstring"""







 for i in range(6		):
  # read in weights + bias of input projection layer of self-attention
  UpperCAmelCase__								=  state_dict.pop(F'''decoder.layers.{i}.self_attn.in_proj_weight'''		)
  UpperCAmelCase__								=  state_dict.pop(F'''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:]



def   lowerCAmelCase						(						):




 """simple docstring"""







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

 return im



@torch.no_grad()
def   lowerCAmelCase						(						_lowerCAmelCase							:				List[str]  ,							_lowerCAmelCase							:				Tuple  ,							_lowerCAmelCase							:				List[Any]  ,							_lowerCAmelCase							:				int  ,							_lowerCAmelCase							:				List[str]  ,							_lowerCAmelCase							:				Tuple  ,							_lowerCAmelCase							:				int  ,							):




 """simple docstring"""







 UpperCAmelCase__								=  DeformableDetrConfig()
 # set config attributes
 if single_scale:
  UpperCAmelCase__								=  1
 UpperCAmelCase__								=  dilation
 UpperCAmelCase__								=  with_box_refine
 UpperCAmelCase__								=  two_stage
 # set labels
 UpperCAmelCase__								=  91
 UpperCAmelCase__								=  "huggingface/label-files"
 UpperCAmelCase__								=  "coco-detection-id2label.json"
 UpperCAmelCase__								=  json.load(open(cached_download(hf_hub_url(lowercase_  ,							lowercase_  ,							repo_type="dataset"		)		)  ,							"r"		)		)
 UpperCAmelCase__								=  {int(lowercase_		): v for k, v in idalabel.items()}
 UpperCAmelCase__								=  idalabel
 UpperCAmelCase__								=  {v: k for k, v in idalabel.items()}

 # load image processor
 UpperCAmelCase__								=  DeformableDetrImageProcessor(format="coco_detection"		)

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

 logger.info("Converting model..."		)

 # load original state dict
 UpperCAmelCase__								=  torch.load(lowercase_  ,							map_location="cpu"		)["model"]
 # rename keys
 for key in state_dict.copy().keys():
  UpperCAmelCase__								=  state_dict.pop(lowercase_		)
  UpperCAmelCase__								=  val
 # query, key and value matrices need special treatment
 read_in_q_k_v(lowercase_		)
 # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
 UpperCAmelCase__								=  "model."
 for key in state_dict.copy().keys():
  if not key.startswith("class_embed"		) and not key.startswith("bbox_embed"		):
   UpperCAmelCase__								=  state_dict.pop(lowercase_		)
   UpperCAmelCase__								=  val
    # finally, create HuggingFace model and load state dict
 UpperCAmelCase__								=  DeformableDetrForObjectDetection(lowercase_		)
 model.load_state_dict(lowercase_		)
 model.eval()

 UpperCAmelCase__								=  "cuda" if torch.cuda.is_available() else "cpu"
 model.to(lowercase_		)
 # verify our conversion
 UpperCAmelCase__								=  model(pixel_values.to(lowercase_		)		)

 UpperCAmelCase__								=  torch.tensor(
     [[-9.6_645, -4.3_449, -5.8_705], [-9.7_035, -3.8_504, -5.0_724], [-10.5634, -5.3_379, -7.5_116]]		)
 UpperCAmelCase__								=  torch.tensor([[0.8_693, 0.2_289, 0.2_492], [0.3_150, 0.5_489, 0.5_845], [0.5_563, 0.7_580, 0.8_518]]		)

 if single_scale:
  UpperCAmelCase__								=  torch.tensor(
      [[-9.9_051, -4.2_541, -6.4_852], [-9.6_947, -4.0_854, -6.8_033], [-10.0665, -5.8_470, -7.7_003]]		)
  UpperCAmelCase__								=  torch.tensor([[0.7_292, 0.4_991, 0.5_532], [0.7_959, 0.2_426, 0.4_236], [0.7_582, 0.3_518, 0.4_451]]		)

 if single_scale and dilation:
  UpperCAmelCase__								=  torch.tensor(
      [[-8.9_652, -4.1_074, -5.6_635], [-9.0_596, -4.9_447, -6.6_075], [-10.1178, -4.5_275, -6.2_671]]		)
  UpperCAmelCase__								=  torch.tensor([[0.7_665, 0.4_130, 0.4_769], [0.8_364, 0.1_841, 0.3_391], [0.6_261, 0.3_895, 0.7_978]]		)

 if with_box_refine:
  UpperCAmelCase__								=  torch.tensor(
      [[-8.8_895, -5.4_187, -6.8_153], [-8.4_706, -6.1_668, -7.6_184], [-9.0_042, -5.5_359, -6.9_141]]		)
  UpperCAmelCase__								=  torch.tensor([[0.7_828, 0.2_208, 0.4_323], [0.0_892, 0.5_996, 0.1_319], [0.5_524, 0.6_389, 0.8_914]]		)

 if with_box_refine and two_stage:
  UpperCAmelCase__								=  torch.tensor(
      [[-6.7_108, -4.3_213, -6.3_777], [-8.9_014, -6.1_799, -6.7_240], [-6.9_315, -4.4_735, -6.2_298]]		)
  UpperCAmelCase__								=  torch.tensor([[0.2_583, 0.5_499, 0.4_683], [0.7_652, 0.9_068, 0.4_882], [0.5_490, 0.2_763, 0.0_564]]		)

 print("Logits:"  ,							outputs.logits[0, :3, :3]		)

 assert torch.allclose(outputs.logits[0, :3, :3]  ,							expected_logits.to(lowercase_		)  ,							atol=1E-4		)
 assert torch.allclose(outputs.pred_boxes[0, :3, :3]  ,							expected_boxes.to(lowercase_		)  ,							atol=1E-4		)

 print("Everything ok!"		)

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

 # Push to hub
 if push_to_hub:
  UpperCAmelCase__								=  "deformable-detr"
  model_name += "-single-scale" if single_scale else ""
  model_name += "-dc5" if dilation else ""
  model_name += "-with-box-refine" if with_box_refine else ""
  model_name += "-two-stage" if two_stage else ""
  print("Pushing model to hub..."		)
  model.push_to_hub(repo_path_or_name=lowercase_  ,							organization="nielsr"  ,							commit_message="Add model"		)


if __name__ == "__main__":
  _lowerCAmelCase  :				Union[str, Any]						=				argparse.ArgumentParser()

  parser.add_argument(
      "--checkpoint_path",
      type=str,
      default="/home/niels/checkpoints/deformable_detr/r50_deformable_detr-checkpoint.pth",
      help="Path to Pytorch checkpoint (.pth file) you'd like to convert.",
  )
  parser.add_argument("--single_scale", action="store_true", help="Whether to set config.num_features_levels = 1.")
  parser.add_argument("--dilation", action="store_true", help="Whether to set config.dilation=True.")
  parser.add_argument("--with_box_refine", action="store_true", help="Whether to set config.with_box_refine=True.")
  parser.add_argument("--two_stage", action="store_true", help="Whether to set config.two_stage=True.")
  parser.add_argument(
      "--pytorch_dump_folder_path",
      default=None,
      type=str,
      required=True,
      help="Path to the folder to output PyTorch model.",
  )
  parser.add_argument(
      "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
  )
  _lowerCAmelCase  :				Optional[Any]						=				parser.parse_args()
  convert_deformable_detr_checkpoint(
      args.checkpoint_path,
      args.single_scale,
      args.dilation,
      args.with_box_refine,
      args.two_stage,
      args.pytorch_dump_folder_path,
      args.push_to_hub,
  )











import unittest

from transformers import YosoConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      YosoForMaskedLM,
      YosoForMultipleChoice,
      YosoForQuestionAnswering,
      YosoForSequenceClassification,
      YosoForTokenClassification,
      YosoModel,
  )
  from transformers.models.yoso.modeling_yoso import YOSO_PRETRAINED_MODEL_ARCHIVE_LIST
class      _UpperCamelCase :
      def __init__(     self				:int	, lowerCamelCase				:int	, lowerCamelCase				:Optional[Any]=13	, lowerCamelCase				:Optional[Any]=7	, lowerCamelCase				:Union[str, Any]=True	, lowerCamelCase				:Tuple=True	, lowerCamelCase				:Optional[Any]=True	, lowerCamelCase				:str=True	, lowerCamelCase				:Any=99	, lowerCamelCase				:Optional[int]=32	, lowerCamelCase				:Any=5	, lowerCamelCase				:List[Any]=4	, lowerCamelCase				:Optional[Any]=37	, lowerCamelCase				:Tuple="gelu"	, lowerCamelCase				:Optional[int]=0.1	, lowerCamelCase				:Any=0.1	, lowerCamelCase				:List[Any]=512	, lowerCamelCase				:Tuple=16	, lowerCamelCase				:Optional[int]=2	, lowerCamelCase				:str=0.02	, lowerCamelCase				:Dict=3	, lowerCamelCase				:Union[str, Any]=4	, lowerCamelCase				:Union[str, Any]=None	, )       ->     List[Any]:
       UpperCAmelCase__								=  parent
       UpperCAmelCase__								=  batch_size
       UpperCAmelCase__								=  seq_length
       UpperCAmelCase__								=  is_training
       UpperCAmelCase__								=  use_input_mask
       UpperCAmelCase__								=  use_token_type_ids
       UpperCAmelCase__								=  use_labels
       UpperCAmelCase__								=  vocab_size
       UpperCAmelCase__								=  hidden_size
       UpperCAmelCase__								=  num_hidden_layers
       UpperCAmelCase__								=  num_attention_heads
       UpperCAmelCase__								=  intermediate_size
       UpperCAmelCase__								=  hidden_act
       UpperCAmelCase__								=  hidden_dropout_prob
       UpperCAmelCase__								=  attention_probs_dropout_prob
       UpperCAmelCase__								=  max_position_embeddings
       UpperCAmelCase__								=  type_vocab_size
       UpperCAmelCase__								=  type_sequence_label_size
       UpperCAmelCase__								=  initializer_range
       UpperCAmelCase__								=  num_labels
       UpperCAmelCase__								=  num_choices
       UpperCAmelCase__								=  scope
      def 		UpperCAmelCase_  (     self				:Union[str, Any] )       ->     Any:
       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__								=  self.get_config()

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
      def 		UpperCAmelCase_  (     self				:List[Any] )       ->     Optional[Any]:
       return YosoConfig(
           vocab_size=self.vocab_size	, hidden_size=self.hidden_size	, num_hidden_layers=self.num_hidden_layers	, num_attention_heads=self.num_attention_heads	, intermediate_size=self.intermediate_size	, hidden_act=self.hidden_act	, hidden_dropout_prob=self.hidden_dropout_prob	, attention_probs_dropout_prob=self.attention_probs_dropout_prob	, max_position_embeddings=self.max_position_embeddings	, type_vocab_size=self.type_vocab_size	, is_decoder=lowerCamelCase	, initializer_range=self.initializer_range	, )
      def 		UpperCAmelCase_  (     self				:int )       ->     Any:
       UpperCAmelCase__								=  self.get_config()
       UpperCAmelCase__								=  300
       return config
      def 		UpperCAmelCase_  (     self				:List[str] )       ->     Tuple:
       (
           (
           UpperCAmelCase__
       )	,	(
           UpperCAmelCase__
       )	,	(
           UpperCAmelCase__
       )	,	(
           UpperCAmelCase__
       )	,	(
           UpperCAmelCase__
       )	,	(
           UpperCAmelCase__
       )	,	(
           UpperCAmelCase__
       )	,	
       )								=  self.prepare_config_and_inputs()

       UpperCAmelCase__								=  True
       UpperCAmelCase__								=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
       UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]	, vocab_size=2 )

       return (
           config,
           input_ids,
           token_type_ids,
           input_mask,
           sequence_labels,
           token_labels,
           choice_labels,
           encoder_hidden_states,
           encoder_attention_mask,
       )
      def 		UpperCAmelCase_  (     self				:Any	, lowerCamelCase				:List[str]	, lowerCamelCase				:List[str]	, lowerCamelCase				:str	, lowerCamelCase				:Union[str, Any]	, lowerCamelCase				:Optional[int]	, lowerCamelCase				:Any	, lowerCamelCase				:Tuple )       ->     List[Any]:
       UpperCAmelCase__								=  YosoModel(config=lowerCamelCase )
       model.to(lowerCamelCase )
       model.eval()
       UpperCAmelCase__								=  model(lowerCamelCase	, attention_mask=lowerCamelCase	, token_type_ids=lowerCamelCase )
       UpperCAmelCase__								=  model(lowerCamelCase	, token_type_ids=lowerCamelCase )
       UpperCAmelCase__								=  model(lowerCamelCase )
       self.parent.assertEqual(result.last_hidden_state.shape	, (self.batch_size, self.seq_length, self.hidden_size) )
      def 		UpperCAmelCase_  (     self				:Any	, lowerCamelCase				:str	, lowerCamelCase				:Dict	, lowerCamelCase				:Any	, lowerCamelCase				:Any	, lowerCamelCase				:int	, lowerCamelCase				:Dict	, lowerCamelCase				:Any	, lowerCamelCase				:Union[str, Any]	, lowerCamelCase				:List[Any]	, )       ->     Optional[int]:
       UpperCAmelCase__								=  True
       UpperCAmelCase__								=  YosoModel(lowerCamelCase )
       model.to(lowerCamelCase )
       model.eval()
       UpperCAmelCase__								=  model(
           lowerCamelCase	, attention_mask=lowerCamelCase	, token_type_ids=lowerCamelCase	, encoder_hidden_states=lowerCamelCase	, encoder_attention_mask=lowerCamelCase	, )
       UpperCAmelCase__								=  model(
           lowerCamelCase	, attention_mask=lowerCamelCase	, token_type_ids=lowerCamelCase	, encoder_hidden_states=lowerCamelCase	, )
       UpperCAmelCase__								=  model(lowerCamelCase	, attention_mask=lowerCamelCase	, token_type_ids=lowerCamelCase )
       self.parent.assertEqual(result.last_hidden_state.shape	, (self.batch_size, self.seq_length, self.hidden_size) )
      def 		UpperCAmelCase_  (     self				:Tuple	, lowerCamelCase				:List[Any]	, lowerCamelCase				:Tuple	, lowerCamelCase				:List[str]	, lowerCamelCase				:Optional[Any]	, lowerCamelCase				:Tuple	, lowerCamelCase				:Optional[Any]	, lowerCamelCase				:int )       ->     Optional[int]:
       UpperCAmelCase__								=  YosoForMaskedLM(config=lowerCamelCase )
       model.to(lowerCamelCase )
       model.eval()
       UpperCAmelCase__								=  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 		UpperCAmelCase_  (     self				:Dict	, lowerCamelCase				:str	, lowerCamelCase				:Dict	, lowerCamelCase				:List[Any]	, lowerCamelCase				:Optional[int]	, lowerCamelCase				:Union[str, Any]	, lowerCamelCase				:int	, lowerCamelCase				:Optional[Any] )       ->     Any:
       UpperCAmelCase__								=  YosoForQuestionAnswering(config=lowerCamelCase )
       model.to(lowerCamelCase )
       model.eval()
       UpperCAmelCase__								=  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 		UpperCAmelCase_  (     self				:Optional[Any]	, lowerCamelCase				:Optional[Any]	, lowerCamelCase				:Any	, lowerCamelCase				:List[str]	, lowerCamelCase				:Dict	, lowerCamelCase				:str	, lowerCamelCase				:Tuple	, lowerCamelCase				:int )       ->     Dict:
       UpperCAmelCase__								=  self.num_labels
       UpperCAmelCase__								=  YosoForSequenceClassification(lowerCamelCase )
       model.to(lowerCamelCase )
       model.eval()
       UpperCAmelCase__								=  model(lowerCamelCase	, attention_mask=lowerCamelCase	, token_type_ids=lowerCamelCase	, labels=lowerCamelCase )
       self.parent.assertEqual(result.logits.shape	, (self.batch_size, self.num_labels) )
      def 		UpperCAmelCase_  (     self				:Tuple	, lowerCamelCase				:str	, lowerCamelCase				:int	, lowerCamelCase				:Optional[Any]	, lowerCamelCase				:List[str]	, lowerCamelCase				:Dict	, lowerCamelCase				:Optional[int]	, lowerCamelCase				:str )       ->     Dict:
       UpperCAmelCase__								=  self.num_labels
       UpperCAmelCase__								=  YosoForTokenClassification(config=lowerCamelCase )
       model.to(lowerCamelCase )
       model.eval()
       UpperCAmelCase__								=  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 		UpperCAmelCase_  (     self				:Optional[int]	, lowerCamelCase				:Tuple	, lowerCamelCase				:Tuple	, lowerCamelCase				:int	, lowerCamelCase				:Tuple	, lowerCamelCase				:List[str]	, lowerCamelCase				:Optional[Any]	, lowerCamelCase				:List[str] )       ->     List[Any]:
       UpperCAmelCase__								=  self.num_choices
       UpperCAmelCase__								=  YosoForMultipleChoice(config=lowerCamelCase )
       model.to(lowerCamelCase )
       model.eval()
       UpperCAmelCase__								=  input_ids.unsqueeze(1 ).expand(-1	, self.num_choices	, -1 ).contiguous()
       UpperCAmelCase__								=  token_type_ids.unsqueeze(1 ).expand(-1	, self.num_choices	, -1 ).contiguous()
       UpperCAmelCase__								=  input_mask.unsqueeze(1 ).expand(-1	, self.num_choices	, -1 ).contiguous()
       UpperCAmelCase__								=  model(
           lowerCamelCase	, attention_mask=lowerCamelCase	, token_type_ids=lowerCamelCase	, labels=lowerCamelCase	, )
       self.parent.assertEqual(result.logits.shape	, (self.batch_size, self.num_choices) )







      def 		UpperCAmelCase_  (     self				:int )       ->     List[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_torch
class      _UpperCamelCase (			lowerCAmelCase		,    lowerCAmelCase		,    unittest.TestCase						):
      UpperCAmelCase_            =						(
          (
              YosoModel,
              YosoForMaskedLM,
              YosoForMultipleChoice,
              YosoForQuestionAnswering,
              YosoForSequenceClassification,
              YosoForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      UpperCAmelCase_            =						False
      UpperCAmelCase_            =						False
      UpperCAmelCase_            =						False

      UpperCAmelCase_            =						()
      UpperCAmelCase_            =						(
          {
              """feature-extraction""": YosoModel,
              """fill-mask""": YosoForMaskedLM,
              """question-answering""": YosoForQuestionAnswering,
              """text-classification""": YosoForSequenceClassification,
              """token-classification""": YosoForTokenClassification,
              """zero-shot""": YosoForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      def 		UpperCAmelCase_  (     self				:Any )       ->     Optional[Any]:
       UpperCAmelCase__								=  YosoModelTester(self )
       UpperCAmelCase__								=  ConfigTester(self	, config_class=lowerCamelCase	, hidden_size=37 )
      def 		UpperCAmelCase_  (     self				:List[Any] )       ->     List[str]:
       self.config_tester.run_common_tests()
      def 		UpperCAmelCase_  (     self				:Optional[int] )       ->     List[Any]:
       UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*lowerCamelCase )
      def 		UpperCAmelCase_  (     self				:Tuple )       ->     Optional[Any]:
       UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
        UpperCAmelCase__								=  type
        self.model_tester.create_and_check_model(*lowerCamelCase )
      def 		UpperCAmelCase_  (     self				:Tuple )       ->     Optional[int]:
       UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase )
      def 		UpperCAmelCase_  (     self				:Union[str, Any] )       ->     Dict:
       UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*lowerCamelCase )
      def 		UpperCAmelCase_  (     self				:str )       ->     Optional[Any]:
       UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*lowerCamelCase )
      def 		UpperCAmelCase_  (     self				:Any )       ->     List[str]:
       UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase )
      def 		UpperCAmelCase_  (     self				:Any )       ->     str:
       UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*lowerCamelCase )
      @slow
      def 		UpperCAmelCase_  (     self				:List[Any] )       ->     Optional[int]:
       for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase__								=  YosoModel.from_pretrained(lowerCamelCase )
        self.assertIsNotNone(lowerCamelCase )







      def 		UpperCAmelCase_  (     self				:int )       ->     List[str]:
       return



@require_torch
class      _UpperCamelCase (			unittest.TestCase						):
      @slow
      def 		UpperCAmelCase_  (     self				:Optional[Any] )       ->     Optional[int]:
       UpperCAmelCase__								=  YosoModel.from_pretrained("uw-madison/yoso-4096" )
       UpperCAmelCase__								=  torch.tensor([[0, 1, 2, 3, 4, 5]] )

       with torch.no_grad():
        UpperCAmelCase__								=  model(lowerCamelCase )[0]

       UpperCAmelCase__								=  torch.Size((1, 6, 768) )
       self.assertEqual(output.shape	, lowerCamelCase )

       UpperCAmelCase__								=  torch.tensor(
           [[[-0.06_11, 0.12_42, 0.08_40], [0.02_80, -0.00_48, 0.11_25], [0.01_06, 0.02_26, 0.07_51]]] )

       self.assertTrue(torch.allclose(output[:, :3, :3]	, lowerCamelCase	, atol=1e-4 ) )
      @slow
      def 		UpperCAmelCase_  (     self				:Optional[Any] )       ->     List[Any]:
       UpperCAmelCase__								=  YosoForMaskedLM.from_pretrained("uw-madison/yoso-4096" )
       UpperCAmelCase__								=  torch.tensor([[0, 1, 2, 3, 4, 5]] )

       with torch.no_grad():
        UpperCAmelCase__								=  model(lowerCamelCase )[0]

       UpperCAmelCase__								=  5_0265

       UpperCAmelCase__								=  torch.Size((1, 6, vocab_size) )
       self.assertEqual(output.shape	, lowerCamelCase )

       UpperCAmelCase__								=  torch.tensor(
           [[[-2.13_13, -3.72_85, -2.24_07], [-2.70_47, -3.33_14, -2.64_08], [0.06_29, -2.51_66, -0.33_56]]] )

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







      @slow
      def 		UpperCAmelCase_  (     self				:Any )       ->     str:
       UpperCAmelCase__								=  YosoForMaskedLM.from_pretrained("uw-madison/yoso-4096" )
       UpperCAmelCase__								=  torch.arange(4096 ).unsqueeze(0 )

       with torch.no_grad():
        UpperCAmelCase__								=  model(lowerCamelCase )[0]

       UpperCAmelCase__								=  5_0265

       UpperCAmelCase__								=  torch.Size((1, 4096, vocab_size) )
       self.assertEqual(output.shape	, lowerCamelCase )

       UpperCAmelCase__								=  torch.tensor(
           [[[-2.39_14, -4.37_42, -5.09_56], [-4.09_88, -4.23_84, -7.04_06], [-3.14_27, -3.71_92, -6.68_00]]] )

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





