



"""simple docstring"""

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


lowercase__			   =  logging.get_logger(__name__)

lowercase__			   =  {
    "google/realm-cc-news-pretrained-embedder": (
        "https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/config.json"
    ),
    "google/realm-cc-news-pretrained-encoder": (
        "https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/config.json"
    ),
    "google/realm-cc-news-pretrained-scorer": (
        "https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/config.json"
    ),
    "google/realm-cc-news-pretrained-openqa": (
        "https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/config.json"
    ),
    "google/realm-orqa-nq-openqa": "https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/config.json",
    "google/realm-orqa-nq-reader": "https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/config.json",
    "google/realm-orqa-wq-openqa": "https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/config.json",
    "google/realm-orqa-wq-reader": "https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/config.json",
    # See all REALM models at https://huggingface.co/models?filter=realm
}




class SCREAMING_SNAKE_CASE__					(		__snake_case      ):
						_lowerCAmelCase							=					"realm"



						def __init__(self						,     _lowercase=30522						,     _lowercase=768						,     _lowercase=128						,     _lowercase=12						,     _lowercase=12						,     _lowercase=8						,     _lowercase=3072						,     _lowercase="gelu_new"						,     _lowercase=0.1						,     _lowercase=0.1						,     _lowercase=512						,     _lowercase=2						,     _lowercase=0.02						,     _lowercase=1e-12						,     _lowercase=256						,     _lowercase=10						,     _lowercase=1e-3						,     _lowercase=5						,     _lowercase=320						,     _lowercase=13353718						,     _lowercase=5000						,     _lowercase=1						,     _lowercase=0						,     _lowercase=2						,     **_lowercase						,     ):





													'''simple docstring'''

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

													# Common config
													__a		:		Any									=       vocab_size
													__a		:		Union[str, Any]									=       max_position_embeddings
													__a		:		int									=       hidden_size
													__a		:		Optional[int]									=       retriever_proj_size
													__a		:		int									=       num_hidden_layers
													__a		:		List[Any]									=       num_attention_heads
													__a		:		Union[str, Any]									=       num_candidates
													__a		:		int									=       intermediate_size
													__a		:		List[Any]									=       hidden_act
													__a		:		Dict									=       hidden_dropout_prob
													__a		:		int									=       attention_probs_dropout_prob
													__a		:		List[Any]									=       initializer_range
													__a		:		str									=       type_vocab_size
													__a		:		List[Any]									=       layer_norm_eps

													# Reader config
													__a		:		Any									=       span_hidden_size
													__a		:		List[Any]									=       max_span_width
													__a		:		Tuple									=       reader_layer_norm_eps
													__a		:		Dict									=       reader_beam_size
													__a		:		List[Any]									=       reader_seq_len

													# Retrieval config
													__a		:		str									=       num_block_records
													__a		:		int									=       searcher_beam_size






"""simple docstring"""

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 SCREAMING_SNAKE_CASE__					(		__snake_case      ):



						def __init__(self						,     _lowercase						,     _lowercase=None						,     _lowercase=None						,     _lowercase=0							):





													'''simple docstring'''

													__a		:		Any									=       1.0 if scale is None else scale
													__a		:		str									=       0.0 if loc is None else loc

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



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

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



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

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





						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													return self.variance.sqrt()







class SCREAMING_SNAKE_CASE__					(		nn.Module      ):



						def __init__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     **_lowercase							):





													'''simple docstring'''

													super().__init__(**_lowercase							)
													__a		:		str									=       args_dim
													__a		:		List[Any]									=       nn.ModuleList([nn.Linear(_lowercase						,     _lowercase							) for dim in args_dim.values()]							)
													__a		:		Dict									=       domain_map





						def     lowerCAmelCase__(self						,     _lowercase							):





													'''simple docstring'''

													__a		:		List[Any]									=       [proj(_lowercase							) for proj in self.proj]

													return self.domain_map(*_lowercase							)







class SCREAMING_SNAKE_CASE__					(		nn.Module      ):



						def __init__(self						,     _lowercase							):





													'''simple docstring'''

													super().__init__()
													__a		:		Optional[int]									=       function





						def     lowerCAmelCase__(self						,     _lowercase						,     *_lowercase							):





													'''simple docstring'''

													return self.function(_lowercase						,     *_lowercase							)







class SCREAMING_SNAKE_CASE__					:
						_lowerCAmelCase							=					42
						_lowerCAmelCase							=					42
						_lowerCAmelCase							=					42



						def __init__(self						,     _lowercase = 1							):





													'''simple docstring'''

													__a		:		Optional[int]									=       dim
													__a		:		str									=       {k: dim * self.args_dim[k] for k in self.args_dim}



						def     lowerCAmelCase__(self						,     _lowercase							):





													'''simple docstring'''

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



						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase = None						,     _lowercase = None						,     ):





													'''simple docstring'''

													__a		:		Tuple									=       self._base_distribution(_lowercase							)
													if loc is None and scale is None:
																				return distr
													else:
																				return AffineTransformed(_lowercase						,     loc=_lowercase						,     scale=_lowercase						,     event_dim=self.event_dim							)



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

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



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													return len(self.event_shape							)



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													return 0.0



						def     lowerCAmelCase__(self						,     _lowercase							):





													'''simple docstring'''

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



						def     lowerCAmelCase__(self						,     *_lowercase							):





													'''simple docstring'''

													raise NotImplementedError()





						@staticmethod
						def     lowerCAmelCase__(_lowercase							):





													'''simple docstring'''

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







class SCREAMING_SNAKE_CASE__					(		__snake_case      ):
						_lowerCAmelCase							=					{"df": 1, "loc": 1, "scale": 1}
						_lowerCAmelCase							=					StudentT



						@classmethod
						def     lowerCAmelCase__(cls						,     _lowercase						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

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







class SCREAMING_SNAKE_CASE__					(		__snake_case      ):
						_lowerCAmelCase							=					{"loc": 1, "scale": 1}
						_lowerCAmelCase							=					Normal



						@classmethod
						def     lowerCAmelCase__(cls						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		str									=       cls.squareplus(_lowercase							).clamp_min(torch.finfo(scale.dtype							).eps							)
													return loc.squeeze(-1							), scale.squeeze(-1							)







class SCREAMING_SNAKE_CASE__					(		__snake_case      ):
						_lowerCAmelCase							=					{"total_count": 1, "logits": 1}
						_lowerCAmelCase							=					NegativeBinomial



						@classmethod
						def     lowerCAmelCase__(cls						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		Union[str, Any]									=       cls.squareplus(_lowercase							)
													return total_count.squeeze(-1							), logits.squeeze(-1							)



						def     lowerCAmelCase__(self						,     _lowercase							):





													'''simple docstring'''

													__a							,     __a		:		Optional[Any]									=       distr_args
													if self.dim == 1:
																				return self.distribution_class(total_count=_lowercase						,     logits=_lowercase							)
													else:
																				return Independent(self.distribution_class(total_count=_lowercase						,     logits=_lowercase							)						,     1							)





						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase = None						,     _lowercase = None							):





													'''simple docstring'''

													__a							,     __a		:		List[Any]									=       distr_args

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

													return self._base_distribution((total_count, logits)							)

