'''simple docstring'''

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

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


_lowerCamelCase :				Tuple														= logging.get_logger(__name__)

_lowerCamelCase :				List[str]														= {
    """facebook/m2m100_418M""": """https://huggingface.co/facebook/m2m100_418M/resolve/main/config.json""",
    # See all M2M100 models at https://huggingface.co/models?filter=m2m_100
}




class    lowerCamelCase__  (   __snake_case     ):
   __UpperCAmelCase           =						"""m2m_100"""
   __UpperCAmelCase           =						["""past_key_values"""]
   __UpperCAmelCase           =						{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}



   def __init__(    self					,      lowerCAmelCase__=128_112					,      lowerCAmelCase__=1_024					,      lowerCAmelCase__=12					,      lowerCAmelCase__=4_096					,      lowerCAmelCase__=16					,      lowerCAmelCase__=12					,      lowerCAmelCase__=4_096					,      lowerCAmelCase__=16					,      lowerCAmelCase__=0.05					,      lowerCAmelCase__=0.05					,      lowerCAmelCase__=True					,      lowerCAmelCase__=True					,      lowerCAmelCase__="relu"					,      lowerCAmelCase__=1_024					,      lowerCAmelCase__=0.1					,      lowerCAmelCase__=0.1					,      lowerCAmelCase__=0.0					,      lowerCAmelCase__=0.02					,      lowerCAmelCase__=2					,      lowerCAmelCase__=True					,      lowerCAmelCase__=1					,      lowerCAmelCase__=0					,      lowerCAmelCase__=2					,      **lowerCAmelCase__					,      )				->     str:
        """simple docstring"""





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

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




class    lowerCamelCase__  (   __snake_case     ):



   @property
   def  _UpperCamelCase		(    self     )				->     Mapping[str, Mapping[int, str]]:
        """simple docstring"""





        _UpperCamelCase  :Union[str, Any]       =OrderedDict(
            [
                ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
            ]     )

        if self.use_past:
             _UpperCamelCase  :List[Any]       ={0: """batch"""}
             _UpperCamelCase  :Optional[Any]       ={0: """batch""", 1: """past_decoder_sequence + sequence"""}
        else:
             _UpperCamelCase  :str       ={0: """batch""", 1: """decoder_sequence"""}
             _UpperCamelCase  :List[Any]       ={0: """batch""", 1: """decoder_sequence"""}

        if self.use_past:
             self.fill_with_past_key_values_(lowerCAmelCase__					,      direction="""inputs"""     )
        return common_inputs



   def  _UpperCamelCase		(    self					,      lowerCAmelCase__					,      lowerCAmelCase__ = -1					,      lowerCAmelCase__ = -1					,      lowerCAmelCase__ = False					,      lowerCAmelCase__ = None					,      )				->     Mapping[str, Any]:
        """simple docstring"""





        _UpperCamelCase  :Any       =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  :Dict       =tokenizer.num_special_tokens_to_add(lowerCAmelCase__     )
        _UpperCamelCase  :Dict       =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  :Optional[int]       =[""" """.join([tokenizer.unk_token]     ) * seq_length] * batch_size
        _UpperCamelCase  :int       =dict(tokenizer(lowerCAmelCase__					,      return_tensors=lowerCAmelCase__     )     )
        return common_inputs



   def  _UpperCamelCase		(    self					,      lowerCAmelCase__					,      lowerCAmelCase__ = -1					,      lowerCAmelCase__ = -1					,      lowerCAmelCase__ = False					,      lowerCAmelCase__ = None					,      )				->     Mapping[str, Any]:
        """simple docstring"""





        _UpperCamelCase  :Union[str, Any]       =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__     )

        # Generate decoder inputs
        _UpperCamelCase  :int       =seq_length if not self.use_past else 1
        _UpperCamelCase  :Optional[int]       =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__     )
        _UpperCamelCase  :List[str]       ={f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
        _UpperCamelCase  :Tuple       =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  :Any       =common_inputs["""input_ids"""].shape
             _UpperCamelCase  :Tuple       =common_inputs["""decoder_input_ids"""].shape[1]
             _UpperCamelCase       ,      _UpperCamelCase  :Any       =self.num_attention_heads
             _UpperCamelCase  :List[str]       =(
                 batch,
                 num_encoder_attention_heads,
                 encoder_seq_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )
             _UpperCamelCase  :Tuple       =decoder_seq_length + 3
             _UpperCamelCase  :Tuple       =(
                 batch,
                 num_decoder_attention_heads,
                 decoder_past_length,
                 self._config.hidden_size // num_decoder_attention_heads,
             )

             _UpperCamelCase  :Tuple       =torch.cat(
                 [common_inputs["""decoder_attention_mask"""], torch.ones(lowerCAmelCase__					,      lowerCAmelCase__     )]					,      dim=1     )

             _UpperCamelCase  :List[Any]       =[]
             # If the number of encoder and decoder layers are present in the model configuration, both are considered
             _UpperCamelCase       ,      _UpperCamelCase  :int       =self.num_layers
             _UpperCamelCase  :Dict       =min(lowerCAmelCase__					,      lowerCAmelCase__     )
             _UpperCamelCase  :Union[str, Any]       =max(lowerCAmelCase__					,      lowerCAmelCase__     ) - min_num_layers
             _UpperCamelCase  :Union[str, Any]       ="""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  :List[Any]       =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




   __UpperCAmelCase           =						_generate_dummy_inputs_for_default_and_seqaseq_lm
'''simple docstring'''

from typing import Optional

from ..utils.logging import get_logger
from .audio_classificiation import AudioClassification
from .automatic_speech_recognition import AutomaticSpeechRecognition
from .base import TaskTemplate
from .image_classification import ImageClassification
from .language_modeling import LanguageModeling
from .question_answering import QuestionAnsweringExtractive
from .summarization import Summarization
from .text_classification import TextClassification


_lowerCamelCase :				int														= [
    """AutomaticSpeechRecognition""",
    """AudioClassification""",
    """ImageClassification""",
    """LanguageModeling""",
    """QuestionAnsweringExtractive""",
    """Summarization""",
    """TaskTemplate""",
    """TextClassification""",
]

_lowerCamelCase :				int														= get_logger(__name__)


_lowerCamelCase :				List[str]														= {
    AutomaticSpeechRecognition.task: AutomaticSpeechRecognition,
    AudioClassification.task: AudioClassification,
    ImageClassification.task: ImageClassification,
    LanguageModeling.task: LanguageModeling,
    QuestionAnsweringExtractive.task: QuestionAnsweringExtractive,
    Summarization.task: Summarization,
    TextClassification.task: TextClassification,
}



def  _lowerCAmelCase				(				__a      )  ->      Optional[TaskTemplate]:
     '''simple docstring'''







     _UpperCamelCase  :List[str]       =task_template_dict.get("""task"""      )
     if task_name is None:
          logger.warning(F'''Couldn\'t find template for task \'{task_name}\'. Available templates: {list(__a      )}'''      )
          return None
     _UpperCamelCase  :List[str]       =NAMEaTEMPLATE.get(__a      )
     return template.from_dict(__a      )