





import argparse
import json
from pathlib import Path

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

from transformers import YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


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

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :str     )      ->						YolosConfig:
  _lowercase							= YolosConfig()

  # size of the architecture
  if "yolos_ti" in yolos_name:
    _lowercase							= 192
    _lowercase							= 768
    _lowercase							= 12
    _lowercase							= 3
    _lowercase							= [800, 1333]
    _lowercase							= False
  elif yolos_name == "yolos_s_dWr":
    _lowercase							= 330
    _lowercase							= 14
    _lowercase							= 6
    _lowercase							= 1320
  elif "yolos_s" in yolos_name:
    _lowercase							= 384
    _lowercase							= 1536
    _lowercase							= 12
    _lowercase							= 6
  elif "yolos_b" in yolos_name:
    _lowercase							= [800, 1344]

  _lowercase							= 91
  _lowercase							= 'huggingface/label-files'
  _lowercase							= 'coco-detection-id2label.json'
  _lowercase							= json.load(open(hf_hub_download(snake_case__			,							snake_case__			,							repo_type='dataset'     )			,							'r'     )     )
  _lowercase							= {int(snake_case__     ): v for k, v in idalabel.items()}
  _lowercase							= idalabel
  _lowercase							= {v: k for k, v in idalabel.items()}

  return config

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :dict			,							snake_case__  :YolosConfig			,							snake_case__  :bool = False     )      ->						Tuple:
  for i in range(config.num_hidden_layers     ):
    # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
    _lowercase							= state_dict.pop(F"""blocks.{i}.attn.qkv.weight"""     )
    _lowercase							= state_dict.pop(F"""blocks.{i}.attn.qkv.bias"""     )
    # next, add query, keys and values (in that order) to the state dict
    _lowercase							= in_proj_weight[: config.hidden_size, :]
    _lowercase							= in_proj_bias[: config.hidden_size]
    _lowercase							= in_proj_weight[
        config.hidden_size : config.hidden_size * 2, :
    ]
    _lowercase							= in_proj_bias[
        config.hidden_size : config.hidden_size * 2
    ]
    _lowercase							= in_proj_weight[-config.hidden_size :, :]
    _lowercase							= in_proj_bias[-config.hidden_size :]

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :str     )      ->						str:
  if "backbone" in name:
    _lowercase							= name.replace('backbone'			,							'vit'     )
  if "cls_token" in name:
    _lowercase							= name.replace('cls_token'			,							'embeddings.cls_token'     )
  if "det_token" in name:
    _lowercase							= name.replace('det_token'			,							'embeddings.detection_tokens'     )
  if "mid_pos_embed" in name:
    _lowercase							= name.replace('mid_pos_embed'			,							'encoder.mid_position_embeddings'     )
  if "pos_embed" in name:
    _lowercase							= name.replace('pos_embed'			,							'embeddings.position_embeddings'     )
  if "patch_embed.proj" in name:
    _lowercase							= name.replace('patch_embed.proj'			,							'embeddings.patch_embeddings.projection'     )
  if "blocks" in name:
    _lowercase							= name.replace('blocks'			,							'encoder.layer'     )
  if "attn.proj" in name:
    _lowercase							= name.replace('attn.proj'			,							'attention.output.dense'     )
  if "attn" in name:
    _lowercase							= name.replace('attn'			,							'attention.self'     )
  if "norm1" in name:
    _lowercase							= name.replace('norm1'			,							'layernorm_before'     )
  if "norm2" in name:
    _lowercase							= name.replace('norm2'			,							'layernorm_after'     )
  if "mlp.fc1" in name:
    _lowercase							= name.replace('mlp.fc1'			,							'intermediate.dense'     )
  if "mlp.fc2" in name:
    _lowercase							= name.replace('mlp.fc2'			,							'output.dense'     )
  if "class_embed" in name:
    _lowercase							= name.replace('class_embed'			,							'class_labels_classifier'     )
  if "bbox_embed" in name:
    _lowercase							= name.replace('bbox_embed'			,							'bbox_predictor'     )
  if "vit.norm" in name:
    _lowercase							= name.replace('vit.norm'			,							'vit.layernorm'     )

  return name

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :dict			,							snake_case__  :YolosForObjectDetection     )      ->						dict:
  for key in orig_state_dict.copy().keys():
    _lowercase							= orig_state_dict.pop(snake_case__     )

    if "qkv" in key:
      _lowercase							= key.split('.'     )
      _lowercase							= int(key_split[2]     )
      _lowercase							= model.vit.encoder.layer[layer_num].attention.attention.all_head_size
      if "weight" in key:
        _lowercase							= val[:dim, :]
        _lowercase							= val[
            dim : dim * 2, :
        ]
        _lowercase							= val[-dim:, :]
      else:
        _lowercase							= val[:dim]
        _lowercase							= val[dim : dim * 2]
        _lowercase							= val[-dim:]
    else:
      _lowercase							= val

  return orig_state_dict

