
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		:
				def __init__(  self		,snake_case		,snake_case=13		,snake_case=7		,snake_case=True		,snake_case=True		,snake_case=True		,snake_case=True		,snake_case=99		,snake_case=32		,snake_case=2		,snake_case=4		,snake_case=37		,snake_case="gelu"		,snake_case=0.1		,snake_case=0.1		,snake_case=512		,snake_case=16		,snake_case=2		,snake_case=0.02		,snake_case=3		,snake_case=4		,snake_case=None		,snake_case=0		,):



											'''simple docstring'''







											lowercase							:					str	     =				parent
											lowercase							:					Union[str, Any]	     =				batch_size
											lowercase							:					int	     =				seq_length
											lowercase							:					List[Any]	     =				is_training
											lowercase							:					str	     =				use_input_mask
											lowercase							:					Union[str, Any]	     =				use_token_type_ids
											lowercase							:					Optional[Any]	     =				use_labels
											lowercase							:					str	     =				vocab_size
											lowercase							:					List[Any]	     =				hidden_size
											lowercase							:					int	     =				num_hidden_layers
											lowercase							:					Dict	     =				num_attention_heads
											lowercase							:					int	     =				intermediate_size
											lowercase							:					Optional[Any]	     =				hidden_act
											lowercase							:					List[Any]	     =				hidden_dropout_prob
											lowercase							:					List[Any]	     =				attention_probs_dropout_prob
											lowercase							:					Optional[int]	     =				max_position_embeddings
											lowercase							:					str	     =				type_vocab_size
											lowercase							:					Optional[Any]	     =				type_sequence_label_size
											lowercase							:					List[str]	     =				initializer_range
											lowercase							:					int	     =				num_labels
											lowercase							:					Tuple	     =				num_choices
											lowercase							:					List[Any]	     =				scope
											lowercase							:					List[str]	     =				projection_dim
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







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

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

											lowercase							:					Union[str, Any]	     =				None
											if self.use_token_type_ids:
																		lowercase							:					Dict	     =				ids_tensor([self.batch_size, self.seq_length]		,self.type_vocab_size				)

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

											lowercase							:					str	     =				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							:					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 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case				):



											'''simple docstring'''







											lowercase							:					Dict	     =				TFDPRContextEncoder(config=_a				)
											lowercase							:					Dict	     =				model(_a		,attention_mask=_a		,token_type_ids=_a				)
											lowercase							:					int	     =				model(_a		,token_type_ids=_a				)
											lowercase							:					Any	     =				model(_a				)
											self.parent.assertEqual(result.pooler_output.shape		,(self.batch_size, self.projection_dim or self.hidden_size)				)
				def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case				):



											'''simple docstring'''







											lowercase							:					str	     =				TFDPRQuestionEncoder(config=_a				)
											lowercase							:					Dict	     =				model(_a		,attention_mask=_a		,token_type_ids=_a				)
											lowercase							:					str	     =				model(_a		,token_type_ids=_a				)
											lowercase							:					List[str]	     =				model(_a				)
											self.parent.assertEqual(result.pooler_output.shape		,(self.batch_size, self.projection_dim or self.hidden_size)				)
				def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case				):



											'''simple docstring'''







											lowercase							:					List[Any]	     =				TFDPRReader(config=_a				)
											lowercase							:					Any	     =				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 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







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






@require_tf
class __snake_case		(				_a ,     _a ,     unittest.TestCase					):
				_a		:						Any=    (
				    (
				        TFDPRContextEncoder,
				        TFDPRQuestionEncoder,
				        TFDPRReader,
				    )
				    if is_tf_available()
				    else ()
				)
				_a		:						Tuple=    {"feature-extraction": TFDPRQuestionEncoder} if is_tf_available() else {}

				_a		:						int=    False
				_a		:						Optional[int]=    False
				_a		:						int=    False
				_a		:						Optional[int]=    False
				_a		:						str=    False
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					Any	     =				TFDPRModelTester(self				)
											lowercase							:					Tuple	     =				ConfigTester(self		,config_class=_a		,hidden_size=37				)
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											self.config_tester.run_common_tests()
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					Optional[int]	     =				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_dpr_context_encoder(*_a				)
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					str	     =				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_dpr_question_encoder(*_a				)
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					int	     =				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_dpr_reader(*_a				)
				@slow
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowercase							:					str	     =				TFDPRContextEncoder.from_pretrained(_a				)
																		self.assertIsNotNone(_a				)

											for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowercase							:					str	     =				TFDPRContextEncoder.from_pretrained(_a				)
																		self.assertIsNotNone(_a				)

											for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowercase							:					Tuple	     =				TFDPRQuestionEncoder.from_pretrained(_a				)
																		self.assertIsNotNone(_a				)

											for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowercase							:					Tuple	     =				TFDPRReader.from_pretrained(_a				)
																		self.assertIsNotNone(_a				)






