



'''simple docstring'''



from typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from configuration_hybrid_clip import HybridCLIPConfig
from flax.core.frozen_dict import FrozenDict

from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
from transformers.modeling_flax_utils import FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
from transformers.utils import logging


_lowercase       :						Union[str, Any]						=logging.get_logger(__name__)






class 			_SCREAMING_SNAKE_CASE				(nn.Module      ):
			A__								=  42
			A__								=  jnp.floataa



			def 				lowerCAmelCase (     self		:						Any			)      ->      Tuple:
							"""simple docstring"""





							snake_case__      :   List[str]      =    self.config.text_config
							snake_case__      :   List[str]      =    self.config.vision_config

							snake_case__      :   Tuple      =    self.config.projection_dim
							snake_case__      :   Optional[int]      =    text_config.hidden_size
							snake_case__      :   Tuple      =    vision_config.hidden_size

							snake_case__      :   Tuple      =    FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
							snake_case__      :   Optional[int]      =    FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__    ,			__UpperCamelCase			).module_class

							snake_case__      :   List[str]      =    text_module(__UpperCamelCase    ,			dtype=self.dtype			)
							snake_case__      :   Optional[Any]      =    vision_module(__UpperCamelCase    ,			dtype=self.dtype			)

							snake_case__      :   Optional[Any]      =    nn.Dense(
							    self.projection_dim    ,			dtype=self.dtype    ,			kernel_init=jax.nn.initializers.normal(0.02			)    ,			use_bias=__UpperCamelCase    ,			)
							snake_case__      :   List[Any]      =    nn.Dense(
							    self.projection_dim    ,			dtype=self.dtype    ,			kernel_init=jax.nn.initializers.normal(0.02			)    ,			use_bias=__UpperCamelCase    ,			)
							snake_case__      :   str      =    self.param('''logit_scale'''    ,			jax.nn.initializers.ones    ,			[]			)







			def __call__(     self		:						List[str]    ,			__UpperCamelCase		:						Union[str, Any]=None    ,			__UpperCamelCase		:						int=None    ,			__UpperCamelCase		:						Union[str, Any]=None    ,			__UpperCamelCase		:						Union[str, Any]=None    ,			__UpperCamelCase		:						List[Any]=None    ,			__UpperCamelCase		:						bool = True    ,			__UpperCamelCase		:						Any=None    ,			__UpperCamelCase		:						int=None    ,			__UpperCamelCase		:						Dict=None    ,			)      ->      int:
							"""simple docstring"""





							snake_case__      :   Optional[int]      =    return_dict if return_dict is not None else self.config.return_dict

							snake_case__      :   List[Any]      =    self.vision_model(
							    pixel_values=__UpperCamelCase    ,			deterministic=__UpperCamelCase    ,			output_attentions=__UpperCamelCase    ,			output_hidden_states=__UpperCamelCase    ,			return_dict=__UpperCamelCase    ,			)

							snake_case__      :   Dict      =    self.text_model(
							    input_ids=__UpperCamelCase    ,			attention_mask=__UpperCamelCase    ,			token_type_ids=__UpperCamelCase    ,			position_ids=__UpperCamelCase    ,			deterministic=__UpperCamelCase    ,			output_attentions=__UpperCamelCase    ,			output_hidden_states=__UpperCamelCase    ,			return_dict=__UpperCamelCase    ,			)

							snake_case__      :   List[Any]      =    vision_outputs[1]
							snake_case__      :   List[Any]      =    self.visual_projection(__UpperCamelCase			)

							snake_case__      :   str      =    text_outputs[1]
							snake_case__      :   Union[str, Any]      =    self.text_projection(__UpperCamelCase			)

							# normalized features
							snake_case__      :   int      =    image_embeds / jnp.linalg.norm(__UpperCamelCase    ,			axis=-1    ,			keepdims=__UpperCamelCase			)
							snake_case__      :   Optional[int]      =    text_embeds / jnp.linalg.norm(__UpperCamelCase    ,			axis=-1    ,			keepdims=__UpperCamelCase			)

							# cosine similarity as logits
							snake_case__      :   int      =    jnp.exp(self.logit_scale			)
							snake_case__      :   str      =    jnp.matmul(__UpperCamelCase    ,			image_embeds.T			) * logit_scale
							snake_case__      :   Union[str, Any]      =    logits_per_text.T

							if not return_dict:
											return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

							return FlaxCLIPOutput(
							    logits_per_image=__UpperCamelCase    ,			logits_per_text=__UpperCamelCase    ,			text_embeds=__UpperCamelCase    ,			image_embeds=__UpperCamelCase    ,			text_model_output=__UpperCamelCase    ,			vision_model_output=__UpperCamelCase    ,			)






