


'''simple docstring'''




from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


__snake_case       		=logging.get_logger(__name__)

__snake_case       		={
    '''facebook/blenderbot_small-90M''': '''https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/config.json''',
    # See all BlenderbotSmall models at https://huggingface.co/models?filter=blenderbot_small
}

class 						UpperCAmelCase_ (    __snake_case						):
     lowerCamelCase							:      str            =	'''blenderbot-small'''
     lowerCamelCase							:      List[str]            =	['''past_key_values''']
     lowerCamelCase							:      Any            =	{'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}





     def __init__(					self			: Tuple       , UpperCAmelCase__			: Optional[Any]=5_0_2_6_5       , UpperCAmelCase__			: Any=5_1_2       , UpperCAmelCase__			: int=8       , UpperCAmelCase__			: Optional[int]=2_0_4_8       , UpperCAmelCase__			: Optional[Any]=1_6       , UpperCAmelCase__			: Optional[Any]=8       , UpperCAmelCase__			: Dict=2_0_4_8       , UpperCAmelCase__			: Union[str, Any]=1_6       , UpperCAmelCase__			: List[str]=0.0       , UpperCAmelCase__			: Tuple=0.0       , UpperCAmelCase__			: Tuple=True       , UpperCAmelCase__			: Optional[int]=True       , UpperCAmelCase__			: Optional[Any]="gelu"       , UpperCAmelCase__			: int=5_1_2       , UpperCAmelCase__			: Dict=0.1       , UpperCAmelCase__			: str=0.0       , UpperCAmelCase__			: Tuple=0.0       , UpperCAmelCase__			: Tuple=0.02       , UpperCAmelCase__			: List[Any]=1       , UpperCAmelCase__			: List[str]=False       , UpperCAmelCase__			: Dict=0       , UpperCAmelCase__			: Tuple=1       , UpperCAmelCase__			: str=2       , UpperCAmelCase__			: List[Any]=2       , **UpperCAmelCase__			: List[Any]       , )							->    Optional[Any]:
        lowerCAmelCase	=     vocab_size
        lowerCAmelCase	=     max_position_embeddings
        lowerCAmelCase	=     d_model
        lowerCAmelCase	=     encoder_ffn_dim
        lowerCAmelCase	=     encoder_layers
        lowerCAmelCase	=     encoder_attention_heads
        lowerCAmelCase	=     decoder_ffn_dim
        lowerCAmelCase	=     decoder_layers
        lowerCAmelCase	=     decoder_attention_heads
        lowerCAmelCase	=     dropout
        lowerCAmelCase	=     attention_dropout
        lowerCAmelCase	=     activation_dropout
        lowerCAmelCase	=     activation_function
        lowerCAmelCase	=     init_std
        lowerCAmelCase	=     encoder_layerdrop
        lowerCAmelCase	=     decoder_layerdrop
        lowerCAmelCase	=     use_cache
        lowerCAmelCase	=     encoder_layers
        lowerCAmelCase	=     scale_embedding  # scale factor will be sqrt(d_model) if True

        super().__init__(
            pad_token_id=_SCREAMING_SNAKE_CASE       , bos_token_id=_SCREAMING_SNAKE_CASE       , eos_token_id=_SCREAMING_SNAKE_CASE       , is_encoder_decoder=_SCREAMING_SNAKE_CASE       , decoder_start_token_id=_SCREAMING_SNAKE_CASE       , forced_eos_token_id=_SCREAMING_SNAKE_CASE       , **_SCREAMING_SNAKE_CASE       , )




class 						UpperCAmelCase_ (    __snake_case						):





     @property
     def 				__UpperCAmelCase (					self			: Optional[int]       )							->    List[str]:




        if self.task in ["default", "seq2seq-lm"]:
           lowerCAmelCase	=     OrderedDict(
               [
                   ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                   ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
               ]       )

           if self.use_past:
              lowerCAmelCase	=     {0: 'batch'}
              lowerCAmelCase	=     {0: 'batch', 1: 'past_decoder_sequence + sequence'}
           else:
              lowerCAmelCase	=     {0: 'batch', 1: 'decoder_sequence'}
              lowerCAmelCase	=     {0: 'batch', 1: 'decoder_sequence'}

           if self.use_past:
              self.fill_with_past_key_values_(_SCREAMING_SNAKE_CASE       , direction='inputs'       )
        elif self.task == "causal-lm":
           # TODO: figure this case out.
           lowerCAmelCase	=     OrderedDict(
               [
                   ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                   ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
               ]       )
           if self.use_past:
              lowerCAmelCase	,    lowerCAmelCase	=     self.num_layers
              for i in range(_SCREAMING_SNAKE_CASE       ):
                 lowerCAmelCase	=     {0: 'batch', 2: 'past_sequence + sequence'}
                 lowerCAmelCase	=     {0: 'batch', 2: 'past_sequence + sequence'}
        else:
           lowerCAmelCase	=     OrderedDict(
               [
                   ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                   ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
                   ('decoder_input_ids', {0: 'batch', 1: 'decoder_sequence'}),
                   ('decoder_attention_mask', {0: 'batch', 1: 'decoder_sequence'}),
               ]       )

        return common_inputs





     @property
     def 				__UpperCAmelCase (					self			: str       )							->    str:




        if self.task in ["default", "seq2seq-lm"]:
           lowerCAmelCase	=     super().outputs
        else:
           lowerCAmelCase	=     super(_SCREAMING_SNAKE_CASE       , self       ).outputs
           if self.use_past:
              lowerCAmelCase	,    lowerCAmelCase	=     self.num_layers
              for i in range(_SCREAMING_SNAKE_CASE       ):
                 lowerCAmelCase	=     {0: 'batch', 2: 'past_sequence + sequence'}
                 lowerCAmelCase	=     {0: 'batch', 2: 'past_sequence + sequence'}
        return common_outputs





     def 				__UpperCAmelCase (					self			: List[Any]       , UpperCAmelCase__			: PreTrainedTokenizer       , UpperCAmelCase__			: int = -1       , UpperCAmelCase__			: int = -1       , UpperCAmelCase__			: bool = False       , UpperCAmelCase__			: Optional[TensorType] = None       , )							->    List[Any]:
        lowerCAmelCase	=     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       )

        # Generate decoder inputs
        lowerCAmelCase	=     seq_length if not self.use_past else 1
        lowerCAmelCase	=     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       )
        lowerCAmelCase	=     {F'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
        lowerCAmelCase	=     dict(**_SCREAMING_SNAKE_CASE       , **_SCREAMING_SNAKE_CASE       )

        if self.use_past:
           if not is_torch_available():
              raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'       )
           else:
              import torch
           lowerCAmelCase	,    lowerCAmelCase	=     common_inputs['input_ids'].shape
           lowerCAmelCase	=     common_inputs['decoder_input_ids'].shape[1]
           lowerCAmelCase	,    lowerCAmelCase	=     self.num_attention_heads
           lowerCAmelCase	=     (
               batch,
               num_encoder_attention_heads,
               encoder_seq_length,
               self._config.hidden_size // num_encoder_attention_heads,
           )
           lowerCAmelCase	=     decoder_seq_length + 3
           lowerCAmelCase	=     (
               batch,
               num_decoder_attention_heads,
               decoder_past_length,
               self._config.hidden_size // num_decoder_attention_heads,
           )

           lowerCAmelCase	=     torch.cat(
               [common_inputs['decoder_attention_mask'], torch.ones(_SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       )]       , dim=1       )

           lowerCAmelCase	=     []
           # If the number of encoder and decoder layers are present in the model configuration, both are considered
           lowerCAmelCase	,    lowerCAmelCase	=     self.num_layers
           lowerCAmelCase	=     min(_SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       )
           lowerCAmelCase	=     max(_SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       ) - min_num_layers
           lowerCAmelCase	=     'encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

           for _ in range(_SCREAMING_SNAKE_CASE       ):
              common_inputs["past_key_values"].append(
                  (
                      torch.zeros(_SCREAMING_SNAKE_CASE       ),
                      torch.zeros(_SCREAMING_SNAKE_CASE       ),
                      torch.zeros(_SCREAMING_SNAKE_CASE       ),
                      torch.zeros(_SCREAMING_SNAKE_CASE       ),
                  )       )
           # TODO: test this.
           lowerCAmelCase	=     encoder_shape if remaining_side_name == 'encoder' else decoder_shape
           for _ in range(_SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       ):
              common_inputs["past_key_values"].append((torch.zeros(_SCREAMING_SNAKE_CASE       ), torch.zeros(_SCREAMING_SNAKE_CASE       ))       )
        return common_inputs





     def 				__UpperCAmelCase (					self			: int       , UpperCAmelCase__			: PreTrainedTokenizer       , UpperCAmelCase__			: int = -1       , UpperCAmelCase__			: int = -1       , UpperCAmelCase__			: bool = False       , UpperCAmelCase__			: Optional[TensorType] = None       , )							->    Dict:
        lowerCAmelCase	=     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       )

        if self.use_past:
           if not is_torch_available():
              raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'       )
           else:
              import torch
           lowerCAmelCase	,    lowerCAmelCase	=     common_inputs['input_ids'].shape
           # Not using the same length for past_key_values
           lowerCAmelCase	=     seqlen + 2
           lowerCAmelCase	,    lowerCAmelCase	=     self.num_layers
           lowerCAmelCase	,    lowerCAmelCase	=     self.num_attention_heads
           lowerCAmelCase	=     (
               batch,
               num_encoder_attention_heads,
               past_key_values_length,
               self._config.hidden_size // num_encoder_attention_heads,
           )

           lowerCAmelCase	=     common_inputs['attention_mask'].dtype
           lowerCAmelCase	=     torch.cat(
               [common_inputs['attention_mask'], torch.ones(_SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , dtype=_SCREAMING_SNAKE_CASE       )]       , dim=1       )
           lowerCAmelCase	=     [
               (torch.zeros(_SCREAMING_SNAKE_CASE       ), torch.zeros(_SCREAMING_SNAKE_CASE       )) for _ in range(_SCREAMING_SNAKE_CASE       )
           ]
        return common_inputs





     def 				__UpperCAmelCase (					self			: Union[str, Any]       , UpperCAmelCase__			: PreTrainedTokenizer       , UpperCAmelCase__			: int = -1       , UpperCAmelCase__			: int = -1       , UpperCAmelCase__			: bool = False       , UpperCAmelCase__			: Optional[TensorType] = None       , )							->    Union[str, Any]:
        lowerCAmelCase	=     compute_effective_axis_dimension(
            _SCREAMING_SNAKE_CASE       , fixed_dimension=OnnxConfig.default_fixed_batch       , num_token_to_add=0       )

        # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
        lowerCAmelCase	=     tokenizer.num_special_tokens_to_add(_SCREAMING_SNAKE_CASE       )
        lowerCAmelCase	=     compute_effective_axis_dimension(
            _SCREAMING_SNAKE_CASE       , fixed_dimension=OnnxConfig.default_fixed_sequence       , num_token_to_add=_SCREAMING_SNAKE_CASE       )

        # Generate dummy inputs according to compute batch and sequence
        lowerCAmelCase	=     [' '.join([tokenizer.unk_token]       ) * seq_length] * batch_size
        lowerCAmelCase	=     dict(tokenizer(_SCREAMING_SNAKE_CASE       , return_tensors=_SCREAMING_SNAKE_CASE       )       )
        return common_inputs





     def 				__UpperCAmelCase (					self			: Dict       , UpperCAmelCase__			: PreTrainedTokenizer       , UpperCAmelCase__			: int = -1       , UpperCAmelCase__			: int = -1       , UpperCAmelCase__			: bool = False       , UpperCAmelCase__			: Optional[TensorType] = None       , )							->    Any:




        if self.task in ["default", "seq2seq-lm"]:
           lowerCAmelCase	=     self._generate_dummy_inputs_for_default_and_seqaseq_lm(
               _SCREAMING_SNAKE_CASE       , batch_size=_SCREAMING_SNAKE_CASE       , seq_length=_SCREAMING_SNAKE_CASE       , is_pair=_SCREAMING_SNAKE_CASE       , framework=_SCREAMING_SNAKE_CASE       )

        elif self.task == "causal-lm":
           lowerCAmelCase	=     self._generate_dummy_inputs_for_causal_lm(
               _SCREAMING_SNAKE_CASE       , batch_size=_SCREAMING_SNAKE_CASE       , seq_length=_SCREAMING_SNAKE_CASE       , is_pair=_SCREAMING_SNAKE_CASE       , framework=_SCREAMING_SNAKE_CASE       )
        else:
           lowerCAmelCase	=     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
               _SCREAMING_SNAKE_CASE       , batch_size=_SCREAMING_SNAKE_CASE       , seq_length=_SCREAMING_SNAKE_CASE       , is_pair=_SCREAMING_SNAKE_CASE       , framework=_SCREAMING_SNAKE_CASE       )

        return common_inputs





     def 				__UpperCAmelCase (					self			: Dict       , UpperCAmelCase__			: Tuple       , UpperCAmelCase__			: int       , UpperCAmelCase__			: int       , UpperCAmelCase__			: Any       )							->    str:




        if self.task in ["default", "seq2seq-lm"]:
           lowerCAmelCase	=     super()._flatten_past_key_values_(_SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       )
        else:
           lowerCAmelCase	=     super(_SCREAMING_SNAKE_CASE       , self       )._flatten_past_key_values_(
               _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE       )













