


"""simple docstring"""



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

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


__snake_case   			=    logging.get_logger(__name__)

__snake_case   			=    {
    """facebook/bart-large""": """https://huggingface.co/facebook/bart-large/resolve/main/config.json""",
    # See all BART models at https://huggingface.co/models?filter=bart
}







class _lowerCAmelCase				(      __UpperCamelCase	):
   __UpperCAmelCase :       Optional[int]							  =  '''bart'''
   __UpperCAmelCase :       Union[str, Any]							  =  ['''past_key_values''']
   __UpperCAmelCase :       Union[str, Any]							  =  {'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}







   def __init__(	self						,	UpperCamelCase__=5_0265						,	UpperCamelCase__=1024						,	UpperCamelCase__=12						,	UpperCamelCase__=4096						,	UpperCamelCase__=16						,	UpperCamelCase__=12						,	UpperCamelCase__=4096						,	UpperCamelCase__=16						,	UpperCamelCase__=0.0						,	UpperCamelCase__=0.0						,	UpperCamelCase__="gelu"						,	UpperCamelCase__=1024						,	UpperCamelCase__=0.1						,	UpperCamelCase__=0.0						,	UpperCamelCase__=0.0						,	UpperCamelCase__=0.02						,	UpperCamelCase__=0.0						,	UpperCamelCase__=False						,	UpperCamelCase__=True						,	UpperCamelCase__=3						,	UpperCamelCase__=1						,	UpperCamelCase__=0						,	UpperCamelCase__=2						,	UpperCamelCase__=True						,	UpperCamelCase__=2						,	UpperCamelCase__=2						,	**UpperCamelCase__						,	)	->		Union[str, Any]:

       '''simple docstring'''




       snake_case    :      Union[str, Any]      					=    vocab_size
       snake_case    :      Tuple      					=    max_position_embeddings
       snake_case    :      Optional[Any]      					=    d_model
       snake_case    :      int      					=    encoder_ffn_dim
       snake_case    :      int      					=    encoder_layers
       snake_case    :      Optional[Any]      					=    encoder_attention_heads
       snake_case    :      Dict      					=    decoder_ffn_dim
       snake_case    :      Union[str, Any]      					=    decoder_layers
       snake_case    :      List[Any]      					=    decoder_attention_heads
       snake_case    :      List[str]      					=    dropout
       snake_case    :      List[Any]      					=    attention_dropout
       snake_case    :      Optional[Any]      					=    activation_dropout
       snake_case    :      str      					=    activation_function
       snake_case    :      List[str]      					=    init_std
       snake_case    :      Optional[Any]      					=    encoder_layerdrop
       snake_case    :      Optional[int]      					=    decoder_layerdrop
       snake_case    :      int      					=    classifier_dropout
       snake_case    :      Union[str, Any]      					=    use_cache
       snake_case    :      str      					=    encoder_layers
       snake_case    :      str      					=    scale_embedding  # scale factor will be sqrt(d_model) if True

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

       # ensure backward compatibility for BART CNN models
       if self.forced_bos_token_id is None and kwargs.get("force_bos_token_to_be_generated"						,	UpperCamelCase__						):
           snake_case    :      Dict      					=    self.bos_token_id
           warnings.warn(
               F'Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions. '
               "The config can simply be saved and uploaded again to be fixed."						)







class _lowerCAmelCase				(      __UpperCamelCase	):







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

       '''simple docstring'''




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

           if self.use_past:
               snake_case    :      Tuple      					=    {0: 'batch'}
               snake_case    :      Tuple      					=    {0: 'batch', 1: 'past_decoder_sequence + sequence'}
           else:
               snake_case    :      Optional[Any]      					=    {0: 'batch', 1: 'decoder_sequence'}
               snake_case    :      Dict      					=    {0: 'batch', 1: 'decoder_sequence'}

           if self.use_past:
               self.fill_with_past_key_values_(UpperCamelCase__						,	direction="inputs"						)
       elif self.task == "causal-lm":
           # TODO: figure this case out.
           snake_case    :      int      					=    OrderedDict(
               [
                   ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                   ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
               ]						)
           if self.use_past:
               snake_case    :      Optional[Any]      					=    self.num_layers
               for i in range(UpperCamelCase__						):
                   snake_case    :      List[Any]      					=    {0: 'batch', 2: 'past_sequence + sequence'}
                   snake_case    :      Optional[Any]      					=    {0: 'batch', 2: 'past_sequence + sequence'}
       else:
           snake_case    :      Union[str, Any]      					=    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  lowerCamelCase		(	self						)	->		Mapping[str, Mapping[int, str]]:

       '''simple docstring'''




       if self.task in ["default", "seq2seq-lm"]:
           snake_case    :      int      					=    super().outputs
       else:
           snake_case    :      Optional[Any]      					=    super(UpperCamelCase__						,	self						).outputs
           if self.use_past:
               snake_case    :      List[str]      					=    self.num_layers
               for i in range(UpperCamelCase__						):
                   snake_case    :      Union[str, Any]      					=    {0: 'batch', 2: 'past_sequence + sequence'}
                   snake_case    :      List[str]      					=    {0: 'batch', 2: 'past_sequence + sequence'}
       return common_outputs







   def  lowerCamelCase		(	self						,	UpperCamelCase__						,	UpperCamelCase__ = -1						,	UpperCamelCase__ = -1						,	UpperCamelCase__ = False						,	UpperCamelCase__ = None						,	)	->		Mapping[str, Any]:

       '''simple docstring'''




       snake_case    :      List[Any]      					=    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						)

       # Generate decoder inputs
       snake_case    :      List[str]      					=    seq_length if not self.use_past else 1
       snake_case    :      Optional[int]      					=    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						)
       snake_case    :      List[str]      					=    {F'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
       snake_case    :      Optional[int]      					=    dict(**UpperCamelCase__						,	**UpperCamelCase__						)

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

           snake_case    :      Dict      					=    torch.cat(
               [common_inputs["decoder_attention_mask"], torch.ones(UpperCamelCase__						,	UpperCamelCase__						)]						,	dim=1						)

           snake_case    :      Any      					=    []
           # If the number of encoder and decoder layers are present in the model configuration, both are considered
           snake_case    :      Tuple      					=    self.num_layers
           snake_case    :      List[Any]      					=    min(UpperCamelCase__						,	UpperCamelCase__						)
           snake_case    :      Tuple      					=    max(UpperCamelCase__						,	UpperCamelCase__						) - min_num_layers
           snake_case    :      str      					=    'encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

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







   def  lowerCamelCase		(	self						,	UpperCamelCase__						,	UpperCamelCase__ = -1						,	UpperCamelCase__ = -1						,	UpperCamelCase__ = False						,	UpperCamelCase__ = None						,	)	->		Mapping[str, Any]:

       '''simple docstring'''




       snake_case    :      Any      					=    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						)

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

           snake_case    :      Union[str, Any]      					=    common_inputs['attention_mask'].dtype
           snake_case    :      int      					=    torch.cat(
               [common_inputs["attention_mask"], torch.ones(UpperCamelCase__						,	UpperCamelCase__						,	dtype=UpperCamelCase__						)]						,	dim=1						)
           snake_case    :      Any      					=    [
               (torch.zeros(UpperCamelCase__						), torch.zeros(UpperCamelCase__						)) for _ in range(UpperCamelCase__						)
           ]
       return common_inputs







   def  lowerCamelCase		(	self						,	UpperCamelCase__						,	UpperCamelCase__ = -1						,	UpperCamelCase__ = -1						,	UpperCamelCase__ = False						,	UpperCamelCase__ = None						,	)	->		Mapping[str, Any]:

       '''simple docstring'''




       snake_case    :      Tuple      					=    compute_effective_axis_dimension(
           UpperCamelCase__						,	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
       snake_case    :      List[str]      					=    tokenizer.num_special_tokens_to_add(UpperCamelCase__						)
       snake_case    :      List[str]      					=    compute_effective_axis_dimension(
           UpperCamelCase__						,	fixed_dimension=OnnxConfig.default_fixed_sequence						,	num_token_to_add=UpperCamelCase__						)

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







   def  lowerCamelCase		(	self						,	UpperCamelCase__						,	UpperCamelCase__ = -1						,	UpperCamelCase__ = -1						,	UpperCamelCase__ = False						,	UpperCamelCase__ = None						,	)	->		Mapping[str, Any]:

       '''simple docstring'''




       if self.task in ["default", "seq2seq-lm"]:
           snake_case    :      List[Any]      					=    self._generate_dummy_inputs_for_default_and_seqaseq_lm(
               UpperCamelCase__						,	batch_size=UpperCamelCase__						,	seq_length=UpperCamelCase__						,	is_pair=UpperCamelCase__						,	framework=UpperCamelCase__						)

       elif self.task == "causal-lm":
           snake_case    :      int      					=    self._generate_dummy_inputs_for_causal_lm(
               UpperCamelCase__						,	batch_size=UpperCamelCase__						,	seq_length=UpperCamelCase__						,	is_pair=UpperCamelCase__						,	framework=UpperCamelCase__						)
       else:
           snake_case    :      Optional[int]      					=    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
               UpperCamelCase__						,	batch_size=UpperCamelCase__						,	seq_length=UpperCamelCase__						,	is_pair=UpperCamelCase__						,	framework=UpperCamelCase__						)

       return common_inputs







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

       '''simple docstring'''




       if self.task in ["default", "seq2seq-lm"]:
           snake_case    :      Dict      					=    super()._flatten_past_key_values_(UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						)
       else:
           snake_case    :      str      					=    super(UpperCamelCase__						,	self						)._flatten_past_key_values_(
               UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__						)












