





import copy
import os
from typing import TYPE_CHECKING, List, Union


if TYPE_CHECKING:
    pass

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase	:					Any			  =					logging.get_logger(__name__)

__lowerCamelCase	:					Optional[int]			  =					{
    'kakaobrain/align-base': 'https://huggingface.co/kakaobrain/align-base/resolve/main/config.json',
}



class  A__       (	_lowerCamelCase      ):
      _UpperCAmelCase     :Tuple						   =       'align_text_model'





      def __init__(			self	,							A_=3_0522	,							A_=768	,							A_=12	,							A_=12	,							A_=3072	,							A_="gelu"	,							A_=0.1	,							A_=0.1	,							A_=512	,							A_=2	,							A_=0.02	,							A_=1e-12	,							A_=0	,							A_="absolute"	,							A_=True	,							**A_	,							):
       '''simple docstring'''




       super().__init__(**lowercase_       )

       UpperCamelCase     : int	     =       vocab_size
       UpperCamelCase     : List[Any]	     =       hidden_size
       UpperCamelCase     : List[Any]	     =       num_hidden_layers
       UpperCamelCase     : Optional[Any]	     =       num_attention_heads
       UpperCamelCase     : Any	     =       hidden_act
       UpperCamelCase     : Tuple	     =       intermediate_size
       UpperCamelCase     : Any	     =       hidden_dropout_prob
       UpperCamelCase     : Any	     =       attention_probs_dropout_prob
       UpperCamelCase     : Dict	     =       max_position_embeddings
       UpperCamelCase     : Any	     =       type_vocab_size
       UpperCamelCase     : Tuple	     =       initializer_range
       UpperCamelCase     : Optional[int]	     =       layer_norm_eps
       UpperCamelCase     : Any	     =       position_embedding_type
       UpperCamelCase     : List[Any]	     =       use_cache
       UpperCamelCase     : Any	     =       pad_token_id





      @classmethod
      def 	__UpperCamelCase(			cls	,							A_	,							**A_       ):
       '''simple docstring'''




       cls._set_token_in_kwargs(lowercase_       )

       UpperCamelCase     ,   UpperCamelCase     : List[Any]	     =       cls.get_config_dict(lowercase_	,							**lowercase_       )

       # get the text config dict if we are loading from AlignConfig
       if config_dict.get("model_type"       ) == "align":
        UpperCamelCase     : Optional[Any]	     =       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(lowercase_	,							**lowercase_       )



class  A__       (	_lowerCamelCase      ):
      _UpperCAmelCase     :Optional[int]						   =       'align_vision_model'





      def __init__(			self	,							A_ = 3	,							A_ = 600	,							A_ = 2.0	,							A_ = 3.1	,							A_ = 8	,							A_ = [3, 3, 5, 3, 5, 5, 3]	,							A_ = [32, 16, 24, 40, 80, 112, 192]	,							A_ = [16, 24, 40, 80, 112, 192, 320]	,							A_ = []	,							A_ = [1, 2, 2, 2, 1, 2, 1]	,							A_ = [1, 2, 2, 3, 3, 4, 1]	,							A_ = [1, 6, 6, 6, 6, 6, 6]	,							A_ = 0.25	,							A_ = "swish"	,							A_ = 2560	,							A_ = "mean"	,							A_ = 0.02	,							A_ = 0.0_01	,							A_ = 0.99	,							A_ = 0.2	,							**A_	,							):
       '''simple docstring'''




       super().__init__(**lowercase_       )

       UpperCamelCase     : Dict	     =       num_channels
       UpperCamelCase     : Optional[Any]	     =       image_size
       UpperCamelCase     : List[Any]	     =       width_coefficient
       UpperCamelCase     : Optional[Any]	     =       depth_coefficient
       UpperCamelCase     : Any	     =       depth_divisor
       UpperCamelCase     : List[Any]	     =       kernel_sizes
       UpperCamelCase     : Tuple	     =       in_channels
       UpperCamelCase     : Union[str, Any]	     =       out_channels
       UpperCamelCase     : List[Any]	     =       depthwise_padding
       UpperCamelCase     : Union[str, Any]	     =       strides
       UpperCamelCase     : Optional[Any]	     =       num_block_repeats
       UpperCamelCase     : Union[str, Any]	     =       expand_ratios
       UpperCamelCase     : Dict	     =       squeeze_expansion_ratio
       UpperCamelCase     : Dict	     =       hidden_act
       UpperCamelCase     : Any	     =       hidden_dim
       UpperCamelCase     : str	     =       pooling_type
       UpperCamelCase     : Any	     =       initializer_range
       UpperCamelCase     : Optional[Any]	     =       batch_norm_eps
       UpperCamelCase     : Any	     =       batch_norm_momentum
       UpperCamelCase     : List[Any]	     =       drop_connect_rate
       UpperCamelCase     : Tuple	     =       sum(lowercase_       ) * 4





      @classmethod
      def 	__UpperCamelCase(			cls	,							A_	,							**A_       ):
       '''simple docstring'''




       cls._set_token_in_kwargs(lowercase_       )

       UpperCamelCase     ,   UpperCamelCase     : Optional[int]	     =       cls.get_config_dict(lowercase_	,							**lowercase_       )

       # get the vision config dict if we are loading from AlignConfig
       if config_dict.get("model_type"       ) == "align":
        UpperCamelCase     : str	     =       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(lowercase_	,							**lowercase_       )



