


import contextlib
import copy
import random
from typing import Any, Dict, Iterable, Optional, Union

import numpy as np
import torch

from .utils import deprecate, is_transformers_available


if is_transformers_available():
  import transformers







def 						__lowerCAmelCase		(		lowercase						:						int				)			->		Tuple:


    """simple docstring"""





    random.seed(_lowerCAmelCase				)
    np.random.seed(_lowerCAmelCase				)
    torch.manual_seed(_lowerCAmelCase				)
    torch.cuda.manual_seed_all(_lowerCAmelCase				)
    # ^^ safe to call this function even if cuda is not available







class _lowerCAmelCase       :







   def __init__(			self		,			UpperCamelCase__		,			UpperCamelCase__ = 0.9999		,			UpperCamelCase__ = 0.0		,			UpperCamelCase__ = 0		,			UpperCamelCase__ = False		,			UpperCamelCase__ = 1.0		,			UpperCamelCase__ = 2 / 3		,			UpperCamelCase__ = None		,			UpperCamelCase__ = None		,			**UpperCamelCase__		,			)			->	Any:

       '''simple docstring'''




       if isinstance(__A		,			torch.nn.Module		):
           snake_case					:   Any         =					(
               "Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. "
               "Please pass the parameters of the module instead."
           )
           deprecate(
               "passing a `torch.nn.Module` to `ExponentialMovingAverage`"		,			"1.0.0"		,			__A		,			standard_warn=__A		,			)
           snake_case					:   Dict         =					parameters.parameters()

           # set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility
           snake_case					:   Optional[Any]         =					True

       if kwargs.get("max_value"		,			__A		) is not None:
           snake_case					:   Optional[int]         =					"The `max_value` argument is deprecated. Please use `decay` instead."
           deprecate("max_value"		,			"1.0.0"		,			__A		,			standard_warn=__A		)
           snake_case					:   Optional[Any]         =					kwargs["max_value"]

       if kwargs.get("min_value"		,			__A		) is not None:
           snake_case					:   Tuple         =					"The `min_value` argument is deprecated. Please use `min_decay` instead."
           deprecate("min_value"		,			"1.0.0"		,			__A		,			standard_warn=__A		)
           snake_case					:   Any         =					kwargs["min_value"]

       snake_case					:   Optional[int]         =					list(__A		)
       snake_case					:   List[Any]         =					[p.clone().detach() for p in parameters]

       if kwargs.get("device"		,			__A		) is not None:
           snake_case					:   Optional[int]         =					"The `device` argument is deprecated. Please use `to` instead."
           deprecate("device"		,			"1.0.0"		,			__A		,			standard_warn=__A		)
           self.to(device=kwargs["device"]		)

       snake_case					:   Tuple         =					None

       snake_case					:   Dict         =					decay
       snake_case					:   List[Any]         =					min_decay
       snake_case					:   Optional[Any]         =					update_after_step
       snake_case					:   Any         =					use_ema_warmup
       snake_case					:   Union[str, Any]         =					inv_gamma
       snake_case					:   str         =					power
       snake_case					:   Union[str, Any]         =					0
       snake_case					:   Union[str, Any]         =					None  # set in `step()`

       snake_case					:   List[str]         =					model_cls
       snake_case					:   Dict         =					model_config







   @classmethod
   def 				lowerCamelCase	(			cls		,			UpperCamelCase__		,			UpperCamelCase__		)			->	Union[str, Any]:

       '''simple docstring'''




       snake_case					:   Optional[int]         =					model_cls.load_config(__A		,			return_unused_kwargs=__A		)
       snake_case					:   Optional[Any]         =					model_cls.from_pretrained(__A		)

       snake_case					:   Union[str, Any]         =					cls(model.parameters()		,			model_cls=__A		,			model_config=model.config		)

       ema_model.load_state_dict(__A		)
       return ema_model







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

       '''simple docstring'''




       if self.model_cls is None:
           raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__."		)

       if self.model_config is None:
           raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__."		)

       snake_case					:   Tuple         =					self.model_cls.from_config(self.model_config		)
       snake_case					:   List[str]         =					self.state_dict()
       state_dict.pop("shadow_params"		,			__A		)

       model.register_to_config(**__A		)
       self.copy_to(model.parameters()		)
       model.save_pretrained(__A		)







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		)			->	Any:

       '''simple docstring'''




       snake_case					:   str         =					max(0		,			optimization_step - self.update_after_step - 1		)

       if step <= 0:
           return 0.0

       if self.use_ema_warmup:
           snake_case					:   Tuple         =					1 - (1 + step / self.inv_gamma) ** -self.power
       else:
           snake_case					:   List[str]         =					(1 + step) / (10 + step)

       snake_case					:   Union[str, Any]         =					min(__A		,			self.decay		)
       # make sure decay is not smaller than min_decay
       snake_case					:   Union[str, Any]         =					max(__A		,			self.min_decay		)
       return cur_decay_value







   @torch.no_grad()
   def 				lowerCamelCase	(			self		,			UpperCamelCase__		)			->	int:

       '''simple docstring'''




       if isinstance(__A		,			torch.nn.Module		):
           snake_case					:   Dict         =					(
               "Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. "
               "Please pass the parameters of the module instead."
           )
           deprecate(
               "passing a `torch.nn.Module` to `ExponentialMovingAverage.step`"		,			"1.0.0"		,			__A		,			standard_warn=__A		,			)
           snake_case					:   Any         =					parameters.parameters()

       snake_case					:   str         =					list(__A		)

       self.optimization_step += 1

       # Compute the decay factor for the exponential moving average.
       snake_case					:   Dict         =					self.get_decay(self.optimization_step		)
       snake_case					:   Optional[Any]         =					decay
       snake_case					:   List[Any]         =					1 - decay

       snake_case					:   List[Any]         =					contextlib.nullcontext
       if is_transformers_available() and transformers.deepspeed.is_deepspeed_zeroa_enabled():
           import deepspeed

       for s_param, param in zip(self.shadow_params		,			__A		):
           if is_transformers_available() and transformers.deepspeed.is_deepspeed_zeroa_enabled():
               snake_case					:   List[Any]         =					deepspeed.zero.GatheredParameters(__A		,			modifier_rank=__A		)

           with context_manager():
               if param.requires_grad:
                   s_param.sub_(one_minus_decay * (s_param - param)		)
               else:
                   s_param.copy_(__A		)







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		)			->	List[str]:

       '''simple docstring'''




       snake_case					:   Tuple         =					list(__A		)
       for s_param, param in zip(self.shadow_params		,			__A		):
           param.data.copy_(s_param.to(param.device		).data		)







   def 				lowerCamelCase	(			self		,			UpperCamelCase__=None		,			UpperCamelCase__=None		)			->	List[str]:

       '''simple docstring'''




       snake_case					:   int         =					[
           p.to(device=__A		,			dtype=__A		) if p.is_floating_point() else p.to(device=__A		)
           for p in self.shadow_params
       ]







   def 				lowerCamelCase	(			self		)			->	List[Any]:

       '''simple docstring'''




       return {
           "decay": self.decay,
           "min_decay": self.min_decay,
           "optimization_step": self.optimization_step,
           "update_after_step": self.update_after_step,
           "use_ema_warmup": self.use_ema_warmup,
           "inv_gamma": self.inv_gamma,
           "power": self.power,
           "shadow_params": self.shadow_params,
       }







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		)			->	Dict:

       '''simple docstring'''




       snake_case					:   Tuple         =					[param.detach().cpu().clone() for param in parameters]







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		)			->	Union[str, Any]:

       '''simple docstring'''




       if self.temp_stored_params is None:
           raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights " "to `restore()`"		)
       for c_param, param in zip(self.temp_stored_params		,			__A		):
           param.data.copy_(c_param.data		)

       # Better memory-wise.
       snake_case					:   List[str]         =					None







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

       '''simple docstring'''




       snake_case					:   List[str]         =					copy.deepcopy(__A		)

       snake_case					:   Optional[Any]         =					state_dict.get("decay"		,			self.decay		)
       if self.decay < 0.0 or self.decay > 1.0:
           raise ValueError("Decay must be between 0 and 1"		)

       snake_case					:   Dict         =					state_dict.get("min_decay"		,			self.min_decay		)
       if not isinstance(self.min_decay		,			__A		):
           raise ValueError("Invalid min_decay"		)

       snake_case					:   Tuple         =					state_dict.get("optimization_step"		,			self.optimization_step		)
       if not isinstance(self.optimization_step		,			__A		):
           raise ValueError("Invalid optimization_step"		)

       snake_case					:   Any         =					state_dict.get("update_after_step"		,			self.update_after_step		)
       if not isinstance(self.update_after_step		,			__A		):
           raise ValueError("Invalid update_after_step"		)

       snake_case					:   List[Any]         =					state_dict.get("use_ema_warmup"		,			self.use_ema_warmup		)
       if not isinstance(self.use_ema_warmup		,			__A		):
           raise ValueError("Invalid use_ema_warmup"		)

       snake_case					:   Tuple         =					state_dict.get("inv_gamma"		,			self.inv_gamma		)
       if not isinstance(self.inv_gamma		,			(float, int)		):
           raise ValueError("Invalid inv_gamma"		)

       snake_case					:   Union[str, Any]         =					state_dict.get("power"		,			self.power		)
       if not isinstance(self.power		,			(float, int)		):
           raise ValueError("Invalid power"		)

       snake_case					:   Optional[Any]         =					state_dict.get("shadow_params"		,			__A		)
       if shadow_params is not None:
           snake_case					:   str         =					shadow_params
           if not isinstance(self.shadow_params		,			__A		):
               raise ValueError("shadow_params must be a list"		)
           if not all(isinstance(__A		,			torch.Tensor		) for p in self.shadow_params		):
               raise ValueError("shadow_params must all be Tensors"		)