class 			_SCREAMING_SNAKE_CASE				(lowercase__      ):
			A__								=  HybridCLIPConfig
			A__								=  FlaxHybridCLIPModule



			def __init__(     self		:						List[str]    ,			__UpperCamelCase		:						HybridCLIPConfig    ,			__UpperCamelCase		:						Optional[Tuple] = None    ,			__UpperCamelCase		:						int = 0    ,			__UpperCamelCase		:						jnp.dtype = jnp.floataa    ,			**__UpperCamelCase		:						Optional[Any]    ,			)      ->      Optional[Any]:
							"""simple docstring"""





							if input_shape is None:
											snake_case__      :   Any      =    ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

							snake_case__      :   int      =    self.module_class(config=__UpperCamelCase    ,			dtype=__UpperCamelCase    ,			**__UpperCamelCase			)
							super().__init__(__UpperCamelCase    ,			__UpperCamelCase    ,			input_shape=__UpperCamelCase    ,			seed=__UpperCamelCase    ,			dtype=__UpperCamelCase			)



			def 				lowerCAmelCase (     self		:						int    ,			__UpperCamelCase		:						jax.random.PRNGKey    ,			__UpperCamelCase		:						Tuple    ,			__UpperCamelCase		:						FrozenDict = None			)      ->      FrozenDict:
							"""simple docstring"""





							snake_case__      :   Dict      =    jnp.zeros(input_shape[0]    ,			dtype='''i4'''			)
							snake_case__      :   Any      =    jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase			).shape[-1]			)    ,			input_shape[0]			)
							snake_case__      :   str      =    jnp.ones_like(__UpperCamelCase			)
							snake_case__      :   List[str]      =    jnp.ones_like(__UpperCamelCase			)

							snake_case__      :   List[Any]      =    jax.random.normal(__UpperCamelCase    ,			input_shape[1]			)

							snake_case__  ,      snake_case__      :   Dict      =    jax.random.split(__UpperCamelCase			)
							snake_case__      :   Optional[Any]      =    {'''params''': params_rng, '''dropout''': dropout_rng}

							return self.module.init(__UpperCamelCase    ,			__UpperCamelCase    ,			__UpperCamelCase    ,			__UpperCamelCase    ,			__UpperCamelCase    ,			__UpperCamelCase			)["params"]



			def __call__(     self		:						Optional[Any]    ,			__UpperCamelCase		:						Union[str, Any]    ,			__UpperCamelCase		:						Optional[int]    ,			__UpperCamelCase		:						List[str]=None    ,			__UpperCamelCase		:						Optional[Any]=None    ,			__UpperCamelCase		:						Union[str, Any]=None    ,			__UpperCamelCase		:						dict = None    ,			__UpperCamelCase		:						jax.random.PRNGKey = None    ,			__UpperCamelCase		:						bool = False    ,			__UpperCamelCase		:						Optional[bool] = None    ,			__UpperCamelCase		:						Optional[bool] = None    ,			__UpperCamelCase		:						Optional[bool] = None    ,			)      ->      List[str]:
							"""simple docstring"""





							snake_case__      :   str      =    output_attentions if output_attentions is not None else self.config.output_attentions
							snake_case__      :   List[str]      =    (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							snake_case__      :   Dict      =    return_dict if return_dict is not None else self.config.return_dict

							if position_ids is None:
											snake_case__      :   Optional[int]      =    jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase			).shape[-1]			)    ,			input_ids.shape			)

							if token_type_ids is None:
											snake_case__      :   List[Any]      =    jnp.zeros_like(__UpperCamelCase			)

							if attention_mask is None:
											snake_case__      :   Tuple      =    jnp.ones_like(__UpperCamelCase			)

							# Handle any PRNG if needed
							snake_case__      :   Tuple      =    {}
							if dropout_rng is not None:
											snake_case__      :   Optional[int]      =    dropout_rng

							return self.module.apply(
							    {'''params''': params or self.params}    ,			jnp.array(__UpperCamelCase    ,			dtype='''i4'''			)    ,			jnp.array(__UpperCamelCase    ,			dtype=jnp.floataa			)    ,			jnp.array(__UpperCamelCase    ,			dtype='''i4'''			)    ,			jnp.array(__UpperCamelCase    ,			dtype='''i4'''			)    ,			jnp.array(__UpperCamelCase    ,			dtype='''i4'''			)    ,			not train    ,			__UpperCamelCase    ,			__UpperCamelCase    ,			__UpperCamelCase    ,			rngs=__UpperCamelCase    ,			)



			def 				lowerCAmelCase (     self		:						List[Any]    ,			__UpperCamelCase		:						Optional[int]    ,			__UpperCamelCase		:						List[Any]=None    ,			__UpperCamelCase		:						int=None    ,			__UpperCamelCase		:						List[str]=None    ,			__UpperCamelCase		:						dict = None    ,			__UpperCamelCase		:						jax.random.PRNGKey = None    ,			__UpperCamelCase		:						str=False    ,			)      ->      Union[str, Any]:
							"""simple docstring"""





							if position_ids is None:
											snake_case__      :   Optional[Any]      =    jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase			).shape[-1]			)    ,			input_ids.shape			)

							if token_type_ids is None:
											snake_case__      :   Optional[Any]      =    jnp.zeros_like(__UpperCamelCase			)

							if attention_mask is None:
											snake_case__      :   Union[str, Any]      =    jnp.ones_like(__UpperCamelCase			)

							# Handle any PRNG if needed
							snake_case__      :   Any      =    {}
							if dropout_rng is not None:
											snake_case__      :   List[Any]      =    dropout_rng

							def _get_features(__UpperCamelCase		:						str    ,			__UpperCamelCase		:						List[str]    ,			__UpperCamelCase		:						Dict    ,			__UpperCamelCase		:						Union[str, Any]    ,			__UpperCamelCase		:						Union[str, Any]    ,			__UpperCamelCase		:						Optional[Any]			):
											snake_case__      :   Tuple      =    module.text_model(
											    input_ids=__UpperCamelCase    ,			attention_mask=__UpperCamelCase    ,			position_ids=__UpperCamelCase    ,			token_type_ids=__UpperCamelCase    ,			deterministic=__UpperCamelCase    ,			)
											snake_case__      :   str      =    text_outputs[1]
											snake_case__      :   str      =    module.text_projection(__UpperCamelCase			)
											return text_features

							return self.module.apply(
							    {'''params''': params or self.params}    ,			jnp.array(__UpperCamelCase    ,			dtype='''i4'''			)    ,			jnp.array(__UpperCamelCase    ,			dtype='''i4'''			)    ,			jnp.array(__UpperCamelCase    ,			dtype='''i4'''			)    ,			jnp.array(__UpperCamelCase    ,			dtype='''i4'''			)    ,			not train    ,			method=_get_features    ,			rngs=__UpperCamelCase    ,			)



			def 				lowerCAmelCase (     self		:						List[str]    ,			__UpperCamelCase		:						Any    ,			__UpperCamelCase		:						dict = None    ,			__UpperCamelCase		:						jax.random.PRNGKey = None    ,			__UpperCamelCase		:						Optional[int]=False			)      ->      int:
							"""simple docstring"""





							snake_case__      :   List[str]      =    {}
							if dropout_rng is not None:
											snake_case__      :   Optional[Any]      =    dropout_rng

							def _get_features(__UpperCamelCase		:						Any    ,			__UpperCamelCase		:						int    ,			__UpperCamelCase		:						Any			):
											snake_case__      :   List[Any]      =    module.vision_model(pixel_values=__UpperCamelCase    ,			deterministic=__UpperCamelCase			)
											snake_case__      :   List[Any]      =    vision_outputs[1]  # pooled_output
											snake_case__      :   int      =    module.visual_projection(__UpperCamelCase			)
											return image_features

							return self.module.apply(
							    {'''params''': params or self.params}    ,			jnp.array(__UpperCamelCase    ,			dtype=jnp.floataa			)    ,			not train    ,			method=_get_features    ,			rngs=__UpperCamelCase    ,			)







			@classmethod
			def 				lowerCAmelCase (     cls		:						str    ,			__UpperCamelCase		:						str = None    ,			__UpperCamelCase		:						str = None    ,			*__UpperCamelCase		:						Tuple    ,			**__UpperCamelCase		:						List[Any]    ,			)      ->      FlaxPreTrainedModel:
							"""simple docstring"""





							snake_case__      :   Optional[Any]      =    {
							    argument[len('''text_'''			) :]: value for argument, value in kwargs.items() if argument.startswith('''text_'''			)
							}

							snake_case__      :   str      =    {
							    argument[len('''vision_'''			) :]: value for argument, value in kwargs.items() if argument.startswith('''vision_'''			)
							}

							# remove text, vision kwargs from kwargs
							for key in kwargs_text.keys():
											del kwargs["text_" + key]
							for key in kwargs_vision.keys():
											del kwargs["vision_" + key]

							# Load and initialize the text and vision model
							snake_case__      :   Tuple      =    kwargs_text.pop('''model'''    ,			__UpperCamelCase			)
							if text_model is None:
											assert (
											    text_model_name_or_path is not None
											), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
											from transformers import FlaxAutoModel

											if "config" not in kwargs_text:
															from transformers import AutoConfig

															snake_case__      :   str      =    AutoConfig.from_pretrained(__UpperCamelCase			)
															snake_case__      :   Optional[int]      =    text_config

											snake_case__      :   Any      =    FlaxAutoModel.from_pretrained(__UpperCamelCase    ,			*__UpperCamelCase    ,			**__UpperCamelCase			)

							snake_case__      :   Optional[int]      =    kwargs_vision.pop('''model'''    ,			__UpperCamelCase			)
							if vision_model is None:
											assert (
											    vision_model_name_or_path is not None
											), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
											from transformers import FlaxAutoModel

											if "config" not in kwargs_vision:
															from transformers import AutoConfig

															snake_case__      :   Optional[int]      =    AutoConfig.from_pretrained(__UpperCamelCase			)
															snake_case__      :   Any      =    vision_config

											snake_case__      :   List[str]      =    FlaxAutoModel.from_pretrained(__UpperCamelCase    ,			*__UpperCamelCase    ,			**__UpperCamelCase			)

							# instantiate config with corresponding kwargs
							snake_case__      :   str      =    kwargs.pop('''dtype'''    ,			jnp.floataa			)
							snake_case__      :   List[Any]      =    HybridCLIPConfig.from_text_vision_configs(text_model.config    ,			vision_model.config    ,			**__UpperCamelCase			)

							# init model
							snake_case__      :   Optional[int]      =    cls(__UpperCamelCase    ,			*__UpperCamelCase    ,			dtype=__UpperCamelCase    ,			**__UpperCamelCase			)

							if vision_config.model_type == "clip":
											snake_case__      :   Tuple      =    vision_model.params['''vision_model''']
											snake_case__      :   int      =    vision_model.params['''visual_projection''']['''kernel''']
							else:
											snake_case__      :   str      =    vision_model.params

							snake_case__      :   Any      =    text_model.params

							return model