class  A__       (	_lowerCamelCase      ):
      _UpperCAmelCase     :Union[str, Any]						   =       'align'
      _UpperCAmelCase     :Any						   =       True





      def __init__(			self	,							A_=None	,							A_=None	,							A_=640	,							A_=1.0	,							A_=0.02	,							**A_	,							):
       '''simple docstring'''




       super().__init__(**lowercase_       )

       if text_config is None:
        UpperCamelCase     : Dict	     =       {}
        logger.info("text_config is None. Initializing the AlignTextConfig with default values."       )

       if vision_config is None:
        UpperCamelCase     : Any	     =       {}
        logger.info("vision_config is None. Initializing the AlignVisionConfig with default values."       )

       UpperCamelCase     : List[str]	     =       AlignTextConfig(**lowercase_       )
       UpperCamelCase     : Dict	     =       AlignVisionConfig(**lowercase_       )

       UpperCamelCase     : int	     =       projection_dim
       UpperCamelCase     : Tuple	     =       temperature_init_value
       UpperCamelCase     : Any	     =       initializer_range





      @classmethod
      def 	__UpperCamelCase(			cls	,							A_	,							A_	,							**A_       ):
       '''simple docstring'''




       return cls(text_config=text_config.to_dict()	,							vision_config=vision_config.to_dict()	,							**lowercase_       )





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Optional[Any]	     =       copy.deepcopy(self.__dict__       )
       UpperCamelCase     : Optional[int]	     =       self.text_config.to_dict()
       UpperCamelCase     : Optional[int]	     =       self.vision_config.to_dict()
       UpperCamelCase     : int	     =       self.__class__.model_type
       return output







from typing import Callable, Dict, Optional, Tuple

import torch
from torch import nn
from torch.distributions import (
    AffineTransform,
    Distribution,
    Independent,
    NegativeBinomial,
    Normal,
    StudentT,
    TransformedDistribution,
)



class  A__       (	__snake_case      ):





      def __init__(			self	,							A_	,							A_=None	,							A_=None	,							A_=0       ):
       '''simple docstring'''




       UpperCamelCase     : Union[str, Any]	     =       1.0 if scale is None else scale
       UpperCamelCase     : Optional[int]	     =       0.0 if loc is None else loc

       super().__init__(A_	,							[AffineTransform(loc=self.loc	,							scale=self.scale	,							event_dim=A_       )]       )





      @property
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       return self.base_dist.mean * self.scale + self.loc





      @property
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       return self.base_dist.variance * self.scale**2





      @property
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       return self.variance.sqrt()



