



'''simple docstring'''





import copy
import os
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Union


if TYPE_CHECKING:
  from ...processing_utils import ProcessorMixin
  from ...utils import TensorType

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowercase_									=    logging.get_logger(__name__)

lowercase_									=    {
    "google/owlvit-base-patch32": "https://huggingface.co/google/owlvit-base-patch32/resolve/main/config.json",
    "google/owlvit-base-patch16": "https://huggingface.co/google/owlvit-base-patch16/resolve/main/config.json",
    "google/owlvit-large-patch14": "https://huggingface.co/google/owlvit-large-patch14/resolve/main/config.json",
}






class __A     (    A		):



 '''simple docstring'''







 __lowerCamelCase       :			str		=			'owlvit_text_model'



 def __init__(self					,  A=49_408					,  A=512					,  A=2_048					,  A=12					,  A=8					,  A=16					,  A="quick_gelu"					,  A=1E-5					,  A=0.0					,  A=0.02					,  A=1.0					,  A=0					,  A=49_406					,  A=49_407					,  **A					,  )       ->	Union[str, Any]:


      """simple docstring"""



      super().__init__(pad_token_id=A					,  bos_token_id=A					,  eos_token_id=A					,  **A			)

      _a           =    vocab_size
      _a           =    hidden_size
      _a           =    intermediate_size
      _a           =    num_hidden_layers
      _a           =    num_attention_heads
      _a           =    max_position_embeddings
      _a           =    hidden_act
      _a           =    layer_norm_eps
      _a           =    attention_dropout
      _a           =    initializer_range
      _a           =    initializer_factor



 @classmethod
 def    a__    (cls					,  A					,  **A			)       ->	"PretrainedConfig":


      """simple docstring"""



      cls._set_token_in_kwargs(A			)

      _a    ,       _a           =    cls.get_config_dict(A					,  **A			)

      # get the text config dict if we are loading from OwlViTConfig
      if config_dict.get('''model_type'''			) == "owlvit":
           _a           =    config_dict['''text_config''']

      if "model_type" in config_dict and hasattr(cls					,  '''model_type'''			) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
               f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''			)

      return cls.from_dict(A					,  **A			)






class __A     (    A		):



 '''simple docstring'''







 __lowerCamelCase       :			Dict		=			'owlvit_vision_model'



 def __init__(self					,  A=768					,  A=3_072					,  A=12					,  A=12					,  A=3					,  A=768					,  A=32					,  A="quick_gelu"					,  A=1E-5					,  A=0.0					,  A=0.02					,  A=1.0					,  **A					,  )       ->	Union[str, Any]:


      """simple docstring"""



      super().__init__(**A			)

      _a           =    hidden_size
      _a           =    intermediate_size
      _a           =    num_hidden_layers
      _a           =    num_attention_heads
      _a           =    num_channels
      _a           =    image_size
      _a           =    patch_size
      _a           =    hidden_act
      _a           =    layer_norm_eps
      _a           =    attention_dropout
      _a           =    initializer_range
      _a           =    initializer_factor



 @classmethod
 def    a__    (cls					,  A					,  **A			)       ->	"PretrainedConfig":


      """simple docstring"""



      cls._set_token_in_kwargs(A			)

      _a    ,       _a           =    cls.get_config_dict(A					,  **A			)

      # get the vision config dict if we are loading from OwlViTConfig
      if config_dict.get('''model_type'''			) == "owlvit":
           _a           =    config_dict['''vision_config''']

      if "model_type" in config_dict and hasattr(cls					,  '''model_type'''			) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
               f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''			)

      return cls.from_dict(A					,  **A			)






class __A     (    A		):



 '''simple docstring'''







 __lowerCamelCase       :			Optional[int]		=			'owlvit'
 __lowerCamelCase       :			List[str]		=			True



 def __init__(self					,  A=None					,  A=None					,  A=512					,  A=2.6592					,  A=True					,  **A					,  )       ->	Optional[int]:


      """simple docstring"""



      super().__init__(**A			)

      if text_config is None:
           _a           =    {}
           logger.info('''text_config is None. Initializing the OwlViTTextConfig with default values.'''			)

      if vision_config is None:
           _a           =    {}
           logger.info('''vision_config is None. initializing the OwlViTVisionConfig with default values.'''			)

      _a           =    OwlViTTextConfig(**A			)
      _a           =    OwlViTVisionConfig(**A			)

      _a           =    projection_dim
      _a           =    logit_scale_init_value
      _a           =    return_dict
      _a           =    1.0



 @classmethod
 def    a__    (cls					,  A					,  **A			)       ->	"PretrainedConfig":


      """simple docstring"""



      cls._set_token_in_kwargs(A			)

      _a    ,       _a           =    cls.get_config_dict(A					,  **A			)

      if "model_type" in config_dict and hasattr(cls					,  '''model_type'''			) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
               f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''			)

      return cls.from_dict(A					,  **A			)



 @classmethod
 def    a__    (cls					,  A					,  A					,  **A			)       ->	Any:


      """simple docstring"""



      _a           =    {}
      _a           =    text_config
      _a           =    vision_config

      return cls.from_dict(A					,  **A			)



 def    a__    (self			)       ->	Tuple:


      """simple docstring"""



      _a           =    copy.deepcopy(self.__dict__			)
      _a           =    self.text_config.to_dict()
      _a           =    self.vision_config.to_dict()
      _a           =    self.__class__.model_type
      return output