"""simple docstring"""



from collections import OrderedDict
from typing import Mapping

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


__snake_case         =					logging.get_logger(__name__)

__snake_case         =					{
    """google/bigbird-roberta-base""": """https://huggingface.co/google/bigbird-roberta-base/resolve/main/config.json""",
    """google/bigbird-roberta-large""": """https://huggingface.co/google/bigbird-roberta-large/resolve/main/config.json""",
    """google/bigbird-base-trivia-itc""": """https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/config.json""",
    # See all BigBird models at https://huggingface.co/models?filter=big_bird
}







class _lowerCAmelCase       (   snake_case_			):
   __UpperCAmelCase				: Optional[Any]													=						'''big_bird'''







   def __init__(			self		,			UpperCamelCase__=5_0358		,			UpperCamelCase__=768		,			UpperCamelCase__=12		,			UpperCamelCase__=12		,			UpperCamelCase__=3072		,			UpperCamelCase__="gelu_new"		,			UpperCamelCase__=0.1		,			UpperCamelCase__=0.1		,			UpperCamelCase__=4096		,			UpperCamelCase__=2		,			UpperCamelCase__=0.02		,			UpperCamelCase__=1e-12		,			UpperCamelCase__=True		,			UpperCamelCase__=0		,			UpperCamelCase__=1		,			UpperCamelCase__=2		,			UpperCamelCase__=66		,			UpperCamelCase__="block_sparse"		,			UpperCamelCase__=True		,			UpperCamelCase__=False		,			UpperCamelCase__=64		,			UpperCamelCase__=3		,			UpperCamelCase__=None		,			**UpperCamelCase__		,			)			->	Tuple:

       '''simple docstring'''




       super().__init__(
           pad_token_id=UpperCamelCase__		,			bos_token_id=UpperCamelCase__		,			eos_token_id=UpperCamelCase__		,			sep_token_id=UpperCamelCase__		,			**UpperCamelCase__		,			)

       snake_case					:   Union[str, Any]         =					vocab_size
       snake_case					:   List[Any]         =					max_position_embeddings
       snake_case					:   int         =					hidden_size
       snake_case					:   str         =					num_hidden_layers
       snake_case					:   Any         =					num_attention_heads
       snake_case					:   int         =					intermediate_size
       snake_case					:   Union[str, Any]         =					hidden_act
       snake_case					:   Optional[Any]         =					hidden_dropout_prob
       snake_case					:   List[Any]         =					attention_probs_dropout_prob
       snake_case					:   int         =					initializer_range
       snake_case					:   List[str]         =					type_vocab_size
       snake_case					:   Optional[Any]         =					layer_norm_eps
       snake_case					:   Optional[Any]         =					use_cache

       snake_case					:   List[Any]         =					rescale_embeddings
       snake_case					:   Any         =					attention_type
       snake_case					:   List[Any]         =					use_bias
       snake_case					:   int         =					block_size
       snake_case					:   int         =					num_random_blocks
       snake_case					:   Optional[int]         =					classifier_dropout







class _lowerCAmelCase       (   snake_case_			):







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

       '''simple docstring'''




       if self.task == "multiple-choice":
           snake_case					:   int         =					{0: "batch", 1: "choice", 2: "sequence"}
       else:
           snake_case					:   Optional[Any]         =					{0: "batch", 1: "sequence"}
       return OrderedDict(
           [
               ("input_ids", dynamic_axis),
               ("attention_mask", dynamic_axis),
           ]		)