import os
import sys
import tempfile

import torch

from .state import AcceleratorState
from .utils import PrecisionType, PrepareForLaunch, is_mps_available, patch_environment


def        SCREAMING_SNAKE_CASE      (       _UpperCAmelCase						,	_UpperCAmelCase=()						,	_UpperCAmelCase=None						,	_UpperCAmelCase="no"						,	_UpperCAmelCase="29500" )			->     Tuple:
    lowerCamelCase__    :      Dict					 =		False
    lowerCamelCase__    :      Dict					 =		False
    if any(key.startswith('KAGGLE' ) for key in os.environ.keys() ):
        lowerCamelCase__    :      Optional[Any]					 =		True
    elif "IPython" in sys.modules:
        lowerCamelCase__    :      Optional[Any]					 =		'google.colab' in str(sys.modules['IPython'].get_ipython() )

    try:
        lowerCamelCase__    :      List[str]					 =		PrecisionType(mixed_precision.lower() )
    except ValueError:
        raise ValueError(
            F"""Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.""" )

    if (in_colab or in_kaggle) and (os.environ.get('TPU_NAME'						,	_UpperCAmelCase ) is not None):
        # TPU launch
        import torch_xla.distributed.xla_multiprocessing as xmp

        if len(AcceleratorState._shared_state ) > 0:
            raise ValueError(
                'To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside '
                'your training function. Restart your notebook and make sure no cells initializes an '
                '`Accelerator`.' )
        if num_processes is None:
            lowerCamelCase__    :      Optional[Any]					 =		8

        lowerCamelCase__    :      List[str]					 =		PrepareForLaunch(_UpperCAmelCase						,	distributed_type='TPU' )
        print(F"""Launching a training on {num_processes} TPU cores.""" )
        xmp.spawn(_UpperCAmelCase						,	args=_UpperCAmelCase						,	nprocs=_UpperCAmelCase						,	start_method='fork' )
    elif in_colab:
        # No need for a distributed launch otherwise as it's either CPU or one GPU.
        if torch.cuda.is_available():
            print('Launching training on one GPU.' )
        else:
            print('Launching training on one CPU.' )
        function(*_UpperCAmelCase )
    else:
        if num_processes is None:
            raise ValueError(
                'You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call.' )

        if num_processes > 1:
            # Multi-GPU launch
            from torch.multiprocessing import start_processes
            from torch.multiprocessing.spawn import ProcessRaisedException

            if len(AcceleratorState._shared_state ) > 0:
                raise ValueError(
                    'To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized '
                    'inside your training function. Restart your notebook and make sure no cells initializes an '
                    '`Accelerator`.' )

            if torch.cuda.is_initialized():
                raise ValueError(
                    'To launch a multi-GPU training from your notebook, you need to avoid running any instruction '
                    'using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA '
                    'function.' )

            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
            # process here (the other ones will be set be the launcher).
            with patch_environment(
                world_size=_UpperCAmelCase						,	master_addr='127.0.01'						,	master_port=_UpperCAmelCase						,	mixed_precision=_UpperCAmelCase ):
                lowerCamelCase__    :      Optional[int]					 =		PrepareForLaunch(_UpperCAmelCase						,	distributed_type='MULTI_GPU' )
                print(F"""Launching training on {num_processes} GPUs.""" )
                try:
                    start_processes(_UpperCAmelCase						,	args=_UpperCAmelCase						,	nprocs=_UpperCAmelCase						,	start_method='fork' )
                except ProcessRaisedException as e:
                    if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
                        raise RuntimeError(
                            'CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. '
                            'This likely stems from an outside import causing issues once the `notebook_launcher()` is called. '
                            'Please review your imports and test them when running the `notebook_launcher()` to identify '
                            'which one is problematic.' ) from e

        else:
            # No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
            if is_mps_available():
                lowerCamelCase__    :      int					 =		'1'
                print('Launching training on MPS.' )
            elif torch.cuda.is_available():
                print('Launching training on one GPU.' )
            else:
                print('Launching training on CPU.' )
            function(*_UpperCAmelCase )


def        SCREAMING_SNAKE_CASE      (       _UpperCAmelCase						,	_UpperCAmelCase=()						,	_UpperCAmelCase=2 )			->     Optional[Any]:
    from torch.multiprocessing import start_processes

    with tempfile.NamedTemporaryFile() as tmp_file:
        # torch.distributed will expect a few environment variable to be here. We set the ones common to each
        # process here (the other ones will be set be the launcher).
        with patch_environment(
            world_size=_UpperCAmelCase						,	master_addr='127.0.01'						,	master_port='29500'						,	accelerate_mixed_precision='no'						,	accelerate_debug_rdv_file=tmp_file.name						,	accelerate_use_cpu='yes'						,	):
            lowerCamelCase__    :      Optional[Any]					 =		PrepareForLaunch(_UpperCAmelCase						,	debug=_UpperCAmelCase )
            start_processes(_UpperCAmelCase						,	args=_UpperCAmelCase						,	nprocs=_UpperCAmelCase						,	start_method='fork' )