import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


__magic_name__   :  Optional[int]        =							logging.get_logger(__name__)
def 		lowercase__  (				_UpperCamelCase) ->			Dict:
      """simple docstring"""

      UpperCamelCase					       =  r'\w+[.]\d+'
      UpperCamelCase					       =  re.findall(_UpperCamelCase       ,							_UpperCamelCase)
      for pat in pats:
            UpperCamelCase					       =  key.replace(_UpperCamelCase       ,							'_'.join(pat.split('.')))
      return key
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase) ->			Tuple:
      """simple docstring"""

      UpperCamelCase					       =  pt_tuple_key[:-1] + ('scale',)
      if (
          any('norm' in str_ for str_ in pt_tuple_key)
          and (pt_tuple_key[-1] == "bias")
          and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
          and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
      ):
            UpperCamelCase					       =  pt_tuple_key[:-1] + ('scale',)
            return renamed_pt_tuple_key, pt_tensor
      elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
            UpperCamelCase					       =  pt_tuple_key[:-1] + ('scale',)
            return renamed_pt_tuple_key, pt_tensor

      # embedding
      if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
            UpperCamelCase					       =  pt_tuple_key[:-1] + ('embedding',)
            return renamed_pt_tuple_key, pt_tensor

      # conv layer
      UpperCamelCase					       =  pt_tuple_key[:-1] + ('kernel',)
      if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
            UpperCamelCase					       =  pt_tensor.transpose(2       ,							3       ,							1       ,							0)
            return renamed_pt_tuple_key, pt_tensor

      # linear layer
      UpperCamelCase					       =  pt_tuple_key[:-1] + ('kernel',)
      if pt_tuple_key[-1] == "weight":
            UpperCamelCase					       =  pt_tensor.T
            return renamed_pt_tuple_key, pt_tensor

      # old PyTorch layer norm weight
      UpperCamelCase					       =  pt_tuple_key[:-1] + ('weight',)
      if pt_tuple_key[-1] == "gamma":
            return renamed_pt_tuple_key, pt_tensor

      # old PyTorch layer norm bias
      UpperCamelCase					       =  pt_tuple_key[:-1] + ('bias',)
      if pt_tuple_key[-1] == "beta":
            return renamed_pt_tuple_key, pt_tensor

      return pt_tuple_key, pt_tensor
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase=42) ->			Optional[Any]:
      """simple docstring"""

      UpperCamelCase					       =  {k: v.numpy() for k, v in pt_state_dict.items()}

      # Step 2: Since the model is stateless, get random Flax params
      UpperCamelCase					       =  flax_model.init_weights(PRNGKey(_UpperCamelCase))

      UpperCamelCase					       =  flatten_dict(_UpperCamelCase)
      UpperCamelCase					       =  {}

      # Need to change some parameters name to match Flax names
      for pt_key, pt_tensor in pt_state_dict.items():
            UpperCamelCase					       =  rename_key(_UpperCamelCase)
            UpperCamelCase					       =  tuple(renamed_pt_key.split('.'))

            # Correctly rename weight parameters
            UpperCamelCase       ,  UpperCamelCase					       =  rename_key_and_reshape_tensor(_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase)

            if flax_key in random_flax_state_dict:
                  if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                        raise ValueError(
                            F'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
                            F'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.')

        # also add unexpected weight so that warning is thrown
            UpperCamelCase					       =  jnp.asarray(_UpperCamelCase)

      return unflatten_dict(_UpperCamelCase)



