



from __future__ import annotations

import unittest

from transformers import is_tf_available, is_torch_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, SMALL_MODEL_IDENTIFIER, is_pt_tf_cross_test, slow


if is_tf_available():
  from transformers import (
      AutoConfig,
      BertConfig,
      GPTaConfig,
      TaConfig,
      TFAutoModel,
      TFAutoModelForCausalLM,
      TFAutoModelForMaskedLM,
      TFAutoModelForPreTraining,
      TFAutoModelForQuestionAnswering,
      TFAutoModelForSeqaSeqLM,
      TFAutoModelForSequenceClassification,
      TFAutoModelWithLMHead,
      TFBertForMaskedLM,
      TFBertForPreTraining,
      TFBertForQuestionAnswering,
      TFBertForSequenceClassification,
      TFBertModel,
      TFGPTaLMHeadModel,
      TFRobertaForMaskedLM,
      TFTaForConditionalGeneration,
  )
  from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
  from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
  from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST

if is_torch_available():
  from transformers import (
      AutoModel,
      AutoModelForCausalLM,
      AutoModelForMaskedLM,
      AutoModelForPreTraining,
      AutoModelForQuestionAnswering,
      AutoModelForSeqaSeqLM,
      AutoModelForSequenceClassification,
      AutoModelWithLMHead,
      BertForMaskedLM,
      BertForPreTraining,
      BertForQuestionAnswering,
      BertForSequenceClassification,
      BertModel,
      GPTaLMHeadModel,
      RobertaForMaskedLM,
      TaForConditionalGeneration,
  )