def 							SCREAMING_SNAKE_CASE__       (  )      ->						torch.Tensor:
  _lowercase							= 'http://images.cocodataset.org/val2017/000000039769.jpg'
  _lowercase							= Image.open(requests.get(snake_case__			,							stream=snake_case__     ).raw     )
  return im

@torch.no_grad()
def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :str			,							snake_case__  :str			,							snake_case__  :str			,							snake_case__  :bool = False     )      ->						List[Any]:
  _lowercase							= get_yolos_config(snake_case__     )

  # load original state_dict
  _lowercase							= torch.load(snake_case__			,							map_location='cpu'     )['model']

  # load 🤗 model
  _lowercase							= YolosForObjectDetection(snake_case__     )
  model.eval()
  _lowercase							= convert_state_dict(snake_case__			,							snake_case__     )
  model.load_state_dict(snake_case__     )

  # Check outputs on an image, prepared by YolosImageProcessor
  _lowercase							= 800 if yolos_name != 'yolos_ti' else 512
  _lowercase							= YolosImageProcessor(format='coco_detection'			,							size=snake_case__     )
  _lowercase							= image_processor(images=prepare_img()			,							return_tensors='pt'     )
  _lowercase							= model(**snake_case__     )
  _lowercase						,		_lowercase							= outputs.logits, outputs.pred_boxes

  _lowercase						,		_lowercase							= None, None
  if yolos_name == "yolos_ti":
    _lowercase							= torch.tensor(
        [[-39.5022, -11.9820, -17.6888], [-29.9574, -9.9769, -17.7691], [-42.3281, -20.7200, -30.6294]]     )
    _lowercase							= torch.tensor(
        [[0.4021, 0.0836, 0.7979], [0.0184, 0.2609, 0.0364], [0.1781, 0.2004, 0.2095]]     )
  elif yolos_name == "yolos_s_200_pre":
    _lowercase							= torch.tensor(
        [[-24.0248, -10.3024, -14.8290], [-42.0392, -16.8200, -27.4334], [-27.2743, -11.8154, -18.7148]]     )
    _lowercase							= torch.tensor(
        [[0.2559, 0.5455, 0.4706], [0.2989, 0.7279, 0.1875], [0.7732, 0.4017, 0.4462]]     )
  elif yolos_name == "yolos_s_300_pre":
    _lowercase							= torch.tensor(
        [[-36.2220, -14.4385, -23.5457], [-35.6970, -14.7583, -21.3935], [-31.5939, -13.6042, -16.8049]]     )
    _lowercase							= torch.tensor(
        [[0.7614, 0.2316, 0.4728], [0.7168, 0.4495, 0.3855], [0.4996, 0.1466, 0.9996]]     )
  elif yolos_name == "yolos_s_dWr":
    _lowercase							= torch.tensor(
        [[-42.8668, -24.1049, -41.1690], [-34.7456, -14.1274, -24.9194], [-33.7898, -12.1946, -25.6495]]     )
    _lowercase							= torch.tensor(
        [[0.5587, 0.2773, 0.0605], [0.5004, 0.3014, 0.9994], [0.4999, 0.1548, 0.9994]]     )
  elif yolos_name == "yolos_base":
    _lowercase							= torch.tensor(
        [[-40.6064, -24.3084, -32.6447], [-55.1990, -30.7719, -35.5877], [-51.4311, -33.3507, -35.6462]]     )
    _lowercase							= torch.tensor(
        [[0.5555, 0.2794, 0.0655], [0.9049, 0.2664, 0.1894], [0.9183, 0.1984, 0.1635]]     )
  else:
    raise ValueError(F"""Unknown yolos_name: {yolos_name}"""     )

  assert torch.allclose(logits[0, :3, :3]			,							snake_case__			,							atol=1E-4     )
  assert torch.allclose(pred_boxes[0, :3, :3]			,							snake_case__			,							atol=1E-4     )

  Path(snake_case__     ).mkdir(exist_ok=snake_case__     )
  print(F"""Saving model {yolos_name} to {pytorch_dump_folder_path}"""     )
  model.save_pretrained(snake_case__     )
  print(F"""Saving image processor to {pytorch_dump_folder_path}"""     )
  image_processor.save_pretrained(snake_case__     )

  if push_to_hub:
    _lowercase							= {
        'yolos_ti': 'yolos-tiny',
        'yolos_s_200_pre': 'yolos-small',
        'yolos_s_300_pre': 'yolos-small-300',
        'yolos_s_dWr': 'yolos-small-dwr',
        'yolos_base': 'yolos-base',
    }

    print('Pushing to the hub...'     )
    _lowercase							= model_mapping[yolos_name]
    image_processor.push_to_hub(snake_case__			,							organization='hustvl'     )
    model.push_to_hub(snake_case__			,							organization='hustvl'     )


