



import argparse

import requests
import torch
from PIL import Image

from transformers import CLIPProcessor, GroupViTConfig, GroupViTModel



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	Union[str, Any]				)  ->     Optional[Any]:
 # vision encoder
 if "img_encoder.pos_embed" in name:
  _snake_case    :  Any					=					name.replace("""img_encoder.pos_embed"""				,       """vision_model.embeddings.position_embeddings"""				)
 if "img_encoder.patch_embed.proj" in name:
  _snake_case    :  Union[str, Any]					=					name.replace("""img_encoder.patch_embed.proj"""				,       """vision_model.embeddings.patch_embeddings.projection"""				)
 if "img_encoder.patch_embed.norm" in name:
  _snake_case    :  Tuple					=					name.replace("""img_encoder.patch_embed.norm"""				,       """vision_model.embeddings.layernorm"""				)
 if "img_encoder.layers" in name:
  _snake_case    :  Optional[Any]					=					name.replace("""img_encoder.layers"""				,       """vision_model.encoder.stages"""				)
 if "blocks" in name and "res" not in name:
  _snake_case    :  Union[str, Any]					=					name.replace("""blocks"""				,       """layers"""				)
 if "attn" in name and "pre_assign" not in name:
  _snake_case    :  Any					=					name.replace("""attn"""				,       """self_attn"""				)
 if "proj" in name and "self_attn" in name and "text" not in name:
  _snake_case    :  int					=					name.replace("""proj"""				,       """out_proj"""				)
 if "pre_assign_attn.attn.proj" in name:
  _snake_case    :  Union[str, Any]					=					name.replace("""pre_assign_attn.attn.proj"""				,       """pre_assign_attn.attn.out_proj"""				)
 if "norm1" in name:
  _snake_case    :  Any					=					name.replace("""norm1"""				,       """layer_norm1"""				)
 if "norm2" in name and "pre_assign" not in name:
  _snake_case    :  Dict					=					name.replace("""norm2"""				,       """layer_norm2"""				)
 if "img_encoder.norm" in name:
  _snake_case    :  Union[str, Any]					=					name.replace("""img_encoder.norm"""				,       """vision_model.layernorm"""				)
 # text encoder
 if "text_encoder.token_embedding" in name:
  _snake_case    :  Dict					=					name.replace("""text_encoder.token_embedding"""				,       """text_model.embeddings.token_embedding"""				)
 if "text_encoder.positional_embedding" in name:
  _snake_case    :  str					=					name.replace("""text_encoder.positional_embedding"""				,       """text_model.embeddings.position_embedding.weight"""				)
 if "text_encoder.transformer.resblocks." in name:
  _snake_case    :  int					=					name.replace("""text_encoder.transformer.resblocks."""				,       """text_model.encoder.layers."""				)
 if "ln_1" in name:
  _snake_case    :  Union[str, Any]					=					name.replace("""ln_1"""				,       """layer_norm1"""				)
 if "ln_2" in name:
  _snake_case    :  Optional[Any]					=					name.replace("""ln_2"""				,       """layer_norm2"""				)
 if "c_fc" in name:
  _snake_case    :  Union[str, Any]					=					name.replace("""c_fc"""				,       """fc1"""				)
 if "c_proj" in name:
  _snake_case    :  Tuple					=					name.replace("""c_proj"""				,       """fc2"""				)
 if "text_encoder" in name:
  _snake_case    :  Any					=					name.replace("""text_encoder"""				,       """text_model"""				)
 if "ln_final" in name:
  _snake_case    :  str					=					name.replace("""ln_final"""				,       """final_layer_norm"""				)
 # projection layers
 if "img_projector.linear_hidden." in name:
  _snake_case    :  List[Any]					=					name.replace("""img_projector.linear_hidden."""				,       """visual_projection."""				)
 if "img_projector.linear_out." in name:
  _snake_case    :  Optional[int]					=					name.replace("""img_projector.linear_out."""				,       """visual_projection.3."""				)
 if "text_projector.linear_hidden" in name:
  _snake_case    :  List[Any]					=					name.replace("""text_projector.linear_hidden"""				,       """text_projection"""				)
 if "text_projector.linear_out" in name:
  _snake_case    :  Tuple					=					name.replace("""text_projector.linear_out"""				,       """text_projection.3"""				)

 return name



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	List[str]				,       SCREAMING_SNAKE_CASE__       :	int				)  ->     List[str]:
 for key in orig_state_dict.copy().keys():
  _snake_case    :  Tuple					=					orig_state_dict.pop(SCREAMING_SNAKE_CASE__				)

  if "qkv" in key:
   # weights and biases of the key, value and query projections of vision encoder's attention layers require special treatment:
   # we need to split them up into separate matrices/vectors
   _snake_case    :  Any					=					key.split("""."""				)
   _snake_case    :  List[Any]					=					int(key_split[2]				), int(key_split[4]				)
   _snake_case    :  List[Any]					=					config.vision_config.hidden_size
   if "weight" in key:
    _snake_case    :  List[Any]					=					val[:dim, :]
    _snake_case    :  Union[str, Any]					=					val[dim : dim * 2, :]
    _snake_case    :  int					=					val[-dim:, :]
   else:
    _snake_case    :  Union[str, Any]					=					val[:dim]
    _snake_case    :  str					=					val[dim : dim * 2]
    _snake_case    :  Dict					=					val[-dim:]
  elif "in_proj" in key:
   # weights and biases of the key, value and query projections of text encoder's attention layers require special treatment:
   # we need to split them up into separate matrices/vectors
   _snake_case    :  int					=					key.split("""."""				)
   _snake_case    :  Optional[int]					=					int(key_split[3]				)
   _snake_case    :  Union[str, Any]					=					config.text_config.hidden_size
   if "weight" in key:
    _snake_case    :  int					=					val[:dim, :]
    _snake_case    :  Tuple					=					val[
        dim : dim * 2, :
    ]
    _snake_case    :  int					=					val[-dim:, :]
   else:
    _snake_case    :  Tuple					=					val[:dim]
    _snake_case    :  Tuple					=					val[dim : dim * 2]
    _snake_case    :  str					=					val[-dim:]
  else:
   _snake_case    :  Tuple					=					rename_key(SCREAMING_SNAKE_CASE__				)
   # squeeze if necessary
   if (
       "text_projection.0" in new_name
       or "text_projection.3" in new_name
       or "visual_projection.0" in new_name
       or "visual_projection.3" in new_name
   ):
    _snake_case    :  Any					=					val.squeeze_()
   else:
    _snake_case    :  Optional[int]					=					val

 return orig_state_dict



