
'''simple docstring'''


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_ :  List[str]     =					logging.get_logger(__name__)

snake_case_ :  Optional[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 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   _snake_case		   =							'''marian'''
   _snake_case		   =							['''past_key_values''']
   _snake_case		   =							{'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}







   def __init__(			self ,   lowerCamelCase__=5_8_1_0_1 ,   lowerCamelCase__=None ,   lowerCamelCase__=1_0_2_4 ,   lowerCamelCase__=1_2 ,   lowerCamelCase__=4_0_9_6 ,   lowerCamelCase__=1_6 ,   lowerCamelCase__=1_2 ,   lowerCamelCase__=4_0_9_6 ,   lowerCamelCase__=1_6 ,   lowerCamelCase__=0.0 ,   lowerCamelCase__=0.0 ,   lowerCamelCase__=True ,   lowerCamelCase__=True ,   lowerCamelCase__="gelu" ,   lowerCamelCase__=1_0_2_4 ,   lowerCamelCase__=0.1 ,   lowerCamelCase__=0.0 ,   lowerCamelCase__=0.0 ,   lowerCamelCase__=0.02 ,   lowerCamelCase__=5_8_1_0_0 ,   lowerCamelCase__=False ,   lowerCamelCase__=5_8_1_0_0 ,   lowerCamelCase__=0 ,   lowerCamelCase__=0 ,   lowerCamelCase__=True ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      UpperCamelCase       =  vocab_size
      UpperCamelCase       =  decoder_vocab_size or vocab_size
      UpperCamelCase       =  max_position_embeddings
      UpperCamelCase       =  d_model
      UpperCamelCase       =  encoder_ffn_dim
      UpperCamelCase       =  encoder_layers
      UpperCamelCase       =  encoder_attention_heads
      UpperCamelCase       =  decoder_ffn_dim
      UpperCamelCase       =  decoder_layers
      UpperCamelCase       =  decoder_attention_heads
      UpperCamelCase       =  dropout
      UpperCamelCase       =  attention_dropout
      UpperCamelCase       =  activation_dropout
      UpperCamelCase       =  activation_function
      UpperCamelCase       =  init_std
      UpperCamelCase       =  encoder_layerdrop
      UpperCamelCase       =  decoder_layerdrop
      UpperCamelCase       =  use_cache
      UpperCamelCase       =  encoder_layers
      UpperCamelCase       =  scale_embedding  # scale factor will be sqrt(d_model) if True
      UpperCamelCase       =  share_encoder_decoder_embeddings
      super().__init__(
          pad_token_id=lowerCamelCase__ ,   eos_token_id=lowerCamelCase__ ,   is_encoder_decoder=lowerCamelCase__ ,   decoder_start_token_id=lowerCamelCase__ ,   forced_eos_token_id=lowerCamelCase__ ,   **lowerCamelCase__ ,   )




class 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   @property
   # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







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

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

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

      '''simple docstring'''







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







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = -1 ,   lowerCamelCase__ = -1 ,   lowerCamelCase__ = False ,   lowerCamelCase__ = None ,   ):

      '''simple docstring'''







      UpperCamelCase       =  self._generate_dummy_inputs_for_encoder_and_decoder(
          lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )

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

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

         UpperCamelCase       =  torch.cat(
             [common_inputs['''decoder_attention_mask'''], torch.ones(lowerCamelCase__ ,   lowerCamelCase__   )] ,   dim=1   )

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

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







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = -1 ,   lowerCamelCase__ = -1 ,   lowerCamelCase__ = False ,   lowerCamelCase__ = None ,   ):

      '''simple docstring'''







      UpperCamelCase       =  self._generate_dummy_inputs_for_encoder_and_decoder(
          lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )

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

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







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = -1 ,   lowerCamelCase__ = -1 ,   lowerCamelCase__ = False ,   lowerCamelCase__ = None ,   ):

      '''simple docstring'''







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

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







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = -1 ,   lowerCamelCase__ = -1 ,   lowerCamelCase__ = False ,   lowerCamelCase__ = None ,   ):

      '''simple docstring'''







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

      else:
         UpperCamelCase       =  self._generate_dummy_inputs_for_causal_lm(
             lowerCamelCase__ ,   batch_size=lowerCamelCase__ ,   seq_length=lowerCamelCase__ ,   is_pair=lowerCamelCase__ ,   framework=lowerCamelCase__   )

      return common_inputs







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







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







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return 1e-4







'''simple docstring'''


import tensorflow as tf

from ...tf_utils import shape_list


