




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


_lowerCamelCase							:		Optional[int]										=						logging.get_logger(__name__)





class   lowerCamelCase_							( nn.Module							):




					'''simple docstring'''






					lowercase__					=				42
					lowercase__					=				jnp.floataa






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = self.config.text_config
							_snake_case			 = self.config.vision_config

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

							_snake_case			 = FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
							_snake_case			 = FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__,						__UpperCamelCase				).module_class

							_snake_case			 = text_module(__UpperCamelCase,						dtype=self.dtype				)
							_snake_case			 = vision_module(__UpperCamelCase,						dtype=self.dtype				)

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






					def __call__(					self,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase = True,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						):
							_snake_case			 = return_dict if return_dict is not None else self.config.return_dict

							_snake_case			 = self.vision_model(
							    pixel_values=__UpperCamelCase,						deterministic=__UpperCamelCase,						output_attentions=__UpperCamelCase,						output_hidden_states=__UpperCamelCase,						return_dict=__UpperCamelCase,						)

							_snake_case			 = 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			 = vision_outputs[1]
							_snake_case			 = self.visual_projection(__UpperCamelCase				)

							_snake_case			 = text_outputs[1]
							_snake_case			 = self.text_projection(__UpperCamelCase				)

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

							# cosine similarity as logits
							_snake_case			 = jnp.exp(self.logit_scale				)
							_snake_case			 = jnp.matmul(__UpperCamelCase,						image_embeds.T				) * logit_scale
							_snake_case			 = 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   lowerCamelCase_							( UpperCAmelCase_							):




					'''simple docstring'''






					lowercase__					=				HybridCLIPConfig
					lowercase__					=				FlaxHybridCLIPModule






					def __init__(					self,						__UpperCamelCase,						__UpperCamelCase = None,						__UpperCamelCase = 0,						__UpperCamelCase = jnp.floataa,						**__UpperCamelCase,						):
							if input_shape is None:
									_snake_case			 = ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

							_snake_case			 = self.module_class(config=__UpperCamelCase,						dtype=__UpperCamelCase,						**__UpperCamelCase				)
							super().__init__(__UpperCamelCase,						__UpperCamelCase,						input_shape=__UpperCamelCase,						seed=__UpperCamelCase,						dtype=__UpperCamelCase				)






					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase = None				):
							# init input tensor
							_snake_case			 = jnp.zeros(input_shape[0],						dtype="i4"				)
							_snake_case			 = jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase				).shape[-1]				),						input_shape[0]				)
							_snake_case			 = jnp.ones_like(__UpperCamelCase				)
							_snake_case			 = jnp.ones_like(__UpperCamelCase				)

							_snake_case			 = jax.random.normal(__UpperCamelCase,						input_shape[1]				)

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

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






					def __call__(					self,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase = None,						__UpperCamelCase = None,						__UpperCamelCase = False,						__UpperCamelCase = None,						__UpperCamelCase = None,						__UpperCamelCase = None,						):
							_snake_case			 = output_attentions if output_attentions is not None else self.config.output_attentions
							_snake_case			 = (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							_snake_case			 = return_dict if return_dict is not None else self.config.return_dict

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

							if token_type_ids is None:
									_snake_case			 = jnp.zeros_like(__UpperCamelCase				)

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

							# Handle any PRNG if needed
							_snake_case			 = {}
							if dropout_rng is not None:
									_snake_case			 = 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        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase = None,						__UpperCamelCase = None,						__UpperCamelCase=False,						):
							if position_ids is None:
									_snake_case			 = jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase				).shape[-1]				),						input_ids.shape				)

							if token_type_ids is None:
									_snake_case			 = jnp.zeros_like(__UpperCamelCase				)

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

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

							def _get_features(__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase				):
									_snake_case			 = module.text_model(
									    input_ids=__UpperCamelCase,						attention_mask=__UpperCamelCase,						position_ids=__UpperCamelCase,						token_type_ids=__UpperCamelCase,						deterministic=__UpperCamelCase,						)
									_snake_case			 = text_outputs[1]
									_snake_case			 = 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        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase = None,						__UpperCamelCase = None,						__UpperCamelCase=False				):
							_snake_case			 = {}
							if dropout_rng is not None:
									_snake_case			 = dropout_rng

							def _get_features(__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase				):
									_snake_case			 = module.vision_model(pixel_values=__UpperCamelCase,						deterministic=__UpperCamelCase				)
									_snake_case			 = vision_outputs[1]  # pooled_output
									_snake_case			 = 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        __SCREAMING_SNAKE_CASE			(					cls,						__UpperCamelCase = None,						__UpperCamelCase = None,						*__UpperCamelCase,						**__UpperCamelCase,						):
							_snake_case			 = {
							    argument[len("text_"				) :]: value for argument, value in kwargs.items() if argument.startswith("text_"				)
							}

							_snake_case			 = {
							    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			 = 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			 = AutoConfig.from_pretrained(__UpperCamelCase				)
											_snake_case			 = text_config

									_snake_case			 = FlaxAutoModel.from_pretrained(__UpperCamelCase,						*__UpperCamelCase,						**__UpperCamelCase				)

							_snake_case			 = 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			 = AutoConfig.from_pretrained(__UpperCamelCase				)
											_snake_case			 = vision_config

									_snake_case			 = FlaxAutoModel.from_pretrained(__UpperCamelCase,						*__UpperCamelCase,						**__UpperCamelCase				)

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

							# init model
							_snake_case			 = cls(__UpperCamelCase,						*__UpperCamelCase,						dtype=__UpperCamelCase,						**__UpperCamelCase				)

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

							_snake_case			 = text_model.params

							return model











'''simple docstring'''







from __future__ import annotations

from collections.abc import Iterable, Iterator
from typing import Any, Generic, TypeVar

_lowerCamelCase							:		List[str]										=						TypeVar('T', bound=bool)





class   lowerCamelCase_							( Generic[T]							):




					'''simple docstring'''






					def __init__(					self,						__UpperCamelCase				):
							_snake_case			 = value
							_snake_case			 = None
							_snake_case			 = None






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							return self._value






					@staticmethod
					def        __SCREAMING_SNAKE_CASE			(					__UpperCamelCase,						__UpperCamelCase				):
							if not roota:
									return roota

							if not roota:
									return roota

							if roota.value > roota.value:
									_snake_case ,					_snake_case			 = roota, roota

							_snake_case			 = roota
							_snake_case			 = roota.right
							_snake_case			 = roota.left
							_snake_case			 = SkewNode.merge(__UpperCamelCase,						__UpperCamelCase				)

							return result





class   lowerCamelCase_							( Generic[T]							):




					'''simple docstring'''






					def __init__(					self,						__UpperCamelCase = ()				):
							_snake_case			 = None
							if data:
									for item in data:
											self.insert(__UpperCamelCase				)






					def __bool__(					self				):
							return self._root is not None






					def __iter__(					self				):
							_snake_case			 = []
							while self:
									result.append(self.pop()				)

							# Pushing items back to the heap not to clear it.
							for item in result:
									self.insert(__UpperCamelCase				)

							return iter(__UpperCamelCase				)






					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase				):
							_snake_case			 = SkewNode.merge(self._root,						SkewNode(__UpperCamelCase				)				)






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = self.top()
							_snake_case			 = (
							    SkewNode.merge(self._root.left,						self._root.right				) if self._root else None
							)

							return result






					def        __SCREAMING_SNAKE_CASE			(					self				):
							if not self._root:
									raise IndexError("Can't get top element for the empty heap."				)
							return self._root.value






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = None


if __name__ == "__main__":
						import doctest

						doctest.testmod()