def       lowercase       (   )  ->     List[str]:
 _snake_case    :  List[str]					=					"""http://images.cocodataset.org/val2017/000000039769.jpg"""
 _snake_case    :  Dict					=					Image.open(requests.get(SCREAMING_SNAKE_CASE__				,       stream=SCREAMING_SNAKE_CASE__				).raw				)
 return im



@torch.no_grad()
def       lowercase       (   SCREAMING_SNAKE_CASE__       :	Any				,       SCREAMING_SNAKE_CASE__       :	str				,       SCREAMING_SNAKE_CASE__       :	int="groupvit-gcc-yfcc"				,       SCREAMING_SNAKE_CASE__       :	Optional[int]=False				)  ->     Optional[Any]:
 _snake_case    :  Any					=					GroupViTConfig()
 _snake_case    :  List[Any]					=					GroupViTModel(SCREAMING_SNAKE_CASE__				).eval()

 _snake_case    :  Optional[int]					=					torch.load(SCREAMING_SNAKE_CASE__				,       map_location="""cpu"""				)["""model"""]
 _snake_case    :  Tuple					=					convert_state_dict(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				)
 _snake_case    :  int					=					model.load_state_dict(SCREAMING_SNAKE_CASE__				,       strict=SCREAMING_SNAKE_CASE__				)
 assert missing_keys == ["text_model.embeddings.position_ids"]
 assert (unexpected_keys == ["multi_label_logit_scale"]) or (len(SCREAMING_SNAKE_CASE__				) == 0)

 # verify result
 _snake_case    :  List[Any]					=					CLIPProcessor.from_pretrained("""openai/clip-vit-base-patch32"""				)
 _snake_case    :  List[Any]					=					prepare_img()
 _snake_case    :  int					=					processor(text=["""a photo of a cat""", """a photo of a dog"""]				,       images=SCREAMING_SNAKE_CASE__				,       padding=SCREAMING_SNAKE_CASE__				,       return_tensors="""pt"""				)

 with torch.no_grad():
  _snake_case    :  Optional[int]					=					model(**SCREAMING_SNAKE_CASE__				)

 if model_name == "groupvit-gcc-yfcc":
  _snake_case    :  Union[str, Any]					=					torch.tensor([[13.3_523, 6.3_6_2_9]]				)
 elif model_name == "groupvit-gcc-redcaps":
  _snake_case    :  Union[str, Any]					=					torch.tensor([[16.1_873, 8.6_2_3_0]]				)
 else:
  raise ValueError(F'''Model name {model_name} not supported.'''				)
 assert torch.allclose(outputs.logits_per_image				,       SCREAMING_SNAKE_CASE__				,       atol=1e-3				)

 processor.save_pretrained(SCREAMING_SNAKE_CASE__				)
 model.save_pretrained(SCREAMING_SNAKE_CASE__				)
 print("""Successfully saved processor and model to"""				,       SCREAMING_SNAKE_CASE__				)

 if push_to_hub:
  print("""Pushing to the hub..."""				)
  processor.push_to_hub(SCREAMING_SNAKE_CASE__				,       organization="""nielsr"""				)
  model.push_to_hub(SCREAMING_SNAKE_CASE__				,       organization="""nielsr"""				)