class 						lowercase__	(  tf.keras.layers.Layer   ):






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=1 ,   lowerCamelCase__=False ,   **lowerCamelCase__   ):

      '''simple docstring'''







      super().__init__(**lowerCamelCase__   )

      UpperCamelCase       =  vocab_size
      UpperCamelCase       =  d_embed
      UpperCamelCase       =  d_proj

      UpperCamelCase       =  cutoffs + [vocab_size]
      UpperCamelCase       =  [0] + self.cutoffs
      UpperCamelCase       =  div_val

      UpperCamelCase       =  self.cutoffs[0]
      UpperCamelCase       =  len(self.cutoffs   ) - 1
      UpperCamelCase       =  self.shortlist_size + self.n_clusters
      UpperCamelCase       =  keep_order

      UpperCamelCase       =  []
      UpperCamelCase       =  []







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      if self.n_clusters > 0:
         UpperCamelCase       =  self.add_weight(
             shape=(self.n_clusters, self.d_embed) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name='''cluster_weight'''   )
         UpperCamelCase       =  self.add_weight(
             shape=(self.n_clusters,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name='''cluster_bias'''   )

      if self.div_val == 1:
         for i in range(len(self.cutoffs   )   ):
            if self.d_proj != self.d_embed:
               UpperCamelCase       =  self.add_weight(
                   shape=(self.d_embed, self.d_proj) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_projs_._{i}' ,   )
               self.out_projs.append(lowerCamelCase__   )
            else:
               self.out_projs.append(lowerCamelCase__   )
            UpperCamelCase       =  self.add_weight(
                shape=(self.vocab_size, self.d_embed) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._weight' ,   )
            UpperCamelCase       =  self.add_weight(
                shape=(self.vocab_size,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._bias' ,   )
            self.out_layers.append((weight, bias)   )
      else:
         for i in range(len(self.cutoffs   )   ):
            UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
            UpperCamelCase       =  self.d_embed // (self.div_val**i)

            UpperCamelCase       =  self.add_weight(
                shape=(d_emb_i, self.d_proj) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_projs_._{i}'   )
            self.out_projs.append(lowerCamelCase__   )
            UpperCamelCase       =  self.add_weight(
                shape=(r_idx - l_idx, d_emb_i) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._weight' ,   )
            UpperCamelCase       =  self.add_weight(
                shape=(r_idx - l_idx,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._bias' ,   )
            self.out_layers.append((weight, bias)   )
      super().build(lowerCamelCase__   )







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None   ):

      '''simple docstring'''







      UpperCamelCase       =  x
      if proj is not None:
         UpperCamelCase       =  tf.einsum('''ibd,ed->ibe''' ,   lowerCamelCase__ ,   lowerCamelCase__   )
      return tf.einsum('''ibd,nd->ibn''' ,   lowerCamelCase__ ,   lowerCamelCase__   ) + b







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  shape_list(lowerCamelCase__   )
      UpperCamelCase       =  tf.range(lp_size[0] ,   dtype=target.dtype   )
      UpperCamelCase       =  tf.stack([r, target] ,   1   )
      return tf.gather_nd(lowerCamelCase__ ,   lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=True ,   lowerCamelCase__=False   ):

      '''simple docstring'''







      UpperCamelCase       =  0
      if self.n_clusters == 0:
         UpperCamelCase       =  self._logit(lowerCamelCase__ ,   self.out_layers[0][0] ,   self.out_layers[0][1] ,   self.out_projs[0]   )
         if target is not None:
            UpperCamelCase       =  tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowerCamelCase__ ,   logits=lowerCamelCase__   )
         UpperCamelCase       =  tf.nn.log_softmax(lowerCamelCase__ ,   axis=-1   )
      else:
         UpperCamelCase       =  shape_list(lowerCamelCase__   )
         UpperCamelCase       =  []
         UpperCamelCase       =  tf.zeros(hidden_sizes[:2]   )
         for i in range(len(self.cutoffs   )   ):
            UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
            if target is not None:
               UpperCamelCase       =  (target >= l_idx) & (target < r_idx)
               UpperCamelCase       =  tf.where(lowerCamelCase__   )
               UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   ) - l_idx

            if self.div_val == 1:
               UpperCamelCase       =  self.out_layers[0][0][l_idx:r_idx]
               UpperCamelCase       =  self.out_layers[0][1][l_idx:r_idx]
            else:
               UpperCamelCase       =  self.out_layers[i][0]
               UpperCamelCase       =  self.out_layers[i][1]

            if i == 0:
               UpperCamelCase       =  tf.concat([cur_W, self.cluster_weight] ,   0   )
               UpperCamelCase       =  tf.concat([cur_b, self.cluster_bias] ,   0   )

               UpperCamelCase       =  self._logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   self.out_projs[0]   )
               UpperCamelCase       =  tf.nn.log_softmax(lowerCamelCase__   )
               out.append(head_logprob[..., : self.cutoffs[0]]   )
               if target is not None:
                  UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   )
                  UpperCamelCase       =  self._gather_logprob(lowerCamelCase__ ,   lowerCamelCase__   )
            else:
               UpperCamelCase       =  self._logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   self.out_projs[i]   )
               UpperCamelCase       =  tf.nn.log_softmax(lowerCamelCase__   )
               UpperCamelCase       =  self.cutoffs[0] + i - 1  # No probability for the head cluster
               UpperCamelCase       =  head_logprob[..., cluster_prob_idx, None] + tail_logprob
               out.append(lowerCamelCase__   )
               if target is not None:
                  UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   )
                  UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   )
                  UpperCamelCase       =  self._gather_logprob(lowerCamelCase__ ,   lowerCamelCase__   )
                  cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
            if target is not None:
               loss += tf.scatter_nd(lowerCamelCase__ ,   -cur_logprob ,   shape_list(lowerCamelCase__   )   )
         UpperCamelCase       =  tf.concat(lowerCamelCase__ ,   axis=-1   )

      if target is not None:
         if return_mean:
            UpperCamelCase       =  tf.reduce_mean(lowerCamelCase__   )
         # Add the training-time loss value to the layer using `self.add_loss()`.
         self.add_loss(lowerCamelCase__   )

         # Log the loss as a metric (we could log arbitrary metrics,
         # including different metrics for training and inference.
         self.add_metric(lowerCamelCase__ ,   name=self.name ,   aggregation='''mean''' if return_mean else ''''''   )

      return out