import argparse

import torch
from clip import load

from transformers import CLIPConfig, CLIPModel






def        lowerCamelCase_  (				UpperCAmelCase_				:  Optional[int]				,							UpperCAmelCase_				:  Tuple						)    ->					str:




	'''simple docstring'''




	_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase :		Tuple								=    pt_attn_layer.in_proj_weight.chunk(3				,							dim=0						)
	_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase :		str								=    pt_attn_layer.in_proj_bias.chunk(3				,							dim=0						)

	_UpperCamelCase :		Dict								=    pt_attn_layer.out_proj.weight
	_UpperCamelCase :		Tuple								=    pt_attn_layer.out_proj.bias

	_UpperCamelCase :		Tuple								=    q_proj
	_UpperCamelCase :		Dict								=    q_proj_bias

	_UpperCamelCase :		Dict								=    k_proj
	_UpperCamelCase :		str								=    k_proj_bias

	_UpperCamelCase :		Dict								=    v_proj
	_UpperCamelCase :		Optional[int]								=    v_proj_bias

	_UpperCamelCase :		Any								=    out_proj_weights
	_UpperCamelCase :		Optional[int]								=    out_proj_bias






def        lowerCamelCase_  (				UpperCAmelCase_				:  Optional[int]				,							UpperCAmelCase_				:  Union[str, Any]						)    ->					List[Any]:




	'''simple docstring'''




	copy_linear(hf_mlp.fca				,							pt_mlp.c_fc						)
	copy_linear(hf_mlp.fca				,							pt_mlp.c_proj						)






