





'''simple docstring'''







from collections import Counter
from functools import total_ordering

from data_structures.heap.heap import Heap


@total_ordering
class        __a   :

							def __init__(				self				:  Union[str, Any]							,  __magic_name__				:  str							,  __magic_name__				:  int )       ->							None:







													"""simple docstring"""
													UpperCAmelCase_      :   Union[str, Any]      =	word
													UpperCAmelCase_      :   Optional[Any]      =	count

							def __eq__(				self				:  Dict							,  __magic_name__				:  object )       ->							bool:







													"""simple docstring"""
													if not isinstance(__magic_name__							,  __magic_name__ ):
																			return NotImplemented
													return self.count == other.count



							def __lt__(				self				:  List[str]							,  __magic_name__				:  object )       ->							bool:







													"""simple docstring"""
													if not isinstance(__magic_name__							,  __magic_name__ ):
																			return NotImplemented
													return self.count < other.count





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      list[str],      SCREAMING_SNAKE_CASE__   :      int )			->						list[str]:
						UpperCAmelCase_      :   Heap[WordCount]      =	Heap()
						UpperCAmelCase_      :   int      =	Counter(SCREAMING_SNAKE_CASE__ )
						heap.build_max_heap(
						    [WordCount(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ ) for word, count in count_by_word.items()] )
						return [heap.extract_max().word for _ in range(min(SCREAMING_SNAKE_CASE__,      len(SCREAMING_SNAKE_CASE__ ) ) )]


if __name__ == "__main__":
					import doctest

					doctest.testmod()











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_      :							List[str]    							=					logging.get_logger(__name__)






class        __lowerCAmelCase					(			nn.Module					):
			snake_case_ :					HybridCLIPConfig
			snake_case_ :					jnp.dtype						=	jnp.floataa







			def UpperCamelCase     (			self       :     Union[str, Any] ):




									"""simple docstring"""


									_UpperCAmelCase						=	self.config.text_config
									_UpperCAmelCase						=	self.config.vision_config

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

									_UpperCAmelCase						=	FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
									_UpperCAmelCase						=	FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__       ,       snake_case__ ).module_class

									_UpperCAmelCase						=	text_module(snake_case__       ,       dtype=self.dtype )
									_UpperCAmelCase						=	vision_module(snake_case__       ,       dtype=self.dtype )

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







			def __call__(			self       :     List[Any]       ,       snake_case__       :     List[Any]=None       ,       snake_case__       :     List[str]=None       ,       snake_case__       :     Optional[Any]=None       ,       snake_case__       :     Dict=None       ,       snake_case__       :     Tuple=None       ,       snake_case__       :     bool = True       ,       snake_case__       :     Dict=None       ,       snake_case__       :     Dict=None       ,       snake_case__       :     Dict=None       ,       ):




									"""simple docstring"""


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

									_UpperCAmelCase						=	self.vision_model(
									    pixel_values=snake_case__       ,       deterministic=snake_case__       ,       output_attentions=snake_case__       ,       output_hidden_states=snake_case__       ,       return_dict=snake_case__       ,       )

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

									_UpperCAmelCase						=	vision_outputs[1]
									_UpperCAmelCase						=	self.visual_projection(snake_case__ )

									_UpperCAmelCase						=	text_outputs[1]
									_UpperCAmelCase						=	self.text_projection(snake_case__ )

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

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






class        __lowerCAmelCase					(			UpperCAmelCase__					):
			snake_case_ :					Union[str, Any]						=	HybridCLIPConfig
			snake_case_ :					Optional[Any]						=	FlaxHybridCLIPModule







			def __init__(			self       :     List[Any]       ,       snake_case__       :     HybridCLIPConfig       ,       snake_case__       :     Optional[Tuple] = None       ,       snake_case__       :     int = 0       ,       snake_case__       :     jnp.dtype = jnp.floataa       ,       **snake_case__       :     Any       ,       ):




									"""simple docstring"""


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

									_UpperCAmelCase						=	self.module_class(config=snake_case__       ,       dtype=snake_case__       ,       **snake_case__ )
									super().__init__(snake_case__       ,       snake_case__       ,       input_shape=snake_case__       ,       seed=snake_case__       ,       dtype=snake_case__ )







			def UpperCamelCase     (			self       :     Union[str, Any]       ,       snake_case__       :     jax.random.PRNGKey       ,       snake_case__       :     Tuple       ,       snake_case__       :     FrozenDict = None ):




									"""simple docstring"""


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

									_UpperCAmelCase						=	jax.random.normal(snake_case__       ,       input_shape[1] )

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

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







			def __call__(			self       :     Union[str, Any]       ,       snake_case__       :     Optional[int]       ,       snake_case__       :     Tuple       ,       snake_case__       :     Optional[int]=None       ,       snake_case__       :     Tuple=None       ,       snake_case__       :     Optional[Any]=None       ,       snake_case__       :     dict = None       ,       snake_case__       :     jax.random.PRNGKey = None       ,       snake_case__       :     bool = False       ,       snake_case__       :     Optional[bool] = None       ,       snake_case__       :     Optional[bool] = None       ,       snake_case__       :     Optional[bool] = None       ,       ):




									"""simple docstring"""


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

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

									if token_type_ids is None:
															_UpperCAmelCase						=	jnp.zeros_like(snake_case__ )

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

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

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







			def UpperCamelCase     (			self       :     int       ,       snake_case__       :     int       ,       snake_case__       :     List[str]=None       ,       snake_case__       :     Union[str, Any]=None       ,       snake_case__       :     str=None       ,       snake_case__       :     dict = None       ,       snake_case__       :     jax.random.PRNGKey = None       ,       snake_case__       :     int=False       ,       ):




									"""simple docstring"""


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

									if token_type_ids is None:
															_UpperCAmelCase						=	jnp.zeros_like(snake_case__ )

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

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

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

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







			def UpperCamelCase     (			self       :     List[Any]       ,       snake_case__       :     Optional[Any]       ,       snake_case__       :     dict = None       ,       snake_case__       :     jax.random.PRNGKey = None       ,       snake_case__       :     str=False ):




									"""simple docstring"""


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

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

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







			@classmethod
			def UpperCamelCase     (			cls       :     Tuple       ,       snake_case__       :     str = None       ,       snake_case__       :     str = None       ,       *snake_case__       :     Any       ,       **snake_case__       :     Tuple       ,       ):




									"""simple docstring"""


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

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

																					_UpperCAmelCase						=	AutoConfig.from_pretrained(snake_case__ )
																					_UpperCAmelCase						=	text_config

															_UpperCAmelCase						=	FlaxAutoModel.from_pretrained(snake_case__       ,       *snake_case__       ,       **snake_case__ )

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

																					_UpperCAmelCase						=	AutoConfig.from_pretrained(snake_case__ )
																					_UpperCAmelCase						=	vision_config

															_UpperCAmelCase						=	FlaxAutoModel.from_pretrained(snake_case__       ,       *snake_case__       ,       **snake_case__ )

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

									# init model
									_UpperCAmelCase						=	cls(snake_case__       ,       *snake_case__       ,       dtype=snake_case__       ,       **snake_case__ )

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

									_UpperCAmelCase						=	text_model.params

									return model




