



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


a_      :      Optional[Any]    =     logging.get_logger(__name__)

a_      :      int    =     {
    'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/config.json',
    # See all Marian models at https://huggingface.co/models?filter=marian
}
class    __UpperCamelCase			(       _lowercase				):

 """simple docstring"""



 _lowercase   :    Union[str, Any]      						=					'''marian'''
 _lowercase   :    int      						=					['''past_key_values''']
 _lowercase   :    Dict      						=					{'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}

 def __init__(	self		,				SCREAMING_SNAKE_CASE=5_8_1_0_1		,				SCREAMING_SNAKE_CASE=None		,				SCREAMING_SNAKE_CASE=1_0_2_4		,				SCREAMING_SNAKE_CASE=1_2		,				SCREAMING_SNAKE_CASE=4_0_9_6		,				SCREAMING_SNAKE_CASE=1_6		,				SCREAMING_SNAKE_CASE=1_2		,				SCREAMING_SNAKE_CASE=4_0_9_6		,				SCREAMING_SNAKE_CASE=1_6		,				SCREAMING_SNAKE_CASE=0.0		,				SCREAMING_SNAKE_CASE=0.0		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE="gelu"		,				SCREAMING_SNAKE_CASE=1_0_2_4		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=0.0		,				SCREAMING_SNAKE_CASE=0.0		,				SCREAMING_SNAKE_CASE=0.02		,				SCREAMING_SNAKE_CASE=5_8_1_0_0		,				SCREAMING_SNAKE_CASE=False		,				SCREAMING_SNAKE_CASE=5_8_1_0_0		,				SCREAMING_SNAKE_CASE=0		,				SCREAMING_SNAKE_CASE=0		,				SCREAMING_SNAKE_CASE=True		,				**SCREAMING_SNAKE_CASE		,				) ->			Optional[Any]:
  a__												= vocab_size
  a__												= decoder_vocab_size or 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
  a__												= share_encoder_decoder_embeddings
  super().__init__(
      pad_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			(       _lowercase				):

 """simple docstring"""



 @property
 # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
 def 						_UpperCAmelCase       (	self      ) ->			Mapping[str, Mapping[int, str]]:
  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_(SCREAMING_SNAKE_CASE		,				direction='''inputs'''      )
  elif self.task == "causal-lm":
   # TODO: figure this case out.
   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(SCREAMING_SNAKE_CASE      ):
     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 						_UpperCAmelCase       (	self      ) ->			Mapping[str, Mapping[int, str]]:
  if self.task in ["default", "seq2seq-lm"]:
   a__												= super().outputs
  else:
   a__												= super(SCREAMING_SNAKE_CASE		,				self      ).outputs
   if self.use_past:
    a__							,			a__												= self.num_layers
    for i in range(SCREAMING_SNAKE_CASE      ):
     a__												= {0: '''batch''', 2: '''past_sequence + sequence'''}
     a__												= {0: '''batch''', 2: '''past_sequence + sequence'''}
  return common_outputs

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE = -1		,				SCREAMING_SNAKE_CASE = -1		,				SCREAMING_SNAKE_CASE = False		,				SCREAMING_SNAKE_CASE = None		,				) ->			Mapping[str, Any]:
  a__												= self._generate_dummy_inputs_for_encoder_and_decoder(
      SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )

  # Generate decoder inputs
  a__												= seq_length if not self.use_past else 1
  a__												= self._generate_dummy_inputs_for_encoder_and_decoder(
      SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
  a__												= {f"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
  a__												= 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
   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 + 3
   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(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )]		,				dim=1      )

   a__												= []
   # If the number of encoder and decoder layers are present in the model configuration, both are considered
   a__							,			a__												= self.num_layers
   a__												= min(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
   a__												= max(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) - min_num_layers
   a__												= '''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.
   a__												= 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		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE = -1		,				SCREAMING_SNAKE_CASE = -1		,				SCREAMING_SNAKE_CASE = False		,				SCREAMING_SNAKE_CASE = None		,				) ->			Mapping[str, Any]:
  a__												= self._generate_dummy_inputs_for_encoder_and_decoder(
      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
   a__							,			a__												= common_inputs['''input_ids'''].shape
   # Not using the same length for past_key_values
   a__												= seqlen + 2
   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(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				dtype=SCREAMING_SNAKE_CASE      )]		,				dim=1      )
   a__												= [
       (torch.zeros(SCREAMING_SNAKE_CASE      ), torch.zeros(SCREAMING_SNAKE_CASE      )) for _ in range(SCREAMING_SNAKE_CASE      )
   ]
  return common_inputs

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE = -1		,				SCREAMING_SNAKE_CASE = -1		,				SCREAMING_SNAKE_CASE = False		,				SCREAMING_SNAKE_CASE = None		,				) ->			Mapping[str, Any]:
  # Copied from OnnxConfig.generate_dummy_inputs
  # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
  # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
  a__												= 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
  a__												= tokenizer.num_special_tokens_to_add(SCREAMING_SNAKE_CASE      )
  a__												= 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
  a__												= [''' '''.join([tokenizer.unk_token]      ) * seq_length] * batch_size
  a__												= dict(tokenizer(SCREAMING_SNAKE_CASE		,				return_tensors=SCREAMING_SNAKE_CASE      )      )
  return common_inputs

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE = -1		,				SCREAMING_SNAKE_CASE = -1		,				SCREAMING_SNAKE_CASE = False		,				SCREAMING_SNAKE_CASE = None		,				) ->			Mapping[str, Any]:
  if self.task in ["default", "seq2seq-lm"]:
   a__												= 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      )

  else:
   a__												= 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      )

  return common_inputs

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			Optional[int]:
  if self.task in ["default", "seq2seq-lm"]:
   a__												= super()._flatten_past_key_values_(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
  else:
   a__												= super(SCREAMING_SNAKE_CASE		,				self      )._flatten_past_key_values_(
       SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )



 @property
 def 						_UpperCAmelCase       (	self      ) ->			float:
  return 1e-4










import logging
import os

from .state import PartialState
class    __UpperCamelCase			(       logging.LoggerAdapter				):

 """simple docstring"""



 @staticmethod
 def 						_UpperCAmelCase       (	SCREAMING_SNAKE_CASE      ) ->			Optional[Any]:
  a__												= PartialState()
  return not main_process_only or (main_process_only and state.is_main_process)



 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      ) ->			Optional[int]:
  if PartialState._shared_state == {}:
   raise RuntimeError(
       '''You must initialize the accelerate state by calling either `PartialState()` or `Accelerator()` before using the logging utility.'''      )
  a__												= kwargs.pop('''main_process_only'''		,				SCREAMING_SNAKE_CASE      )
  a__												= kwargs.pop('''in_order'''		,				SCREAMING_SNAKE_CASE      )

  if self.isEnabledFor(SCREAMING_SNAKE_CASE      ):
   if self._should_log(SCREAMING_SNAKE_CASE      ):
    a__							,			a__												= self.process(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
    self.logger.log(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      )

   elif in_order:
    a__												= PartialState()
    for i in range(state.num_processes      ):
     if i == state.process_index:
      a__							,			a__												= self.process(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
      self.logger.log(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      )
     state.wait_for_everyone()






def     __a					(		__UpperCAmelCase							,  __UpperCAmelCase = None       ):
 if log_level is None:
  a__												= os.environ.get('''ACCELERATE_LOG_LEVEL'''							,  __UpperCAmelCase       )
 a__												= logging.getLogger(__UpperCAmelCase       )
 if log_level is not None:
  logger.setLevel(log_level.upper()       )
  logger.root.setLevel(log_level.upper()       )
 return MultiProcessAdapter(__UpperCAmelCase							,  {}       )