def        lowerCamelCase_  (				UpperCAmelCase_				:  List[str]				,							UpperCAmelCase_				:  List[str]						)    ->					List[Any]:




	'''simple docstring'''




	_UpperCamelCase :		int								=    pt_linear.weight
	_UpperCamelCase :		int								=    pt_linear.bias






def        lowerCamelCase_  (				UpperCAmelCase_				:  str				,							UpperCAmelCase_				:  Union[str, Any]						)    ->					Dict:




	'''simple docstring'''




	copy_linear(hf_layer.layer_norma				,							pt_layer.ln_a						)
	copy_linear(hf_layer.layer_norma				,							pt_layer.ln_a						)

	# copy MLP
	copy_mlp(hf_layer.mlp				,							pt_layer.mlp						)

	# copy attn
	copy_attn_layer(hf_layer.self_attn				,							pt_layer.attn						)






def        lowerCamelCase_  (				UpperCAmelCase_				:  Optional[Any]				,							UpperCAmelCase_				:  List[str]						)    ->					Dict:




	'''simple docstring'''




	for hf_layer, pt_layer in zip(UpperCAmelCase_				,							UpperCAmelCase_						):
		copy_layer(UpperCAmelCase_				,							UpperCAmelCase_						)






def        lowerCamelCase_  (				UpperCAmelCase_				:  Any				,							UpperCAmelCase_				:  Optional[Any]						)    ->					int:




	'''simple docstring'''




	_UpperCamelCase :		Optional[int]								=    pt_model.token_embedding.weight
	_UpperCamelCase :		Tuple								=    pt_model.positional_embedding

	# copy layer norm
	copy_linear(hf_encoder.final_layer_norm				,							pt_model.ln_final						)

	# copy hidden layers
	copy_layers(hf_encoder.encoder.layers				,							pt_model.transformer.resblocks						)