@is_pt_tf_cross_test
class 		snake_case						(	unittest.TestCase  ):



   '''simple docstring'''






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

    """simple docstring"""







    for model_name in ["bert-base-uncased"]:
     _snake_case    :  List[Any]					=					AutoConfig.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  str					=					TFAutoModel.from_pretrained(lowerCAmelCase				,		from_pt=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Optional[Any]					=					AutoModel.from_pretrained(lowerCAmelCase				,		from_tf=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

   @slow
   def      UpperCamelCase_  (					self							:			Optional[Any])				->				Union[str, Any]:

    """simple docstring"""







    for model_name in ["bert-base-uncased"]:
     _snake_case    :  int					=					AutoConfig.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Any					=					TFAutoModelForPreTraining.from_pretrained(lowerCAmelCase				,		from_pt=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  List[str]					=					AutoModelForPreTraining.from_pretrained(lowerCAmelCase				,		from_tf=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

   @slow
   def      UpperCamelCase_  (					self							:			int)				->				Tuple:

    """simple docstring"""







    for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  Optional[int]					=					AutoConfig.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Union[str, Any]					=					TFAutoModelForCausalLM.from_pretrained(lowerCAmelCase				,		from_pt=lowerCAmelCase)
     _snake_case	,    _snake_case    :  List[str]					=					TFAutoModelForCausalLM.from_pretrained(
         lowerCAmelCase				,		output_loading_info=lowerCAmelCase				,		from_pt=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Union[str, Any]					=					AutoModelForCausalLM.from_pretrained(lowerCAmelCase				,		from_tf=lowerCAmelCase)
     _snake_case	,    _snake_case    :  str					=					AutoModelForCausalLM.from_pretrained(
         lowerCAmelCase				,		output_loading_info=lowerCAmelCase				,		from_tf=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

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

    """simple docstring"""







    for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  int					=					AutoConfig.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Any					=					TFAutoModelWithLMHead.from_pretrained(lowerCAmelCase				,		from_pt=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  List[str]					=					AutoModelWithLMHead.from_pretrained(lowerCAmelCase				,		from_tf=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

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

    """simple docstring"""







    for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  Optional[Any]					=					AutoConfig.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Dict					=					TFAutoModelForMaskedLM.from_pretrained(lowerCAmelCase				,		from_pt=lowerCAmelCase)
     _snake_case	,    _snake_case    :  Dict					=					TFAutoModelForMaskedLM.from_pretrained(
         lowerCAmelCase				,		output_loading_info=lowerCAmelCase				,		from_pt=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Tuple					=					AutoModelForMaskedLM.from_pretrained(lowerCAmelCase				,		from_tf=lowerCAmelCase)
     _snake_case	,    _snake_case    :  str					=					AutoModelForMaskedLM.from_pretrained(
         lowerCAmelCase				,		output_loading_info=lowerCAmelCase				,		from_tf=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

   @slow
   def      UpperCamelCase_  (					self							:			Any)				->				Any:

    """simple docstring"""







    for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  List[Any]					=					AutoConfig.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Union[str, Any]					=					TFAutoModelForSeqaSeqLM.from_pretrained(lowerCAmelCase				,		from_pt=lowerCAmelCase)
     _snake_case	,    _snake_case    :  Optional[Any]					=					TFAutoModelForSeqaSeqLM.from_pretrained(
         lowerCAmelCase				,		output_loading_info=lowerCAmelCase				,		from_pt=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Tuple					=					AutoModelForSeqaSeqLM.from_pretrained(lowerCAmelCase				,		from_tf=lowerCAmelCase)
     _snake_case	,    _snake_case    :  int					=					AutoModelForSeqaSeqLM.from_pretrained(
         lowerCAmelCase				,		output_loading_info=lowerCAmelCase				,		from_tf=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

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

    """simple docstring"""







    for model_name in ["bert-base-uncased"]:
     _snake_case    :  Dict					=					AutoConfig.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  str					=					TFAutoModelForSequenceClassification.from_pretrained(lowerCAmelCase				,		from_pt=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Union[str, Any]					=					AutoModelForSequenceClassification.from_pretrained(lowerCAmelCase				,		from_tf=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

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

    """simple docstring"""







    for model_name in ["bert-base-uncased"]:
     _snake_case    :  Optional[int]					=					AutoConfig.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Optional[Any]					=					TFAutoModelForQuestionAnswering.from_pretrained(lowerCAmelCase				,		from_pt=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

     _snake_case    :  Optional[Any]					=					AutoModelForQuestionAnswering.from_pretrained(lowerCAmelCase				,		from_tf=lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)
     self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)

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

    """simple docstring"""







    _snake_case    :  Any					=					TFAutoModelWithLMHead.from_pretrained(lowerCAmelCase				,		from_pt=lowerCAmelCase)
    self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)
    self.assertEqual(model.num_parameters()				,		1_4410)
    self.assertEqual(model.num_parameters(only_trainable=lowerCAmelCase)				,		1_4410)

    _snake_case    :  Dict					=					AutoModelWithLMHead.from_pretrained(lowerCAmelCase				,		from_tf=lowerCAmelCase)
    self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)
    self.assertEqual(model.num_parameters()				,		1_4410)
    self.assertEqual(model.num_parameters(only_trainable=lowerCAmelCase)				,		1_4410)


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

    """simple docstring"""







    _snake_case    :  Optional[int]					=					TFAutoModelWithLMHead.from_pretrained(lowerCAmelCase				,		from_pt=lowerCAmelCase)
    self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)
    self.assertEqual(model.num_parameters()				,		1_4410)
    self.assertEqual(model.num_parameters(only_trainable=lowerCAmelCase)				,		1_4410)

    _snake_case    :  Optional[int]					=					AutoModelWithLMHead.from_pretrained(lowerCAmelCase				,		from_tf=lowerCAmelCase)
    self.assertIsInstance(lowerCAmelCase				,		lowerCAmelCase)
    self.assertEqual(model.num_parameters()				,		1_4410)
    self.assertEqual(model.num_parameters(only_trainable=lowerCAmelCase)				,		1_4410)





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	,    _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	,    _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([[1_3.3_5_2_3, 6.3_6_2_9]]				)
 elif model_name == "groupvit-gcc-redcaps":
  _snake_case    :  Union[str, Any]					=					torch.tensor([[1_6.1_8_7_3, 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)
