





'''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


_a					: Tuple      =    logging.get_logger(__name__)





class     _UpperCAmelCase							(						nn.Module   ):
		a   :   HybridCLIPConfig
		a   :   jnp.dtype     =jnp.floataa
		def   lowerCamelCase__			(  self			):



					'''simple docstring'''

					__lowerCAmelCase											=				self.config.text_config
					__lowerCAmelCase											=				self.config.vision_config

					__lowerCAmelCase											=				self.config.projection_dim
					__lowerCAmelCase											=				text_config.hidden_size
					__lowerCAmelCase											=				vision_config.hidden_size

					__lowerCAmelCase											=				FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
					__lowerCAmelCase											=				FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__,__SCREAMING_SNAKE_CASE			).module_class

					__lowerCAmelCase											=				text_module(__SCREAMING_SNAKE_CASE,dtype=self.dtype			)
					__lowerCAmelCase											=				vision_module(__SCREAMING_SNAKE_CASE,dtype=self.dtype			)

					__lowerCAmelCase											=				nn.Dense(
					    self.projection_dim,dtype=self.dtype,kernel_init=jax.nn.initializers.normal(0.02			),use_bias=__SCREAMING_SNAKE_CASE,)
					__lowerCAmelCase											=				nn.Dense(
					    self.projection_dim,dtype=self.dtype,kernel_init=jax.nn.initializers.normal(0.02			),use_bias=__SCREAMING_SNAKE_CASE,)
					__lowerCAmelCase											=				self.param("""logit_scale""",jax.nn.initializers.ones,[]			)







		def __call__(  self,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE = True,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=None,):



					'''simple docstring'''

					__lowerCAmelCase											=				return_dict if return_dict is not None else self.config.return_dict

					__lowerCAmelCase											=				self.vision_model(
					    pixel_values=__SCREAMING_SNAKE_CASE,deterministic=__SCREAMING_SNAKE_CASE,output_attentions=__SCREAMING_SNAKE_CASE,output_hidden_states=__SCREAMING_SNAKE_CASE,return_dict=__SCREAMING_SNAKE_CASE,)

					__lowerCAmelCase											=				self.text_model(
					    input_ids=__SCREAMING_SNAKE_CASE,attention_mask=__SCREAMING_SNAKE_CASE,token_type_ids=__SCREAMING_SNAKE_CASE,position_ids=__SCREAMING_SNAKE_CASE,deterministic=__SCREAMING_SNAKE_CASE,output_attentions=__SCREAMING_SNAKE_CASE,output_hidden_states=__SCREAMING_SNAKE_CASE,return_dict=__SCREAMING_SNAKE_CASE,)

					__lowerCAmelCase											=				vision_outputs[1]
					__lowerCAmelCase											=				self.visual_projection(__SCREAMING_SNAKE_CASE			)

					__lowerCAmelCase											=				text_outputs[1]
					__lowerCAmelCase											=				self.text_projection(__SCREAMING_SNAKE_CASE			)

					# normalized features
					__lowerCAmelCase											=				image_embeds / jnp.linalg.norm(__SCREAMING_SNAKE_CASE,axis=-1,keepdims=__SCREAMING_SNAKE_CASE			)
					__lowerCAmelCase											=				text_embeds / jnp.linalg.norm(__SCREAMING_SNAKE_CASE,axis=-1,keepdims=__SCREAMING_SNAKE_CASE			)

					# cosine similarity as logits
					__lowerCAmelCase											=				jnp.exp(self.logit_scale			)
					__lowerCAmelCase											=				jnp.matmul(__SCREAMING_SNAKE_CASE,image_embeds.T			) * logit_scale
					__lowerCAmelCase											=				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=__SCREAMING_SNAKE_CASE,logits_per_text=__SCREAMING_SNAKE_CASE,text_embeds=__SCREAMING_SNAKE_CASE,image_embeds=__SCREAMING_SNAKE_CASE,text_model_output=__SCREAMING_SNAKE_CASE,vision_model_output=__SCREAMING_SNAKE_CASE,)





class     _UpperCAmelCase							(						lowerCAmelCase_   ):
		a   :   Dict     =HybridCLIPConfig
		a   :   int     =FlaxHybridCLIPModule
		def __init__(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = 0,__SCREAMING_SNAKE_CASE = jnp.floataa,**__SCREAMING_SNAKE_CASE,):



					'''simple docstring'''

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

					__lowerCAmelCase											=				self.module_class(config=__SCREAMING_SNAKE_CASE,dtype=__SCREAMING_SNAKE_CASE,**__SCREAMING_SNAKE_CASE			)
					super().__init__(__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,input_shape=__SCREAMING_SNAKE_CASE,seed=__SCREAMING_SNAKE_CASE,dtype=__SCREAMING_SNAKE_CASE			)
		def   lowerCamelCase__			(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE = None			):



					'''simple docstring'''

					__lowerCAmelCase											=				jnp.zeros(input_shape[0],dtype="""i4"""			)
					__lowerCAmelCase											=				jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__SCREAMING_SNAKE_CASE			).shape[-1]			),input_shape[0]			)
					__lowerCAmelCase											=				jnp.ones_like(__SCREAMING_SNAKE_CASE			)
					__lowerCAmelCase											=				jnp.ones_like(__SCREAMING_SNAKE_CASE			)

					__lowerCAmelCase											=				jax.random.normal(__SCREAMING_SNAKE_CASE,input_shape[1]			)

					__lowerCAmelCase						,       __lowerCAmelCase											=				jax.random.split(__SCREAMING_SNAKE_CASE			)
					__lowerCAmelCase											=				{"""params""": params_rng, """dropout""": dropout_rng}

					return self.module.init(__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			)["params"]
		def __call__(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = False,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = None,):



					'''simple docstring'''

					__lowerCAmelCase											=				output_attentions if output_attentions is not None else self.config.output_attentions
					__lowerCAmelCase											=				(
					    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
					)
					__lowerCAmelCase											=				return_dict if return_dict is not None else self.config.return_dict

					if position_ids is None:
								__lowerCAmelCase											=				jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__SCREAMING_SNAKE_CASE			).shape[-1]			),input_ids.shape			)

					if token_type_ids is None:
								__lowerCAmelCase											=				jnp.zeros_like(__SCREAMING_SNAKE_CASE			)

					if attention_mask is None:
								__lowerCAmelCase											=				jnp.ones_like(__SCREAMING_SNAKE_CASE			)

					# Handle any PRNG if needed
					__lowerCAmelCase											=				{}
					if dropout_rng is not None:
								__lowerCAmelCase											=				dropout_rng

					return self.module.apply(
					    {"""params""": params or self.params},jnp.array(__SCREAMING_SNAKE_CASE,dtype="""i4"""			),jnp.array(__SCREAMING_SNAKE_CASE,dtype=jnp.floataa			),jnp.array(__SCREAMING_SNAKE_CASE,dtype="""i4"""			),jnp.array(__SCREAMING_SNAKE_CASE,dtype="""i4"""			),jnp.array(__SCREAMING_SNAKE_CASE,dtype="""i4"""			),not train,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,rngs=__SCREAMING_SNAKE_CASE,)
		def   lowerCamelCase__			(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE=False,):



					'''simple docstring'''

					if position_ids is None:
								__lowerCAmelCase											=				jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__SCREAMING_SNAKE_CASE			).shape[-1]			),input_ids.shape			)

					if token_type_ids is None:
								__lowerCAmelCase											=				jnp.zeros_like(__SCREAMING_SNAKE_CASE			)

					if attention_mask is None:
								__lowerCAmelCase											=				jnp.ones_like(__SCREAMING_SNAKE_CASE			)

					# Handle any PRNG if needed
					__lowerCAmelCase											=				{}
					if dropout_rng is not None:
								__lowerCAmelCase											=				dropout_rng

					def _get_features(__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			):
								__lowerCAmelCase											=				module.text_model(
								    input_ids=__SCREAMING_SNAKE_CASE,attention_mask=__SCREAMING_SNAKE_CASE,position_ids=__SCREAMING_SNAKE_CASE,token_type_ids=__SCREAMING_SNAKE_CASE,deterministic=__SCREAMING_SNAKE_CASE,)
								__lowerCAmelCase											=				text_outputs[1]
								__lowerCAmelCase											=				module.text_projection(__SCREAMING_SNAKE_CASE			)
								return text_features

					return self.module.apply(
					    {"""params""": params or self.params},jnp.array(__SCREAMING_SNAKE_CASE,dtype="""i4"""			),jnp.array(__SCREAMING_SNAKE_CASE,dtype="""i4"""			),jnp.array(__SCREAMING_SNAKE_CASE,dtype="""i4"""			),jnp.array(__SCREAMING_SNAKE_CASE,dtype="""i4"""			),not train,method=_get_features,rngs=__SCREAMING_SNAKE_CASE,)
		def   lowerCamelCase__			(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE=False			):



					'''simple docstring'''

					__lowerCAmelCase											=				{}
					if dropout_rng is not None:
								__lowerCAmelCase											=				dropout_rng

					def _get_features(__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			):
								__lowerCAmelCase											=				module.vision_model(pixel_values=__SCREAMING_SNAKE_CASE,deterministic=__SCREAMING_SNAKE_CASE			)
								__lowerCAmelCase											=				vision_outputs[1]  # pooled_output
								__lowerCAmelCase											=				module.visual_projection(__SCREAMING_SNAKE_CASE			)
								return image_features

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







		@classmethod
		def   lowerCamelCase__			(  cls,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = None,*__SCREAMING_SNAKE_CASE,**__SCREAMING_SNAKE_CASE,):



					'''simple docstring'''

					__lowerCAmelCase											=				{
					    argument[len("""text_"""			) :]: value for argument, value in kwargs.items() if argument.startswith("""text_"""			)
					}

					__lowerCAmelCase											=				{
					    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
					__lowerCAmelCase											=				kwargs_text.pop("""model""",__SCREAMING_SNAKE_CASE			)
					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

											__lowerCAmelCase											=				AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE			)
											__lowerCAmelCase											=				text_config

								__lowerCAmelCase											=				FlaxAutoModel.from_pretrained(__SCREAMING_SNAKE_CASE,*__SCREAMING_SNAKE_CASE,**__SCREAMING_SNAKE_CASE			)

					__lowerCAmelCase											=				kwargs_vision.pop("""model""",__SCREAMING_SNAKE_CASE			)
					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

											__lowerCAmelCase											=				AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE			)
											__lowerCAmelCase											=				vision_config

								__lowerCAmelCase											=				FlaxAutoModel.from_pretrained(__SCREAMING_SNAKE_CASE,*__SCREAMING_SNAKE_CASE,**__SCREAMING_SNAKE_CASE			)

					# instantiate config with corresponding kwargs
					__lowerCAmelCase											=				kwargs.pop("""dtype""",jnp.floataa			)
					__lowerCAmelCase											=				HybridCLIPConfig.from_text_vision_configs(text_model.config,vision_model.config,**__SCREAMING_SNAKE_CASE			)

					# init model
					__lowerCAmelCase											=				cls(__SCREAMING_SNAKE_CASE,*__SCREAMING_SNAKE_CASE,dtype=__SCREAMING_SNAKE_CASE,**__SCREAMING_SNAKE_CASE			)

					if vision_config.model_type == "clip":
								__lowerCAmelCase											=				vision_model.params["""vision_model"""]
								__lowerCAmelCase											=				vision_model.params["""visual_projection"""]["""kernel"""]
					else:
								__lowerCAmelCase											=				vision_model.params

					__lowerCAmelCase											=				text_model.params

					return model








'''simple docstring'''







from collections import namedtuple

_a					: Any      =    namedtuple("""from_to""", """from_ to""")

_a					: Dict      =    {
    """millimeter""": """mm""",
    """centimeter""": """cm""",
    """meter""": """m""",
    """kilometer""": """km""",
    """inch""": """in""",
    """inche""": """in""",  # Trailing 's' has been stripped off
    """feet""": """ft""",
    """foot""": """ft""",
    """yard""": """yd""",
    """mile""": """mi""",
}

_a					: Optional[Any]      =    {
    """mm""": from_to(0.001, 1_0_0_0),
    """cm""": from_to(0.01, 1_0_0),
    """m""": from_to(1, 1),
    """km""": from_to(1_0_0_0, 0.001),
    """in""": from_to(0.0254, 39.3701),
    """ft""": from_to(0.3048, 3.28084),
    """yd""": from_to(0.9144, 1.09361),
    """mi""": from_to(1609.34, 0.000621371),
}





def    _lowerCAmelCase						(						lowercase					,	lowercase					,	lowercase	)					->					float:
			__lowerCAmelCase											=				from_type.lower().rstrip("""s"""	)
			__lowerCAmelCase											=				TYPE_CONVERSION.get(lowercase					,	lowercase	)
			__lowerCAmelCase											=				to_type.lower().rstrip("""s"""	)
			__lowerCAmelCase											=				TYPE_CONVERSION.get(lowercase					,	lowercase	)
			if new_from not in METRIC_CONVERSION:
						__lowerCAmelCase											=				(
						    f'Invalid \'from_type\' value: {from_type!r}.\n'
						    f'Conversion abbreviations are: {", ".join(lowercase	)}'
						)
						raise ValueError(lowercase	)
			if new_to not in METRIC_CONVERSION:
						__lowerCAmelCase											=				(
						    f'Invalid \'to_type\' value: {to_type!r}.\n'
						    f'Conversion abbreviations are: {", ".join(lowercase	)}'
						)
						raise ValueError(lowercase	)
			return value * METRIC_CONVERSION[new_from].from_ * METRIC_CONVERSION[new_to].to


if __name__ == "__main__":
	import doctest

	doctest.testmod()