if __name__ == "__main__":
      snake_case 				=       argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          """--yolos_name""",
          default="""yolos_s_200_pre""",
          type=str,
          help=(
              """Name of the YOLOS model you'd like to convert. Should be one of 'yolos_ti', 'yolos_s_200_pre',"""
              """ 'yolos_s_300_pre', 'yolos_s_dWr', 'yolos_base'."""
          ),
      )
      parser.add_argument(
          """--checkpoint_path""", default=None, type=str, help="""Path to the original state dict (.pth file)."""
      )
      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 or not to push the converted model to the 🤗 hub."""
      )

      snake_case 				=       parser.parse_args()
      convert_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)






from __future__ import annotations

import unittest

from transformers import 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 numpy
      import tensorflow as tf

      from transformers import (
          TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
          TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
          TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
          BertConfig,
          DPRConfig,
          TFDPRContextEncoder,
          TFDPRQuestionEncoder,
          TFDPRReader,
      )





class 				A_				:







       """simple docstring"""





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







       def    __UpperCAmelCase (    self	:	List[Any]  )						->					List[str]:
         _lowercase							= ids_tensor([self.batch_size, self.seq_length]						,self.vocab_size  )

         _lowercase							= None
         if self.use_input_mask:
           # follow test_modeling_tf_ctrl.py
           _lowercase							= random_attention_mask([self.batch_size, self.seq_length]  )

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

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

         _lowercase							= BertConfig(
             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=__A						,initializer_range=self.initializer_range						,)
         _lowercase							= DPRConfig(projection_dim=self.projection_dim						,**config.to_dict()  )

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







       def    __UpperCAmelCase (    self	:	List[Any]						,__A	:	Dict						,__A	:	str						,__A	:	Any						,__A	:	Dict						,__A	:	Tuple						,__A	:	Tuple						,__A	:	int  )						->					List[str]:
         _lowercase							= TFDPRContextEncoder(config=__A  )
         _lowercase							= model(__A						,attention_mask=__A						,token_type_ids=__A  )
         _lowercase							= model(__A						,token_type_ids=__A  )
         _lowercase							= model(__A  )
         self.parent.assertEqual(result.pooler_output.shape						,(self.batch_size, self.projection_dim or self.hidden_size)  )







       def    __UpperCAmelCase (    self	:	List[Any]						,__A	:	Optional[Any]						,__A	:	Tuple						,__A	:	Optional[Any]						,__A	:	Union[str, Any]						,__A	:	Tuple						,__A	:	Tuple						,__A	:	Any  )						->					str:
         _lowercase							= TFDPRQuestionEncoder(config=__A  )
         _lowercase							= model(__A						,attention_mask=__A						,token_type_ids=__A  )
         _lowercase							= model(__A						,token_type_ids=__A  )
         _lowercase							= model(__A  )
         self.parent.assertEqual(result.pooler_output.shape						,(self.batch_size, self.projection_dim or self.hidden_size)  )







       def    __UpperCAmelCase (    self	:	int						,__A	:	List[str]						,__A	:	Optional[int]						,__A	:	List[str]						,__A	:	str						,__A	:	str						,__A	:	int						,__A	:	int  )						->					Optional[Any]:
         _lowercase							= TFDPRReader(config=__A  )
         _lowercase							= model(__A						,attention_mask=__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)  )
         self.parent.assertEqual(result.relevance_logits.shape						,(self.batch_size,)  )







       def    __UpperCAmelCase (    self	:	str  )						->					str:
         _lowercase							= self.prepare_config_and_inputs()
         (
             (
             _lowercase
         )						,		(
             _lowercase
         )						,		(
             _lowercase
         )						,		(
             _lowercase
         )						,		(
             _lowercase
         )						,		(
             _lowercase
         )						,		(
             _lowercase
         )						,		
         )							= config_and_inputs
         _lowercase							= {'input_ids': input_ids}
         return config, inputs_dict





@require_tf
class 				A_				(      UpperCAmelCase					,	UpperCAmelCase					,	unittest.TestCase	):







       """simple docstring"""





       SCREAMING_SNAKE_CASE_      : Union[str, Any]      =  (
           (
               TFDPRContextEncoder,
               TFDPRQuestionEncoder,
               TFDPRReader,
           )
           if is_tf_available()
           else ()
       )
       SCREAMING_SNAKE_CASE_      : Optional[Any]      =  {'''feature-extraction''': TFDPRQuestionEncoder} if is_tf_available() else {}

       SCREAMING_SNAKE_CASE_      : List[Any]      =  False
       SCREAMING_SNAKE_CASE_      : Dict      =  False
       SCREAMING_SNAKE_CASE_      : int      =  False
       SCREAMING_SNAKE_CASE_      : List[Any]      =  False
       SCREAMING_SNAKE_CASE_      : Optional[int]      =  False







       def    __UpperCAmelCase (    self	:	List[str]  )						->					Optional[int]:
         _lowercase							= TFDPRModelTester(self  )
         _lowercase							= ConfigTester(self						,config_class=__A						,hidden_size=37  )







       def    __UpperCAmelCase (    self	:	Any  )						->					List[str]:
         self.config_tester.run_common_tests()







       def    __UpperCAmelCase (    self	:	Dict  )						->					Optional[int]:
         _lowercase							= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_context_encoder(*__A  )







       def    __UpperCAmelCase (    self	:	Optional[Any]  )						->					Optional[Any]:
         _lowercase							= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_question_encoder(*__A  )







       def    __UpperCAmelCase (    self	:	Dict  )						->					List[Any]:
         _lowercase							= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_reader(*__A  )







       @slow
       def    __UpperCAmelCase (    self	:	Optional[Any]  )						->					str:
         for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _lowercase							= TFDPRContextEncoder.from_pretrained(__A  )
           self.assertIsNotNone(__A  )

         for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _lowercase							= TFDPRContextEncoder.from_pretrained(__A  )
           self.assertIsNotNone(__A  )

         for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _lowercase							= TFDPRQuestionEncoder.from_pretrained(__A  )
           self.assertIsNotNone(__A  )

         for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _lowercase							= TFDPRReader.from_pretrained(__A  )
           self.assertIsNotNone(__A  )





@require_tf
class 				A_				(      unittest.TestCase	):







       """simple docstring"""





       @slow
       def    __UpperCAmelCase (    self	:	List[Any]  )						->					int:
         _lowercase							= TFDPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base'  )

         _lowercase							= tf.constant(
             [[101, 7592, 1010, 2003, 2026, 3899, 1_0140, 1029, 102]]  )  # [CLS] hello, is my dog cute? [SEP]
         _lowercase							= model(__A  )[0]  # embedding shape = (1, 768)
         # compare the actual values for a slice.
         _lowercase							= tf.constant(
             [
                 [
                     0.03236253,
                     0.12753335,
                     0.16818509,
                     0.00279786,
                     0.3896933,
                     0.24264945,
                     0.2178971,
                     -0.02335227,
                     -0.08481959,
                     -0.14324117,
                 ]
             ]  )
         self.assertTrue(numpy.allclose(output[:, :10].numpy()						,expected_slice.numpy()						,atol=1e-4  )  )