


"""simple docstring"""






from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types


def __lowercase  (   lowerCamelCase_	:   torch.Tensor    ,       lowerCamelCase_	:   torch.Tensor    ,       lowerCamelCase_	:   bool	):
				# The first operation in a checkpoint can't be in-place, but it's
				# nice to have in-place addition during inference. Thus...
				if not inplace:
								SCREAMING_SNAKE_CASE__  				=					ma + ma
				else:
								ma += ma

				return ma


def __lowercase  (   lowerCamelCase_	:   torch.Tensor    ,       lowerCamelCase_	:   List[int]	):
				SCREAMING_SNAKE_CASE__  				=					-1 * len(_a	)
				SCREAMING_SNAKE_CASE__  				=					list(range(len(tensor.shape[:zero_index]	)	)	)
				return tensor.permute(first_inds + [zero_index + i for i in inds]	)


def __lowercase  (   lowerCamelCase_	:   torch.Tensor    ,       lowerCamelCase_	:   int	):
				return t.reshape(t.shape[:-no_dims] + (-1,)	)


def __lowercase  (   lowerCamelCase_	:   torch.Tensor    ,       lowerCamelCase_	:   torch.Tensor    ,       lowerCamelCase_	:   int    ,       lowerCamelCase_	:   float = 1e-4	):
				SCREAMING_SNAKE_CASE__  				=					mask.expand(*value.shape	)
				return torch.sum(mask * value    ,       dim=_a	) / (eps + torch.sum(_a    ,       dim=_a	))


def __lowercase  (   lowerCamelCase_	:   torch.Tensor    ,       lowerCamelCase_	:   torch.types.Number = 2.3125    ,       lowerCamelCase_	:   torch.types.Number = 21.6875    ,       lowerCamelCase_	:   int = 64	):
				SCREAMING_SNAKE_CASE__  				=					torch.linspace(_a    ,       _a    ,       no_bins - 1    ,       device=pts.device	)
				SCREAMING_SNAKE_CASE__  				=					torch.sqrt(torch.sum((pts.unsqueeze(-2	) - pts.unsqueeze(-3	)) ** 2    ,       dim=-1	)	)
				return torch.bucketize(_a    ,       _a	)


def __lowercase  (   lowerCamelCase_	:   Callable[[list], Any]    ,       lowerCamelCase_	:   List[dict]	):
				SCREAMING_SNAKE_CASE__  				=					dicts[0]
				SCREAMING_SNAKE_CASE__  				=					{}
				for k, v in first.items():
								SCREAMING_SNAKE_CASE__  				=					[d[k] for d in dicts]
								if isinstance(_a    ,       _a	):
												SCREAMING_SNAKE_CASE__  				=					dict_multimap(_a    ,       _a	)
								else:
												SCREAMING_SNAKE_CASE__  				=					fn(_a	)

				return new_dict


def __lowercase  (   lowerCamelCase_	:   torch.Tensor    ,       lowerCamelCase_	:   torch.Tensor	):
				SCREAMING_SNAKE_CASE__  				=					v_bins.view(((1,) * len(x.shape	)) + (len(_a	),)	)
				SCREAMING_SNAKE_CASE__  				=					x[..., None] - reshaped_bins
				SCREAMING_SNAKE_CASE__  				=					torch.argmin(torch.abs(_a	)    ,       dim=-1	)
				return nn.functional.one_hot(_a    ,       num_classes=len(_a	)	).float()


def __lowercase  (   lowerCamelCase_	:   torch.Tensor    ,       lowerCamelCase_	:   torch.Tensor    ,       lowerCamelCase_	:   int = 0    ,       lowerCamelCase_	:   int = 0	):
				SCREAMING_SNAKE_CASE__  				=					[]
				for i, s in enumerate(data.shape[:no_batch_dims]	):
								SCREAMING_SNAKE_CASE__  				=					torch.arange(_a	)
								SCREAMING_SNAKE_CASE__  				=					r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape	) - i - 1)))	)
								ranges.append(_a	)

				SCREAMING_SNAKE_CASE__  				=					[slice(_a	) for _ in range(len(data.shape	) - no_batch_dims	)]
				SCREAMING_SNAKE_CASE__  				=					inds
				ranges.extend(_a	)
				# Matt note: Editing this to get around the behaviour of using a list as an array index changing
				# in recent Numpy versions
				return data[tuple(_a	)]


_lowerCamelCase  =	TypeVar('T')


