
import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
		import jax
		import jax.numpy as jnp

		from transformers.generation import (
		    FlaxForcedBOSTokenLogitsProcessor,
		    FlaxForcedEOSTokenLogitsProcessor,
		    FlaxLogitsProcessorList,
		    FlaxMinLengthLogitsProcessor,
		    FlaxTemperatureLogitsWarper,
		    FlaxTopKLogitsWarper,
		    FlaxTopPLogitsWarper,
		)

@require_flax
class __snake_case		(				unittest.TestCase					):
				def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case				):



											'''simple docstring'''







											lowercase							:					List[str]	     =				jnp.ones((batch_size, length)				) / length
											return scores
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					List[Any]	     =				None
											lowercase							:					Dict	     =				20

											lowercase							:					Dict	     =				self._get_uniform_logits(batch_size=2		,length=snake_case				)

											# tweak scores to not be uniform anymore
											lowercase							:					Optional[int]	     =				scores.at[1, 5].set((1 / length) + 0.1				)  # peak, 1st batch
											lowercase							:					Tuple	     =				scores.at[1, 10].set((1 / length) - 0.4				)  # valley, 1st batch

											# compute softmax
											lowercase							:					List[str]	     =				jax.nn.softmax(snake_case		,axis=-1				)

											lowercase							:					Any	     =				FlaxTemperatureLogitsWarper(temperature=0.5				)
											lowercase							:					Optional[int]	     =				FlaxTemperatureLogitsWarper(temperature=1.3				)

											lowercase							:					Optional[int]	     =				jax.nn.softmax(temp_dist_warper_sharper(snake_case		,scores.copy()		,cur_len=snake_case				)		,axis=-1				)
											lowercase							:					Optional[int]	     =				jax.nn.softmax(temp_dist_warper_smoother(snake_case		,scores.copy()		,cur_len=snake_case				)		,axis=-1				)

											# uniform distribution stays uniform
											self.assertTrue(jnp.allclose(probs[0, :]		,warped_prob_sharp[0, :]		,atol=1e-3				)				)
											self.assertTrue(jnp.allclose(probs[0, :]		,warped_prob_smooth[0, :]		,atol=1e-3				)				)

											# sharp peaks get higher, valleys get lower
											self.assertLess(probs[1, :].max()		,warped_prob_sharp[1, :].max()				)
											self.assertGreater(probs[1, :].min()		,warped_prob_sharp[1, :].min()				)

											# smooth peaks get lower, valleys get higher
											self.assertGreater(probs[1, :].max()		,warped_prob_smooth[1, :].max()				)
											self.assertLess(probs[1, :].min()		,warped_prob_smooth[1, :].min()				)
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					Tuple	     =				None
											lowercase							:					List[Any]	     =				10
											lowercase							:					Any	     =				2

											# create ramp distribution
											lowercase							:					List[str]	     =				np.broadcast_to(np.arange(snake_case				)[None, :]		,(batch_size, vocab_size)				).copy()
											lowercase							:					Union[str, Any]	     =				ramp_logits[1:, : vocab_size // 2] + vocab_size

											lowercase							:					Optional[Any]	     =				FlaxTopKLogitsWarper(3				)

											lowercase							:					List[str]	     =				top_k_warp(snake_case		,snake_case		,cur_len=snake_case				)

											# check that correct tokens are filtered
											self.assertListEqual(jnp.isinf(scores[0]				).tolist()		,7 * [True] + 3 * [False]				)
											self.assertListEqual(jnp.isinf(scores[1]				).tolist()		,2 * [True] + 3 * [False] + 5 * [True]				)

											# check special case
											lowercase							:					int	     =				5
											lowercase							:					Dict	     =				FlaxTopKLogitsWarper(top_k=1		,filter_value=0.0		,min_tokens_to_keep=3				)

											lowercase							:					str	     =				np.broadcast_to(np.arange(snake_case				)[None, :]		,(batch_size, length)				).copy()
											lowercase							:					int	     =				top_k_warp_safety_check(snake_case		,snake_case		,cur_len=snake_case				)

											# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
											self.assertListEqual((scores == 0.0).sum(axis=-1				).tolist()		,[2, 2]				)
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					Any	     =				None
											lowercase							:					Dict	     =				10
											lowercase							:					Union[str, Any]	     =				2

											# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
											lowercase							:					Optional[int]	     =				np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]				)				)

											lowercase							:					Optional[int]	     =				FlaxTopPLogitsWarper(0.8				)
											lowercase							:					Dict	     =				np.exp(top_p_warp(snake_case		,snake_case		,cur_len=snake_case				)				)

											# dist should be filtered to keep min num values so that sum is >= top_p
											# exp (-inf) => 0
											lowercase							:					Any	     =				np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]				)
											self.assertTrue(np.allclose(snake_case		,snake_case		,atol=1e-3				)				)

											# check edge cases with negative and extreme logits
											lowercase							:					Dict	     =				np.broadcast_to(np.arange(snake_case				)[None, :]		,(batch_size, vocab_size)				).copy() - (
											    vocab_size // 2
											)

											# make ramp_logits more extreme
											lowercase							:					Optional[int]	     =				ramp_logits[1] * 100.0

											# make sure at least 2 tokens are kept
											lowercase							:					List[Any]	     =				FlaxTopPLogitsWarper(0.9		,min_tokens_to_keep=2		,filter_value=0.0				)
											lowercase							:					Dict	     =				top_p_warp(snake_case		,snake_case		,cur_len=snake_case				)

											# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
											self.assertListEqual((filtered_dist != 0.0).sum(axis=-1				).tolist()		,[3, 2]				)
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					Union[str, Any]	     =				20
											lowercase							:					Any	     =				4
											lowercase							:					List[str]	     =				0

											lowercase							:					Any	     =				FlaxMinLengthLogitsProcessor(min_length=10		,eos_token_id=snake_case				)

											# check that min length is applied at length 5
											lowercase							:					Optional[Any]	     =				ids_tensor((batch_size, 20)		,vocab_size=20				)
											lowercase							:					Tuple	     =				5
											lowercase							:					Optional[Any]	     =				self._get_uniform_logits(snake_case		,snake_case				)
											lowercase							:					str	     =				min_dist_processor(snake_case		,snake_case		,cur_len=snake_case				)
											self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()		,4 * [-float("""inf"""				)]				)

											# check that min length is not applied anymore at length 15
											lowercase							:					Union[str, Any]	     =				self._get_uniform_logits(snake_case		,snake_case				)
											lowercase							:					List[Any]	     =				15
											lowercase							:					str	     =				min_dist_processor(snake_case		,snake_case		,cur_len=snake_case				)
											self.assertFalse(jnp.isinf(snake_case				).any()				)
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					Union[str, Any]	     =				20
											lowercase							:					str	     =				4
											lowercase							:					Dict	     =				0

											lowercase							:					Any	     =				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case				)

											# check that all scores are -inf except the bos_token_id score
											lowercase							:					Any	     =				ids_tensor((batch_size, 1)		,vocab_size=20				)
											lowercase							:					Dict	     =				1
											lowercase							:					Tuple	     =				self._get_uniform_logits(snake_case		,snake_case				)
											lowercase							:					Any	     =				logits_processor(snake_case		,snake_case		,cur_len=snake_case				)
											self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]				).all()				)
											self.assertListEqual(scores[:, bos_token_id].tolist()		,4 * [0]				)  # score for bos_token_id shold be zero

											# check that bos_token_id is not forced if current length is greater than 1
											lowercase							:					int	     =				3
											lowercase							:					List[Any]	     =				self._get_uniform_logits(snake_case		,snake_case				)
											lowercase							:					List[Any]	     =				logits_processor(snake_case		,snake_case		,cur_len=snake_case				)
											self.assertFalse(jnp.isinf(snake_case				).any()				)
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					str	     =				20
											lowercase							:					str	     =				4
											lowercase							:					Dict	     =				0
											lowercase							:					Any	     =				5

											lowercase							:					List[Any]	     =				FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case		,eos_token_id=snake_case				)

											# check that all scores are -inf except the eos_token_id when max_length is reached
											lowercase							:					Any	     =				ids_tensor((batch_size, 4)		,vocab_size=20				)
											lowercase							:					Any	     =				4
											lowercase							:					Tuple	     =				self._get_uniform_logits(snake_case		,snake_case				)
											lowercase							:					List[str]	     =				logits_processor(snake_case		,snake_case		,cur_len=snake_case				)
											self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]				).all()				)
											self.assertListEqual(scores[:, eos_token_id].tolist()		,4 * [0]				)  # score for eos_token_id should be zero

											# check that eos_token_id is not forced if max_length is not reached
											lowercase							:					Dict	     =				3
											lowercase							:					int	     =				self._get_uniform_logits(snake_case		,snake_case				)
											lowercase							:					Dict	     =				logits_processor(snake_case		,snake_case		,cur_len=snake_case				)
											self.assertFalse(jnp.isinf(snake_case				).any()				)
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					Any	     =				4
											lowercase							:					Union[str, Any]	     =				10
											lowercase							:					List[Any]	     =				15
											lowercase							:					Any	     =				2
											lowercase							:					Tuple	     =				1
											lowercase							:					Union[str, Any]	     =				15

											# dummy input_ids and scores
											lowercase							:					str	     =				ids_tensor((batch_size, sequence_length)		,snake_case				)
											lowercase							:					str	     =				input_ids.copy()

											lowercase							:					Any	     =				self._get_uniform_logits(snake_case		,snake_case				)
											lowercase							:					Any	     =				scores.copy()

											# instantiate all dist processors
											lowercase							:					str	     =				FlaxTemperatureLogitsWarper(temperature=0.5				)
											lowercase							:					Union[str, Any]	     =				FlaxTopKLogitsWarper(3				)
											lowercase							:					List[Any]	     =				FlaxTopPLogitsWarper(0.8				)

											# instantiate all logits processors
											lowercase							:					List[Any]	     =				FlaxMinLengthLogitsProcessor(min_length=10		,eos_token_id=snake_case				)
											lowercase							:					str	     =				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case				)
											lowercase							:					Union[str, Any]	     =				FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case		,eos_token_id=snake_case				)

											lowercase							:					Optional[Any]	     =				10

											# no processor list
											lowercase							:					List[str]	     =				temp_dist_warp(snake_case		,snake_case		,cur_len=snake_case				)
											lowercase							:					Optional[Any]	     =				top_k_warp(snake_case		,snake_case		,cur_len=snake_case				)
											lowercase							:					str	     =				top_p_warp(snake_case		,snake_case		,cur_len=snake_case				)
											lowercase							:					List[Any]	     =				min_dist_proc(snake_case		,snake_case		,cur_len=snake_case				)
											lowercase							:					List[Any]	     =				bos_dist_proc(snake_case		,snake_case		,cur_len=snake_case				)
											lowercase							:					Optional[Any]	     =				eos_dist_proc(snake_case		,snake_case		,cur_len=snake_case				)

											# with processor list
											lowercase							:					List[Any]	     =				FlaxLogitsProcessorList(
											    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]				)
											lowercase							:					int	     =				processor(snake_case		,snake_case		,cur_len=snake_case				)

											# scores should be equal
											self.assertTrue(jnp.allclose(snake_case		,snake_case		,atol=1e-3				)				)

											# input_ids should never be changed
											self.assertListEqual(input_ids.tolist()		,input_ids_comp.tolist()				)
				def 				_SCREAMING_SNAKE_CASE	(  self				):



											'''simple docstring'''







											lowercase							:					Dict	     =				4
											lowercase							:					int	     =				10
											lowercase							:					List[Any]	     =				15
											lowercase							:					Tuple	     =				2
											lowercase							:					Optional[int]	     =				1
											lowercase							:					Any	     =				15

											# dummy input_ids and scores
											lowercase							:					int	     =				ids_tensor((batch_size, sequence_length)		,snake_case				)
											lowercase							:					Dict	     =				input_ids.copy()

											lowercase							:					int	     =				self._get_uniform_logits(snake_case		,snake_case				)
											lowercase							:					Optional[Any]	     =				scores.copy()

											# instantiate all dist processors
											lowercase							:					Optional[int]	     =				FlaxTemperatureLogitsWarper(temperature=0.5				)
											lowercase							:					Dict	     =				FlaxTopKLogitsWarper(3				)
											lowercase							:					Optional[Any]	     =				FlaxTopPLogitsWarper(0.8				)

											# instantiate all logits processors
											lowercase							:					Tuple	     =				FlaxMinLengthLogitsProcessor(min_length=10		,eos_token_id=snake_case				)
											lowercase							:					List[Any]	     =				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case				)
											lowercase							:					int	     =				FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case		,eos_token_id=snake_case				)

											lowercase							:					Optional[Any]	     =				10

											# no processor list
											def run_no_processor_list(snake_case		,snake_case		,snake_case				):
																		lowercase							:					Optional[Any]	     =				temp_dist_warp(snake_case		,snake_case		,cur_len=snake_case				)
																		lowercase							:					Optional[Any]	     =				top_k_warp(snake_case		,snake_case		,cur_len=snake_case				)
																		lowercase							:					List[str]	     =				top_p_warp(snake_case		,snake_case		,cur_len=snake_case				)
																		lowercase							:					Union[str, Any]	     =				min_dist_proc(snake_case		,snake_case		,cur_len=snake_case				)
																		lowercase							:					List[str]	     =				bos_dist_proc(snake_case		,snake_case		,cur_len=snake_case				)
																		lowercase							:					Any	     =				eos_dist_proc(snake_case		,snake_case		,cur_len=snake_case				)
																		return scores

											# with processor list
											def run_processor_list(snake_case		,snake_case		,snake_case				):
																		lowercase							:					int	     =				FlaxLogitsProcessorList(
																		    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]				)
																		lowercase							:					Union[str, Any]	     =				processor(snake_case		,snake_case		,cur_len=snake_case				)
																		return scores

											lowercase							:					List[Any]	     =				jax.jit(snake_case				)
											lowercase							:					Dict	     =				jax.jit(snake_case				)

											lowercase							:					Tuple	     =				jitted_run_no_processor_list(snake_case		,snake_case		,snake_case				)
											lowercase							:					Any	     =				jitted_run_processor_list(snake_case		,snake_case		,snake_case				)

											# scores should be equal
											self.assertTrue(jnp.allclose(snake_case		,snake_case		,atol=1e-3				)				)

											# input_ids should never be changed
											self.assertListEqual(input_ids.tolist()		,input_ids_comp.tolist()				)


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)