def        lowerCamelCase_  (				UpperCAmelCase_				:  int				,							UpperCAmelCase_				:  Tuple						)    ->					Union[str, Any]:




	'''simple docstring'''




	_UpperCamelCase :		Any								=    pt_model.text_projection.data.T

	# copy text encoder
	copy_encoder(hf_model.text_model				,							UpperCAmelCase_						)






def        lowerCamelCase_  (				UpperCAmelCase_				:  str				,							UpperCAmelCase_				:  Dict						)    ->					Optional[int]:




	'''simple docstring'''




	_UpperCamelCase :		str								=    pt_model.visual.proj.data.T

	# copy layer norms
	copy_linear(hf_model.vision_model.pre_layrnorm				,							pt_model.visual.ln_pre						)
	copy_linear(hf_model.vision_model.post_layernorm				,							pt_model.visual.ln_post						)

	# copy embeds
	_UpperCamelCase :		Tuple								=    pt_model.visual.conva.weight.data
	_UpperCamelCase :		Dict								=    pt_model.visual.class_embedding
	_UpperCamelCase :		str								=    pt_model.visual.positional_embedding.data

	# copy encoder
	copy_layers(hf_model.vision_model.encoder.layers				,							pt_model.visual.transformer.resblocks						)






@torch.no_grad()
def        lowerCamelCase_  (				UpperCAmelCase_				:  Dict				,							UpperCAmelCase_				:  str				,							UpperCAmelCase_				:  List[str]=None						)    ->					Union[str, Any]:




	'''simple docstring'''




	if config_path is not None:
		_UpperCamelCase :		Dict								=    CLIPConfig.from_pretrained(UpperCAmelCase_						)
	else:
		_UpperCamelCase :		Tuple								=    CLIPConfig(projection_dim=5_1_2				,							text_config={}				,							vision_config={}						)

	_UpperCamelCase :		Tuple								=    CLIPModel(UpperCAmelCase_						).eval()

	_UpperCamelCase						,		_UpperCamelCase :		Optional[Any]								=    load(UpperCAmelCase_				,							device='cpu'				,							jit=UpperCAmelCase_						)
	_UpperCamelCase :		Optional[Any]								=    pt_model.eval()

	copy_text_model_and_projection(UpperCAmelCase_				,							UpperCAmelCase_						)
	copy_vison_model_and_projection(UpperCAmelCase_				,							UpperCAmelCase_						)
	_UpperCamelCase :		Optional[Any]								=    pt_model.logit_scale

	_UpperCamelCase :		Optional[int]								=    torch.arange(0				,							7_7						).unsqueeze(0						)
	_UpperCamelCase :		Dict								=    torch.randn(1				,							3				,							2_2_4				,							2_2_4						)

	_UpperCamelCase :		Optional[Any]								=    hf_model(input_ids=UpperCAmelCase_				,							pixel_values=UpperCAmelCase_				,							return_dict=UpperCAmelCase_						)
	_UpperCamelCase :		Optional[Any]								=    hf_outputs.logits_per_image
	_UpperCamelCase :		Optional[int]								=    hf_outputs.logits_per_text
	_UpperCamelCase						,		_UpperCamelCase :		int								=    pt_model(UpperCAmelCase_				,							UpperCAmelCase_						)

	assert torch.allclose(UpperCAmelCase_				,							UpperCAmelCase_				,							atol=1e-3						)
	assert torch.allclose(UpperCAmelCase_				,							UpperCAmelCase_				,							atol=1e-3						)

	hf_model.save_pretrained(UpperCAmelCase_						)


if __name__ == "__main__":
							lowerCAmelCase__             =		argparse.ArgumentParser()
							parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
							parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
							parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
							lowerCAmelCase__             =		parser.parse_args()

							convert_clip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)