@require_tf
class __snake_case		(				unittest.TestCase					):
				@slow
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







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

											lowercase							:					int	     =				tf.constant(
											    [[101, 7592, 1010, 2003, 2026, 3899, 10140, 1029, 102]]				)  # [CLS] hello, is my dog cute? [SEP]
											lowercase							:					str	     =				model(_a				)[0]  # embedding shape = (1, 768)
											# compare the actual values for a slice.
											lowercase							:					Tuple	     =				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				)				)


import argparse

import requests
import torch
from PIL import Image

from transformers import CLIPProcessor, GroupViTConfig, GroupViTModel







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

       return name







def 						_snake_case(  SCREAMING_SNAKE_CASE__ ,     SCREAMING_SNAKE_CASE__    )			->   List[str]:
       for key in orig_state_dict.copy().keys():
              lowercase							:					List[Any]	     =				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
                     lowercase							:					Dict	     =				key.split("""."""    )
                     lowercase     ,      lowercase							:					Optional[Any]	     =				int(key_split[2]    ), int(key_split[4]    )
                     lowercase							:					List[Any]	     =				config.vision_config.hidden_size
                     if "weight" in key:
                            lowercase							:					str	     =				val[:dim, :]
                            lowercase							:					List[str]	     =				val[dim : dim * 2, :]
                            lowercase							:					Optional[int]	     =				val[-dim:, :]
                     else:
                            lowercase							:					Dict	     =				val[:dim]
                            lowercase							:					Dict	     =				val[dim : dim * 2]
                            lowercase							:					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
                     lowercase							:					int	     =				key.split("""."""    )
                     lowercase							:					Tuple	     =				int(key_split[3]    )
                     lowercase							:					str	     =				config.text_config.hidden_size
                     if "weight" in key:
                            lowercase							:					Optional[int]	     =				val[:dim, :]
                            lowercase							:					Optional[Any]	     =				val[
                                dim : dim * 2, :
                            ]
                            lowercase							:					Optional[int]	     =				val[-dim:, :]
                     else:
                            lowercase							:					Optional[int]	     =				val[:dim]
                            lowercase							:					Dict	     =				val[dim : dim * 2]
                            lowercase							:					List[str]	     =				val[-dim:]
              else:
                     lowercase							:					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
                     ):
                            lowercase							:					str	     =				val.squeeze_()
                     else:
                            lowercase							:					Any	     =				val

       return orig_state_dict







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







@torch.no_grad()
def 						_snake_case(  SCREAMING_SNAKE_CASE__ ,     SCREAMING_SNAKE_CASE__ ,     SCREAMING_SNAKE_CASE__="groupvit-gcc-yfcc" ,     SCREAMING_SNAKE_CASE__=False    )			->   str:
       lowercase							:					Dict	     =				GroupViTConfig()
       lowercase							:					Tuple	     =				GroupViTModel(SCREAMING_SNAKE_CASE__    ).eval()

       lowercase							:					List[str]	     =				torch.load(SCREAMING_SNAKE_CASE__ ,     map_location="""cpu"""    )["""model"""]
       lowercase							:					Optional[Any]	     =				convert_state_dict(SCREAMING_SNAKE_CASE__ ,     SCREAMING_SNAKE_CASE__    )
       lowercase     ,      lowercase							:					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
       lowercase							:					Tuple	     =				CLIPProcessor.from_pretrained("""openai/clip-vit-base-patch32"""    )
       lowercase							:					Optional[Any]	     =				prepare_img()
       lowercase							:					List[Any]	     =				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():
              lowercase							:					Optional[Any]	     =				model(**SCREAMING_SNAKE_CASE__    )

       if model_name == "groupvit-gcc-yfcc":
              lowercase							:					Optional[int]	     =				torch.tensor([[13.3523, 6.3629]]    )
       elif model_name == "groupvit-gcc-redcaps":
              lowercase							:					int	     =				torch.tensor([[16.1873, 8.6230]]    )
       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__":
  lowercase  : List[Any]						   =							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`.""",
  )
  lowercase  : Union[str, Any]						   =							parser.parse_args()

  convert_groupvit_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)