class  A__       (	nn.Module      ):





      def __init__(			self	,							A_	,							A_	,							A_	,							**A_       ):
       '''simple docstring'''




       super().__init__(**A_       )
       UpperCamelCase     : Union[str, Any]	     =       args_dim
       UpperCamelCase     : str	     =       nn.ModuleList([nn.Linear(A_	,							A_       ) for dim in args_dim.values()]       )
       UpperCamelCase     : Union[str, Any]	     =       domain_map





      def 	__UpperCamelCase(			self	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     : List[Any]	     =       [proj(A_       ) for proj in self.proj]

       return self.domain_map(*A_       )



class  A__       (	nn.Module      ):





      def __init__(			self	,							A_       ):
       '''simple docstring'''




       super().__init__()
       UpperCamelCase     : str	     =       function





      def 	__UpperCamelCase(			self	,							A_	,							*A_       ):
       '''simple docstring'''




       return self.function(A_	,							*A_       )



class  A__       :
      _UpperCAmelCase     :type
      _UpperCAmelCase     :int
      _UpperCAmelCase     :Dict[str, int]





      def __init__(			self	,							A_ = 1       ):
       '''simple docstring'''




       UpperCamelCase     : Tuple	     =       dim
       UpperCamelCase     : Union[str, Any]	     =       {k: dim * self.args_dim[k] for k in self.args_dim}





      def 	__UpperCamelCase(			self	,							A_       ):
       '''simple docstring'''




       if self.dim == 1:
        return self.distribution_class(*A_       )
       else:
        return Independent(self.distribution_class(*A_       )	,							1       )





      def 	__UpperCamelCase(			self	,							A_	,							A_ = None	,							A_ = None	,							):
       '''simple docstring'''




       UpperCamelCase     : str	     =       self._base_distribution(A_       )
       if loc is None and scale is None:
        return distr
       else:
        return AffineTransformed(A_	,							loc=A_	,							scale=A_	,							event_dim=self.event_dim       )





      @property
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       return () if self.dim == 1 else (self.dim,)





      @property
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       return len(self.event_shape       )





      @property
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       return 0.0





      def 	__UpperCamelCase(			self	,							A_       ):
       '''simple docstring'''




       return ParameterProjection(
           in_features=A_	,							args_dim=self.args_dim	,							domain_map=LambdaLayer(self.domain_map       )	,							)





      def 	__UpperCamelCase(			self	,							*A_       ):
       '''simple docstring'''




       raise NotImplementedError()





      @staticmethod
      def 	__UpperCamelCase(			A_       ):
       '''simple docstring'''




       return (x + torch.sqrt(torch.square(A_       ) + 4.0       )) / 2.0



class  A__       (	__snake_case      ):
      _UpperCAmelCase     :Dict[str, int]						   =       {"df": 1, "loc": 1, "scale": 1}
      _UpperCAmelCase     :type						   =       StudentT





      @classmethod
      def 	__UpperCamelCase(			cls	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     : Optional[Any]	     =       cls.squareplus(A_       ).clamp_min(torch.finfo(scale.dtype       ).eps       )
       UpperCamelCase     : int	     =       2.0 + cls.squareplus(A_       )
       return df.squeeze(-1       ), loc.squeeze(-1       ), scale.squeeze(-1       )



class  A__       (	__snake_case      ):
      _UpperCAmelCase     :Dict[str, int]						   =       {"loc": 1, "scale": 1}
      _UpperCAmelCase     :type						   =       Normal





      @classmethod
      def 	__UpperCamelCase(			cls	,							A_	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     : Dict	     =       cls.squareplus(A_       ).clamp_min(torch.finfo(scale.dtype       ).eps       )
       return loc.squeeze(-1       ), scale.squeeze(-1       )



class  A__       (	__snake_case      ):
      _UpperCAmelCase     :Dict[str, int]						   =       {"total_count": 1, "logits": 1}
      _UpperCAmelCase     :type						   =       NegativeBinomial





      @classmethod
      def 	__UpperCamelCase(			cls	,							A_	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     : List[str]	     =       cls.squareplus(A_       )
       return total_count.squeeze(-1       ), logits.squeeze(-1       )





      def 	__UpperCamelCase(			self	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     ,   UpperCamelCase     : Optional[int]	     =       distr_args
       if self.dim == 1:
        return self.distribution_class(total_count=A_	,							logits=A_       )
       else:
        return Independent(self.distribution_class(total_count=A_	,							logits=A_       )	,							1       )





      def 	__UpperCamelCase(			self	,							A_	,							A_ = None	,							A_ = None       ):
       '''simple docstring'''




       UpperCamelCase     ,   UpperCamelCase     : Any	     =       distr_args

       if scale is not None:
        # See scaling property of Gamma.
        logits += scale.log()

       return self._base_distribution((total_count, logits)       )