class __A     (    A		):



 '''simple docstring'''







 @property
 def    a__    (self			)       ->	Mapping[str, Mapping[int, str]]:


      """simple docstring"""



      return OrderedDict(
          [
              ('''input_ids''', {0: '''batch''', 1: '''sequence'''}),
              ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
              ('''attention_mask''', {0: '''batch''', 1: '''sequence'''}),
          ]			)



 @property
 def    a__    (self			)       ->	Mapping[str, Mapping[int, str]]:


      """simple docstring"""



      return OrderedDict(
          [
              ('''logits_per_image''', {0: '''batch'''}),
              ('''logits_per_text''', {0: '''batch'''}),
              ('''text_embeds''', {0: '''batch'''}),
              ('''image_embeds''', {0: '''batch'''}),
          ]			)



 @property
 def    a__    (self			)       ->	float:


      """simple docstring"""



      return 1E-4



 def    a__    (self					,  A					,  A = -1					,  A = -1					,  A = None					,  )       ->	Mapping[str, Any]:


      """simple docstring"""



      _a           =    super().generate_dummy_inputs(
          processor.tokenizer					,  batch_size=A					,  seq_length=A					,  framework=A			)
      _a           =    super().generate_dummy_inputs(
          processor.image_processor					,  batch_size=A					,  framework=A			)
      return {**text_input_dict, **image_input_dict}



 @property
 def    a__    (self			)       ->	int:


      """simple docstring"""



      return 14










import tensorflow as tf