def __lowercase  (   lowerCamelCase_	:   Callable[[T], Any]    ,       lowerCamelCase_	:   Dict[Any, Union[dict, list, tuple, T]]    ,       lowerCamelCase_	:   Type[T]	):
				SCREAMING_SNAKE_CASE__  				=					{}
				for k, v in dic.items():
								if isinstance(_a    ,       _a	):
												SCREAMING_SNAKE_CASE__  				=					dict_map(_a    ,       _a    ,       _a	)
								else:
												SCREAMING_SNAKE_CASE__  				=					tree_map(_a    ,       _a    ,       _a	)

				return new_dict


@overload
def __lowercase  (   lowerCamelCase_	:   Callable[[T], Any]    ,       lowerCamelCase_	:   T    ,       lowerCamelCase_	:   Type[T]	):
				...


@overload
def __lowercase  (   lowerCamelCase_	:   Callable[[T], Any]    ,       lowerCamelCase_	:   dict    ,       lowerCamelCase_	:   Type[T]	):
				...


@overload
def __lowercase  (   lowerCamelCase_	:   Callable[[T], Any]    ,       lowerCamelCase_	:   list    ,       lowerCamelCase_	:   Type[T]	):
				...


@overload
def __lowercase  (   lowerCamelCase_	:   Callable[[T], Any]    ,       lowerCamelCase_	:   tuple    ,       lowerCamelCase_	:   Type[T]	):
				...


def __lowercase  (   lowerCamelCase_	:   str    ,       lowerCamelCase_	:   str    ,       lowerCamelCase_	:   Optional[Any]	):
				if isinstance(_a    ,       _a	):
								return dict_map(_a    ,       _a    ,       _a	)
				elif isinstance(_a    ,       _a	):
								return [tree_map(_a    ,       _a    ,       _a	) for x in tree]
				elif isinstance(_a    ,       _a	):
								return tuple(tree_map(_a    ,       _a    ,       _a	) for x in tree	)
				elif isinstance(_a    ,       _a	):
								return fn(_a	)
				else:
								print(type(_a	)	)
								raise ValueError("Not supported"	)

_lowerCamelCase  =	partial(tree_map, leaf_type=torch.Tensor)







import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig
from ..clip.configuration_clip import CLIPVisionConfig


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





class   lowercase__     (				__A	):
  __UpperCamelCase						  =	"""vision-text-dual-encoder"""
  __UpperCamelCase						  =	True







  def __init__(						self     ,    _lowercase=512     ,    _lowercase=2.6592     ,    **_lowercase    ):
       super().__init__(**_lowercase    )

       if "vision_config" not in kwargs:
            raise ValueError("""`vision_config` can not be `None`."""    )

       if "text_config" not in kwargs:
            raise ValueError("""`text_config` can not be `None`."""    )

       lowerCAmelCase_       :							Dict			 =		kwargs.pop("""vision_config"""    )
       lowerCAmelCase_       :							Optional[int]			 =		kwargs.pop("""text_config"""    )

       lowerCAmelCase_       :							Tuple			 =		vision_config.pop("""model_type"""    )
       lowerCAmelCase_       :							str			 =		text_config.pop("""model_type"""    )

       if vision_model_type == "clip":
            lowerCAmelCase_       :							Dict			 =		AutoConfig.for_model(_lowercase     ,    **_lowercase    ).vision_config
       elif vision_model_type == "clip_vision_model":
            lowerCAmelCase_       :							Any			 =		CLIPVisionConfig(**_lowercase    )
       else:
            lowerCAmelCase_       :							Tuple			 =		AutoConfig.for_model(_lowercase     ,    **_lowercase    )

       lowerCAmelCase_       :							List[Any]			 =		AutoConfig.for_model(_lowercase     ,    **_lowercase    )

       lowerCAmelCase_       :							Union[str, Any]			 =		projection_dim
       lowerCAmelCase_       :							Optional[int]			 =		logit_scale_init_value







  @classmethod
  def 						UpperCAmelCase__		(						cls     ,    _lowercase     ,    _lowercase     ,    **_lowercase    ):
       return cls(vision_config=vision_config.to_dict()     ,    text_config=text_config.to_dict()     ,    **_lowercase    )







  def 						UpperCAmelCase__		(						self    ):
       lowerCAmelCase_       :							int			 =		copy.deepcopy(self.__dict__    )
       lowerCAmelCase_       :							List[str]			 =		self.vision_config.to_dict()
       lowerCAmelCase_       :							List[str]			 =		self.text_config.to_dict()
       lowerCAmelCase_       :							str			 =		self.__class__.model_type
       return output