if __name__ == "__main__":
  a__    =	argparse.ArgumentParser()
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to dump the processor and PyTorch model."""
  )
  parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to GroupViT checkpoint""")
  parser.add_argument(
      """--model_name""",
      default="""groupvit-gccy-fcc""",
      type=str,
      help="""Name of the model. Expecting either 'groupvit-gcc-yfcc' or 'groupvit-gcc-redcaps'""",
  )
  parser.add_argument(
      """--push_to_hub""",
      action="""store_true""",
      help="""Whether or not to push the converted model and processor to the 🤗 hub using the provided `model_name`.""",
  )
  a__    =	parser.parse_args()

  convert_groupvit_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.model_name, 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 		snake_case						:



   '''simple docstring'''






   def __init__(					self							:			Optional[Any]				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			Optional[Any]=13				,		lowerCAmelCase							:			int=7				,		lowerCAmelCase							:			Optional[int]=True				,		lowerCAmelCase							:			Optional[Any]=True				,		lowerCAmelCase							:			List[str]=True				,		lowerCAmelCase							:			str=True				,		lowerCAmelCase							:			Any=99				,		lowerCAmelCase							:			Dict=32				,		lowerCAmelCase							:			List[str]=2				,		lowerCAmelCase							:			str=4				,		lowerCAmelCase							:			List[str]=37				,		lowerCAmelCase							:			Dict="gelu"				,		lowerCAmelCase							:			Optional[int]=0.1				,		lowerCAmelCase							:			Any=0.1				,		lowerCAmelCase							:			List[Any]=512				,		lowerCAmelCase							:			Optional[int]=16				,		lowerCAmelCase							:			Dict=2				,		lowerCAmelCase							:			str=0.02				,		lowerCAmelCase							:			Optional[Any]=3				,		lowerCAmelCase							:			Optional[Any]=4				,		lowerCAmelCase							:			Tuple=None				,		lowerCAmelCase							:			List[Any]=0				,		)				->				int:

    """simple docstring"""







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

   def      UpperCamelCase_  (					self							:			Dict)				->				List[str]:

    """simple docstring"""







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

    _snake_case    :  List[str]					=					None
    if self.use_input_mask:
     # follow test_modeling_tf_ctrl.py
     _snake_case    :  List[Any]					=					random_attention_mask([self.batch_size, self.seq_length])

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

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

    _snake_case    :  Union[str, Any]					=					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=lowerCAmelCase				,		initializer_range=self.initializer_range				,		)
    _snake_case    :  Optional[int]					=					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							:			Optional[int]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			List[str])				->				List[Any]:

    """simple docstring"""







    _snake_case    :  Dict					=					TFDPRContextEncoder(config=lowerCAmelCase)
    _snake_case    :  Union[str, Any]					=					model(lowerCAmelCase				,		attention_mask=lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  Optional[int]					=					model(lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  List[str]					=					model(lowerCAmelCase)
    self.parent.assertEqual(result.pooler_output.shape				,		(self.batch_size, self.projection_dim or self.hidden_size))

   def      UpperCamelCase_  (					self							:			List[Any]				,		lowerCAmelCase							:			int				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			Any				,		lowerCAmelCase							:			Any				,		lowerCAmelCase							:			Union[str, Any]				,		lowerCAmelCase							:			Optional[int])				->				int:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					TFDPRQuestionEncoder(config=lowerCAmelCase)
    _snake_case    :  List[Any]					=					model(lowerCAmelCase				,		attention_mask=lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  List[Any]					=					model(lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  Tuple					=					model(lowerCAmelCase)
    self.parent.assertEqual(result.pooler_output.shape				,		(self.batch_size, self.projection_dim or self.hidden_size))

   def      UpperCamelCase_  (					self							:			List[str]				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Union[str, Any]				,		lowerCAmelCase							:			Dict				,		lowerCAmelCase							:			Optional[int]				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			str)				->				Tuple:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					TFDPRReader(config=lowerCAmelCase)
    _snake_case    :  Any					=					model(lowerCAmelCase				,		attention_mask=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))
    self.parent.assertEqual(result.relevance_logits.shape				,		(self.batch_size,))


   def      UpperCamelCase_  (					self							:			str)				->				Any:

    """simple docstring"""







    _snake_case    :  List[Any]					=					self.prepare_config_and_inputs()
    (
        (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    
    )    :  Union[str, Any]					=					config_and_inputs
    _snake_case    :  Union[str, Any]					=					{"""input_ids""": input_ids}
    return config, inputs_dict







@require_tf
class 		snake_case						(	SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,unittest.TestCase  ):



   '''simple docstring'''






   snake_case_					: Any          =		(
       (
           TFDPRContextEncoder,
           TFDPRQuestionEncoder,
           TFDPRReader,
       )
       if is_tf_available()
       else ()
   )
   snake_case_					: List[str]          =		{"""feature-extraction""": TFDPRQuestionEncoder} if is_tf_available() else {}

   snake_case_					: Tuple          =		False
   snake_case_					: str          =		False
   snake_case_					: Tuple          =		False
   snake_case_					: int          =		False
   snake_case_					: List[Any]          =		False

   def      UpperCamelCase_  (					self							:			Tuple)				->				str:

    """simple docstring"""







    _snake_case    :  List[str]					=					TFDPRModelTester(self)
    _snake_case    :  Dict					=					ConfigTester(self				,		config_class=lowerCAmelCase				,		hidden_size=37)

   def      UpperCamelCase_  (					self							:			List[Any])				->				Optional[Any]:

    """simple docstring"""







    self.config_tester.run_common_tests()

   def      UpperCamelCase_  (					self							:			List[str])				->				Optional[int]:

    """simple docstring"""







    _snake_case    :  Any					=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_dpr_context_encoder(*lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			List[str])				->				Optional[Any]:

    """simple docstring"""







    _snake_case    :  str					=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_dpr_question_encoder(*lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			List[str])				->				Tuple:

    """simple docstring"""







    _snake_case    :  str					=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_dpr_reader(*lowerCAmelCase)


   @slow
   def      UpperCamelCase_  (					self							:			List[str])				->				Any:

    """simple docstring"""







    for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  List[Any]					=					TFDPRContextEncoder.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)

    for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  List[str]					=					TFDPRContextEncoder.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)

    for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  Any					=					TFDPRQuestionEncoder.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)

    for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  Union[str, Any]					=					TFDPRReader.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)







@require_tf
class 		snake_case						(	unittest.TestCase  ):



   '''simple docstring'''






   @slow
   def      UpperCamelCase_  (					self							:			List[str])				->				Optional[int]:

    """simple docstring"""







    _snake_case    :  List[str]					=					TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base""")

    _snake_case    :  str					=					tf.constant(
        [[101, 7592, 1010, 2003, 2026, 3899, 1_0140, 1029, 102]])  # [CLS] hello, is my dog cute? [SEP]
    _snake_case    :  Optional[int]					=					model(lowerCAmelCase)[0]  # embedding shape = (1, 768)
    # compare the actual values for a slice.
    _snake_case    :  Union[str, Any]					=					tf.constant(
        [
            [
                0.03_236_253,
                0.12_753_335,
                0.16_818_509,
                0.00_279_786,
                0.3_896_933,
                0.24_264_945,
                0.2_178_971,
                -0.02_335_227,
                -0.08_481_959,
                -0.14_324_117,
            ]
        ])
    self.assertTrue(numpy.allclose(output[:, :10].numpy()				,		expected_slice.numpy()				,		atol=1E-4))