from ...tf_utils import shape_list
class     UpperCamelCase__					(     tf.keras.layers.Layer				):






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__,   snake_case__,   snake_case__,   snake_case__=1,   snake_case__=False,   **snake_case__  )       -> Optional[int]:






									"""simple docstring"""
									super().__init__(**snake_case__  )

									lowercase_	:   Tuple  				=     vocab_size
									lowercase_	:   Union[str, Any]  				=     d_embed
									lowercase_	:   Optional[Any]  				=     d_proj

									lowercase_	:   Optional[Any]  				=     cutoffs + [vocab_size]
									lowercase_	:   Optional[int]  				=     [0] + self.cutoffs
									lowercase_	:   Union[str, Any]  				=     div_val

									lowercase_	:   Union[str, Any]  				=     self.cutoffs[0]
									lowercase_	:   List[str]  				=     len(self.cutoffs  ) - 1
									lowercase_	:   List[str]  				=     self.shortlist_size + self.n_clusters
									lowercase_	:   Tuple  				=     keep_order

									lowercase_	:   Dict  				=     []
									lowercase_	:   int  				=     []






		def 						snake_case__				(							self,   snake_case__  )       -> Optional[int]:






									"""simple docstring"""
									if self.n_clusters > 0:
																lowercase_	:   List[str]  				=     self.add_weight(
																    shape=(self.n_clusters, self.d_embed),   initializer="""zeros""",   trainable=snake_case__,   name="""cluster_weight"""  )
																lowercase_	:   List[Any]  				=     self.add_weight(
																    shape=(self.n_clusters,),   initializer="""zeros""",   trainable=snake_case__,   name="""cluster_bias"""  )

									if self.div_val == 1:
																for i in range(len(self.cutoffs  )  ):
																							if self.d_proj != self.d_embed:
																														lowercase_	:   Any  				=     self.add_weight(
																														    shape=(self.d_embed, self.d_proj),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_projs_._{i}""",   )
																														self.out_projs.append(snake_case__  )
																							else:
																														self.out_projs.append(snake_case__  )
																							lowercase_	:   List[str]  				=     self.add_weight(
																							    shape=(self.vocab_size, self.d_embed),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._weight""",   )
																							lowercase_	:   List[str]  				=     self.add_weight(
																							    shape=(self.vocab_size,),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._bias""",   )
																							self.out_layers.append((weight, bias)  )
									else:
																for i in range(len(self.cutoffs  )  ):
																							lowercase_				,		lowercase_	:   Optional[Any]  				=     self.cutoff_ends[i], self.cutoff_ends[i + 1]
																							lowercase_	:   Union[str, Any]  				=     self.d_embed // (self.div_val**i)

																							lowercase_	:   int  				=     self.add_weight(
																							    shape=(d_emb_i, self.d_proj),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_projs_._{i}"""  )
																							self.out_projs.append(snake_case__  )
																							lowercase_	:   Any  				=     self.add_weight(
																							    shape=(r_idx - l_idx, d_emb_i),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._weight""",   )
																							lowercase_	:   Tuple  				=     self.add_weight(
																							    shape=(r_idx - l_idx,),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._bias""",   )
																							self.out_layers.append((weight, bias)  )
									super().build(snake_case__  )






		@staticmethod
		def 						snake_case__				(							snake_case__,   snake_case__,   snake_case__,   snake_case__=None  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   Dict  				=     x
									if proj is not None:
																lowercase_	:   List[Any]  				=     tf.einsum("""ibd,ed->ibe""",   snake_case__,   snake_case__  )
									return tf.einsum("""ibd,nd->ibn""",   snake_case__,   snake_case__  ) + b






		@staticmethod
		def 						snake_case__				(							snake_case__,   snake_case__  )       -> List[str]:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     shape_list(snake_case__  )
									lowercase_	:   Optional[Any]  				=     tf.range(lp_size[0],   dtype=target.dtype  )
									lowercase_	:   str  				=     tf.stack([r, target],   1  )
									return tf.gather_nd(snake_case__,   snake_case__  )






		def 						snake_case__				(							self,   snake_case__,   snake_case__,   snake_case__=True,   snake_case__=False  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   Any  				=     0
									if self.n_clusters == 0:
																lowercase_	:   Union[str, Any]  				=     self._logit(snake_case__,   self.out_layers[0][0],   self.out_layers[0][1],   self.out_projs[0]  )
																if target is not None:
																							lowercase_	:   str  				=     tf.nn.sparse_softmax_cross_entropy_with_logits(labels=snake_case__,   logits=snake_case__  )
																lowercase_	:   Optional[int]  				=     tf.nn.log_softmax(snake_case__,   axis=-1  )
									else:
																lowercase_	:   int  				=     shape_list(snake_case__  )
																lowercase_	:   List[Any]  				=     []
																lowercase_	:   Tuple  				=     tf.zeros(hidden_sizes[:2]  )
																for i in range(len(self.cutoffs  )  ):
																							lowercase_				,		lowercase_	:   Optional[int]  				=     self.cutoff_ends[i], self.cutoff_ends[i + 1]
																							if target is not None:
																														lowercase_	:   Tuple  				=     (target >= l_idx) & (target < r_idx)
																														lowercase_	:   Union[str, Any]  				=     tf.where(snake_case__  )
																														lowercase_	:   Optional[Any]  				=     tf.boolean_mask(snake_case__,   snake_case__  ) - l_idx

																							if self.div_val == 1:
																														lowercase_	:   Optional[Any]  				=     self.out_layers[0][0][l_idx:r_idx]
																														lowercase_	:   Union[str, Any]  				=     self.out_layers[0][1][l_idx:r_idx]
																							else:
																														lowercase_	:   Dict  				=     self.out_layers[i][0]
																														lowercase_	:   int  				=     self.out_layers[i][1]

																							if i == 0:
																														lowercase_	:   Optional[int]  				=     tf.concat([cur_W, self.cluster_weight],   0  )
																														lowercase_	:   Optional[Any]  				=     tf.concat([cur_b, self.cluster_bias],   0  )

																														lowercase_	:   List[str]  				=     self._logit(snake_case__,   snake_case__,   snake_case__,   self.out_projs[0]  )
																														lowercase_	:   List[str]  				=     tf.nn.log_softmax(snake_case__  )
																														out.append(head_logprob[..., : self.cutoffs[0]]  )
																														if target is not None:
																																					lowercase_	:   int  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Optional[int]  				=     self._gather_logprob(snake_case__,   snake_case__  )
																							else:
																														lowercase_	:   List[str]  				=     self._logit(snake_case__,   snake_case__,   snake_case__,   self.out_projs[i]  )
																														lowercase_	:   Dict  				=     tf.nn.log_softmax(snake_case__  )
																														lowercase_	:   Optional[int]  				=     self.cutoffs[0] + i - 1  # No probability for the head cluster
																														lowercase_	:   int  				=     head_logprob[..., cluster_prob_idx, None] + tail_logprob
																														out.append(snake_case__  )
																														if target is not None:
																																					lowercase_	:   Optional[Any]  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Optional[int]  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Union[str, Any]  				=     self._gather_logprob(snake_case__,   snake_case__  )
																																					cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
																							if target is not None:
																														loss += tf.scatter_nd(snake_case__,   -cur_logprob,   shape_list(snake_case__  )  )
																lowercase_	:   List[str]  				=     tf.concat(snake_case__,   axis=-1  )

									if target is not None:
																if return_mean:
																							lowercase_	:   Tuple  				=     tf.reduce_mean(snake_case__  )
																# Add the training-time loss value to the layer using `self.add_loss()`.
																self.add_loss(snake_case__  )

																# Log the loss as a metric (we could log arbitrary metrics,
																# including different metrics for training and inference.
																self.add_metric(snake_case__,   name=self.name,   aggregation="""mean""" if return_mean else """"""  )

									return out