



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


__lowerCamelCase       			=   logging.get_logger(__name__)

__lowerCamelCase       			=   {
    """facebook/blenderbot-3B""": """https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json""",
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}



class 						_snake_case (     lowerCamelCase__			):





 """simple docstring"""



 lowerCamelCase_       =     'blenderbot'
 lowerCamelCase_       =     ['past_key_values']
 lowerCamelCase_       =     {'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}
 def __init__(  self		,      a=8_0_0_8		,      a=1_2_8		,      a=2		,      a=1_0_2_4_0		,      a=3_2		,      a=2_4		,      a=1_0_2_4_0		,      a=3_2		,      a=0.0		,      a=0.0		,      a=True		,      a=True		,      a="gelu"		,      a=2_5_6_0		,      a=0.1		,      a=0.0		,      a=0.0		,      a=0.02		,      a=1		,      a=False		,      a=0		,      a=1		,      a=2		,      a=3		,      a=2		,      **a		,      )		->							Any:




   """simple docstring"""







   _A          =      vocab_size
   _A          =      max_position_embeddings
   _A          =      d_model
   _A          =      encoder_ffn_dim
   _A          =      encoder_layers
   _A          =      encoder_attention_heads
   _A          =      decoder_ffn_dim
   _A          =      decoder_layers
   _A          =      decoder_attention_heads
   _A          =      dropout
   _A          =      attention_dropout
   _A          =      activation_dropout
   _A          =      activation_function
   _A          =      init_std
   _A          =      encoder_layerdrop
   _A          =      decoder_layerdrop
   _A          =      use_cache
   _A          =      encoder_layers
   _A          =      scale_embedding  # scale factor will be sqrt(d_model) if True

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







class 						_snake_case (     lowerCamelCase__			):





 """simple docstring"""



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




   """simple docstring"""







   if self.task in ["default", "seq2seq-lm"]:
     _A          =      OrderedDict(
         [
             ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
             ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
         ]   )
     if self.use_past:
       _A          =      {0: '''batch'''}
       _A          =      {0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
     else:
       _A          =      {0: '''batch''', 1: '''decoder_sequence'''}
       _A          =      {0: '''batch''', 1: '''decoder_sequence'''}
     if self.use_past:
       self.fill_with_past_key_values_(lowercase_		,      direction='''inputs'''   )
   elif self.task == "causal-lm":
     _A          =      OrderedDict(
         [
             ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
             ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
         ]   )
     if self.use_past:
       _A      ,				_A          =      self.num_layers
       for i in range(lowercase_   ):
         _A          =      {0: '''batch''', 2: '''past_sequence + sequence'''}
         _A          =      {0: '''batch''', 2: '''past_sequence + sequence'''}
   else:
     _A          =      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
 # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
 def lowercase_						(  self   )		->							Mapping[str, Mapping[int, str]]:




   """simple docstring"""







   if self.task in ["default", "seq2seq-lm"]:
     _A          =      super().outputs
   else:
     _A          =      super(lowercase_		,      self   ).outputs
     if self.use_past:
       _A      ,				_A          =      self.num_layers
       for i in range(lowercase_   ):
         _A          =      {0: '''batch''', 2: '''past_sequence + sequence'''}
         _A          =      {0: '''batch''', 2: '''past_sequence + sequence'''}
   return common_outputs
 def lowercase_						(  self		,      a		,      a = -1		,      a = -1		,      a = False		,      a = None		,      )		->							Mapping[str, Any]:




   """simple docstring"""







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

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

     for _ in range(lowercase_   ):
       common_inputs["past_key_values"].append(
           (
               torch.zeros(lowercase_   ),
               torch.zeros(lowercase_   ),
               torch.zeros(lowercase_   ),
               torch.zeros(lowercase_   ),
           )   )
   return common_inputs
 def lowercase_						(  self		,      a		,      a = -1		,      a = -1		,      a = False		,      a = None		,      )		->							Mapping[str, Any]:




   """simple docstring"""







   _A          =      self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
       lowercase_		,      lowercase_		,      lowercase_		,      lowercase_		,      lowercase_   )

   if self.use_past:
     if not is_torch_available():
       raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''   )
     else:
       import torch
     _A      ,				_A          =      common_inputs['''input_ids'''].shape
     _A          =      seqlen
     _A      ,				_A          =      self.num_layers
     _A      ,				_A          =      self.num_attention_heads
     _A          =      (
         batch,
         num_encoder_attention_heads,
         past_key_values_length,
         self._config.hidden_size // num_encoder_attention_heads,
     )
     _A          =      common_inputs['''attention_mask'''].dtype
     _A          =      torch.cat(
         [common_inputs['''attention_mask'''], torch.ones(lowercase_		,      lowercase_		,      dtype=lowercase_   )]		,      dim=1   )
     _A          =      [
         (torch.zeros(lowercase_   ), torch.zeros(lowercase_   )) for _ in range(lowercase_   )
     ]
   return common_inputs
 def lowercase_						(  self		,      a		,      a = -1		,      a = -1		,      a = False		,      a = None		,      )		->							Mapping[str, Any]:




   """simple docstring"""







   _A          =      compute_effective_axis_dimension(
       lowercase_		,      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
   _A          =      tokenizer.num_special_tokens_to_add(lowercase_   )
   _A          =      compute_effective_axis_dimension(
       lowercase_		,      fixed_dimension=OnnxConfig.default_fixed_sequence		,      num_token_to_add=lowercase_   )

   # Generate dummy inputs according to compute batch and sequence
   _A          =      [''' '''.join([tokenizer.unk_token]   ) * seq_length] * batch_size
   _A          =      dict(tokenizer(lowercase_		,      return_tensors=lowercase_   )   )
   return common_inputs
 def lowercase_						(  self		,      a		,      a = -1		,      a = -1		,      a = False		,      a = None		,      )		->							Mapping[str, Any]:




   """simple docstring"""







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

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

   return common_inputs
 def lowercase_						(  self		,      a		,      a		,      a		,      a   )		->							Optional[int]:




   """simple docstring"""







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



 def lowercase_						(  self		,      a		,      a   )		->							Tuple:




   """simple docstring"""







   if direction not in ["inputs", "outputs"]:
     raise ValueError(f'''direction must either be \"inputs\" or \"outputs\", but {direction} was given'''   )

   _A          =      '''past_key_values''' if direction == '''inputs''' else '''present'''
   _A      ,				_A          =      self.num_layers

   _A          =      '''past_encoder_sequence'''
   _A          =      '''past_decoder_sequence''' if direction == '''inputs''' else '''past_decoder_sequence + sequence'''

   for i in range(lowercase_   ):
     _A          =      {0: '''batch''', 2: decoder_sequence}
     _A          =      {0: '''batch''', 2: decoder_sequence}
     _A          =      {0: '''batch''', 2: encoder_sequence}
     _A          =      {0: '''batch''', 2: encoder_sequence}




import importlib.metadata
import subprocess
import sys




def       UpperCAmelCase__	(							__snake_case = False							)				->	List[Any]:
  _A          =      False
  if "IPython" in sys.modules:
    _A          =      '''google.colab''' in str(sys.modules['''IPython'''].get_ipython()							)

  if in_colab:
    if upgrade:
      _A          =      ['''pip''', '''install''', '''-U''', '''torch''']
      subprocess.run(__snake_case  ,			check=__snake_case							)
    # get the current version of torch
    _A          =      importlib.metadata.version('''torch'''							)
    _A          =      torch_version[: torch_version.rindex('''.'''							)]
    _A          =      F'''https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl'''
    _A          =      ['''pip''', '''install''', xla_wheel]
    subprocess.run(__snake_case  ,			check=__snake_case							)
  else:
    raise RuntimeError('''`install_xla` utility works only on google colab.'''							)