


from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings

@maybe_allow_in_graph
class 		__a				(    nn.Module      ):
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=0.0      ,   _SCREAMING_SNAKE_CASE = None      ,   _SCREAMING_SNAKE_CASE = "geglu"      ,   _SCREAMING_SNAKE_CASE = None      ,   _SCREAMING_SNAKE_CASE = False      ,   _SCREAMING_SNAKE_CASE = False      ,   _SCREAMING_SNAKE_CASE = False      ,   _SCREAMING_SNAKE_CASE = False      ,   _SCREAMING_SNAKE_CASE = True      ,   _SCREAMING_SNAKE_CASE = "layer_norm"      ,   _SCREAMING_SNAKE_CASE = False      ,   )			->	Optional[Any]:




									"""simple docstring"""





									super().__init__()
									_UpperCAmelCase					=      only_cross_attention

									_UpperCAmelCase					=      (num_embeds_ada_norm is not None) and norm_type == 'ada_norm_zero'
									_UpperCAmelCase					=      (num_embeds_ada_norm is not None) and norm_type == 'ada_norm'

									if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
																raise ValueError(
																    f'''`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to'''
																    f''' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.'''						)

									# Define 3 blocks. Each block has its own normalization layer.
									# 1. Self-Attn
									if self.use_ada_layer_norm:
																_UpperCAmelCase					=      AdaLayerNorm(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									elif self.use_ada_layer_norm_zero:
																_UpperCAmelCase					=      AdaLayerNormZero(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									else:
																_UpperCAmelCase					=      nn.LayerNorm(_SCREAMING_SNAKE_CASE      ,   elementwise_affine=_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      Attention(
									    query_dim=_SCREAMING_SNAKE_CASE      ,   heads=_SCREAMING_SNAKE_CASE      ,   dim_head=_SCREAMING_SNAKE_CASE      ,   dropout=_SCREAMING_SNAKE_CASE      ,   bias=_SCREAMING_SNAKE_CASE      ,   cross_attention_dim=cross_attention_dim if only_cross_attention else None      ,   upcast_attention=_SCREAMING_SNAKE_CASE      ,   )

									# 2. Cross-Attn
									if cross_attention_dim is not None or double_self_attention:
																# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
																# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
																# the second cross attention block.
																_UpperCAmelCase					=      (
																    AdaLayerNorm(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
																    if self.use_ada_layer_norm
																    else nn.LayerNorm(_SCREAMING_SNAKE_CASE      ,   elementwise_affine=_SCREAMING_SNAKE_CASE						)
																)
																_UpperCAmelCase					=      Attention(
																    query_dim=_SCREAMING_SNAKE_CASE      ,   cross_attention_dim=cross_attention_dim if not double_self_attention else None      ,   heads=_SCREAMING_SNAKE_CASE      ,   dim_head=_SCREAMING_SNAKE_CASE      ,   dropout=_SCREAMING_SNAKE_CASE      ,   bias=_SCREAMING_SNAKE_CASE      ,   upcast_attention=_SCREAMING_SNAKE_CASE      ,   )  # is self-attn if encoder_hidden_states is none
									else:
																_UpperCAmelCase					=      None
																_UpperCAmelCase					=      None

									# 3. Feed-forward
									_UpperCAmelCase					=      nn.LayerNorm(_SCREAMING_SNAKE_CASE      ,   elementwise_affine=_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      FeedForward(_SCREAMING_SNAKE_CASE      ,   dropout=_SCREAMING_SNAKE_CASE      ,   activation_fn=_SCREAMING_SNAKE_CASE      ,   final_dropout=_SCREAMING_SNAKE_CASE						)

									# let chunk size default to None
									_UpperCAmelCase					=      None
									_UpperCAmelCase					=      0
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Dict:




									"""simple docstring"""





									_UpperCAmelCase					=      chunk_size
									_UpperCAmelCase					=      dim






		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE = None      ,   _SCREAMING_SNAKE_CASE = None      ,   _SCREAMING_SNAKE_CASE = None      ,   _SCREAMING_SNAKE_CASE = None      ,   _SCREAMING_SNAKE_CASE = None      ,   _SCREAMING_SNAKE_CASE = None      ,   )			->	Dict:




									"""simple docstring"""





									if self.use_ada_layer_norm:
																_UpperCAmelCase					=      self.norma(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									elif self.use_ada_layer_norm_zero:
																_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase					=      self.norma(
																    _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   hidden_dtype=hidden_states.dtype						)
									else:
																_UpperCAmelCase					=      self.norma(_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      cross_attention_kwargs if cross_attention_kwargs is not None else {}

									_UpperCAmelCase					=      self.attna(
									    _SCREAMING_SNAKE_CASE      ,   encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE      ,   )
									if self.use_ada_layer_norm_zero:
																_UpperCAmelCase					=      gate_msa.unsqueeze(1						) * attn_output
									_UpperCAmelCase					=      attn_output + hidden_states

									# 2. Cross-Attention
									if self.attna is not None:
																_UpperCAmelCase					=      (
																    self.norma(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						) if self.use_ada_layer_norm else self.norma(_SCREAMING_SNAKE_CASE						)
																)

																_UpperCAmelCase					=      self.attna(
																    _SCREAMING_SNAKE_CASE      ,   encoder_hidden_states=_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE      ,   )
																_UpperCAmelCase					=      attn_output + hidden_states

									# 3. Feed-forward
									_UpperCAmelCase					=      self.norma(_SCREAMING_SNAKE_CASE						)

									if self.use_ada_layer_norm_zero:
																_UpperCAmelCase					=      norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

									if self._chunk_size is not None:
																# "feed_forward_chunk_size" can be used to save memory
																if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
																							raise ValueError(
																							    f'''`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.'''						)

																_UpperCAmelCase					=      norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
																_UpperCAmelCase					=      torch.cat(
																    [self.ff(_SCREAMING_SNAKE_CASE						) for hid_slice in norm_hidden_states.chunk(_SCREAMING_SNAKE_CASE      ,   dim=self._chunk_dim						)]      ,   dim=self._chunk_dim      ,   )
									else:
																_UpperCAmelCase					=      self.ff(_SCREAMING_SNAKE_CASE						)

									if self.use_ada_layer_norm_zero:
																_UpperCAmelCase					=      gate_mlp.unsqueeze(1						) * ff_output

									_UpperCAmelCase					=      ff_output + hidden_states

									return hidden_states



class 		__a				(    nn.Module      ):
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE = None      ,   _SCREAMING_SNAKE_CASE = 4      ,   _SCREAMING_SNAKE_CASE = 0.0      ,   _SCREAMING_SNAKE_CASE = "geglu"      ,   _SCREAMING_SNAKE_CASE = False      ,   )			->	Any:




									"""simple docstring"""





									super().__init__()
									_UpperCAmelCase					=      int(dim * mult						)
									_UpperCAmelCase					=      dim_out if dim_out is not None else dim

									if activation_fn == "gelu":
																_UpperCAmelCase					=      GELU(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									if activation_fn == "gelu-approximate":
																_UpperCAmelCase					=      GELU(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   approximate='tanh'						)
									elif activation_fn == "geglu":
																_UpperCAmelCase					=      GEGLU(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									elif activation_fn == "geglu-approximate":
																_UpperCAmelCase					=      ApproximateGELU(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      nn.ModuleList([]						)
									# project in
									self.net.append(_SCREAMING_SNAKE_CASE						)
									# project dropout
									self.net.append(nn.Dropout(_SCREAMING_SNAKE_CASE						)						)
									# project out
									self.net.append(nn.Linear(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)						)
									# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
									if final_dropout:
																self.net.append(nn.Dropout(_SCREAMING_SNAKE_CASE						)						)






		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE						)			->	str:




									"""simple docstring"""





									for module in self.net:
																_UpperCAmelCase					=      module(_SCREAMING_SNAKE_CASE						)
									return hidden_states



class 		__a				(    nn.Module      ):
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE = "none"						)			->	List[Any]:




									"""simple docstring"""





									super().__init__()
									_UpperCAmelCase					=      nn.Linear(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      approximate
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[Any]:




									"""simple docstring"""





									if gate.device.type != "mps":
																return F.gelu(_SCREAMING_SNAKE_CASE      ,   approximate=self.approximate						)
									# mps: gelu is not implemented for float16
									return F.gelu(gate.to(dtype=torch.floataa						)      ,   approximate=self.approximate						).to(dtype=gate.dtype						)






		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE						)			->	List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.proj(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      self.gelu(_SCREAMING_SNAKE_CASE						)
									return hidden_states



class 		__a				(    nn.Module      ):
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	str:




									"""simple docstring"""





									super().__init__()
									_UpperCAmelCase					=      nn.Linear(_SCREAMING_SNAKE_CASE      ,   dim_out * 2						)
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE						)			->	Tuple:




									"""simple docstring"""





									if gate.device.type != "mps":
																return F.gelu(_SCREAMING_SNAKE_CASE						)
									# mps: gelu is not implemented for float16
									return F.gelu(gate.to(dtype=torch.floataa						)						).to(dtype=gate.dtype						)






		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase		,	_UpperCAmelCase					=      self.proj(_SCREAMING_SNAKE_CASE						).chunk(2      ,   dim=-1						)
									return hidden_states * self.gelu(_SCREAMING_SNAKE_CASE						)



class 		__a				(    nn.Module      ):
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[int]:




									"""simple docstring"""





									super().__init__()
									_UpperCAmelCase					=      nn.Linear(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)






		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE						)			->	int:




									"""simple docstring"""





									_UpperCAmelCase					=      self.proj(_SCREAMING_SNAKE_CASE						)
									return x * torch.sigmoid(1.702 * x						)



class 		__a				(    nn.Module      ):
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	int:




									"""simple docstring"""





									super().__init__()
									_UpperCAmelCase					=      nn.Embedding(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      nn.SiLU()
									_UpperCAmelCase					=      nn.Linear(_SCREAMING_SNAKE_CASE      ,   embedding_dim * 2						)
									_UpperCAmelCase					=      nn.LayerNorm(_SCREAMING_SNAKE_CASE      ,   elementwise_affine=_SCREAMING_SNAKE_CASE						)






		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase					=      self.linear(self.silu(self.emb(_SCREAMING_SNAKE_CASE						)						)						)
									_UpperCAmelCase		,	_UpperCAmelCase					=      torch.chunk(_SCREAMING_SNAKE_CASE      ,   2						)
									_UpperCAmelCase					=      self.norm(_SCREAMING_SNAKE_CASE						) * (1 + scale) + shift
									return x



class 		__a				(    nn.Module      ):
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Tuple:




									"""simple docstring"""





									super().__init__()

									_UpperCAmelCase					=      CombinedTimestepLabelEmbeddings(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      nn.SiLU()
									_UpperCAmelCase					=      nn.Linear(_SCREAMING_SNAKE_CASE      ,   6 * embedding_dim      ,   bias=_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      nn.LayerNorm(_SCREAMING_SNAKE_CASE      ,   elementwise_affine=_SCREAMING_SNAKE_CASE      ,   eps=1e-6						)






		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=None						)			->	int:




									"""simple docstring"""





									_UpperCAmelCase					=      self.linear(self.silu(self.emb(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   hidden_dtype=_SCREAMING_SNAKE_CASE						)						)						)
									_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase					=      emb.chunk(6      ,   dim=1						)
									_UpperCAmelCase					=      self.norm(_SCREAMING_SNAKE_CASE						) * (1 + scale_msa[:, None]) + shift_msa[:, None]
									return x, gate_msa, shift_mlp, scale_mlp, gate_mlp



class 		__a				(    nn.Module      ):
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE = None      ,   _SCREAMING_SNAKE_CASE = 1e-5						)			->	str:




									"""simple docstring"""





									super().__init__()
									_UpperCAmelCase					=      num_groups
									_UpperCAmelCase					=      eps

									if act_fn is None:
																_UpperCAmelCase					=      None
									else:
																_UpperCAmelCase					=      get_activation(_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      nn.Linear(_SCREAMING_SNAKE_CASE      ,   out_dim * 2						)






		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Union[str, Any]:




									"""simple docstring"""





									if self.act:
																_UpperCAmelCase					=      self.act(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      self.linear(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      emb[:, :, None, None]
									_UpperCAmelCase		,	_UpperCAmelCase					=      emb.chunk(2      ,   dim=1						)

									_UpperCAmelCase					=      F.group_norm(_SCREAMING_SNAKE_CASE      ,   self.num_groups      ,   eps=self.eps						)
									_UpperCAmelCase					=      x * (1 + scale) + shift
									return x




import copy
import tempfile
import unittest

from transformers import MaMaaaConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device
from transformers.utils import cached_property

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import MaMaaaForConditionalGeneration, MaMaaaModel, MaMaaaTokenizer
		from transformers.models.mam_aaa.modeling_mam_aaa import MaMaaaDecoder, MaMaaaEncoder

def 	lowerCAmelCase__			(      a__:		Dict      ,  a__:		Dict      ,  a__:		Any      ,  a__:		Optional[int]=None      ,  a__:		str=None      ,  a__:		List[Any]=None      ,  a__:		Optional[int]=None      ,  a__:		Union[str, Any]=None      ,  )  ->		Tuple:


							'''simple docstring'''





							if attention_mask is None:
														_UpperCAmelCase					=      input_ids.ne(config.pad_token_id	)
							if decoder_attention_mask is None:
														_UpperCAmelCase					=      decoder_input_ids.ne(config.pad_token_id	)
							if head_mask is None:
														_UpperCAmelCase					=      torch.ones(config.encoder_layers      ,  config.encoder_attention_heads      ,  device=a__	)
							if decoder_head_mask is None:
														_UpperCAmelCase					=      torch.ones(config.decoder_layers      ,  config.decoder_attention_heads      ,  device=a__	)
							if cross_attn_head_mask is None:
														_UpperCAmelCase					=      torch.ones(config.decoder_layers      ,  config.decoder_attention_heads      ,  device=a__	)
							return {
							    "input_ids": input_ids,
							    "decoder_input_ids": decoder_input_ids,
							    "attention_mask": attention_mask,
							    "decoder_attention_mask": attention_mask,
							    "head_mask": head_mask,
							    "decoder_head_mask": decoder_head_mask,
							    "cross_attn_head_mask": cross_attn_head_mask,
							}

class 		__a				:
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=13      ,   _SCREAMING_SNAKE_CASE=7      ,   _SCREAMING_SNAKE_CASE=True      ,   _SCREAMING_SNAKE_CASE=False      ,   _SCREAMING_SNAKE_CASE=99      ,   _SCREAMING_SNAKE_CASE=16      ,   _SCREAMING_SNAKE_CASE=2      ,   _SCREAMING_SNAKE_CASE=4      ,   _SCREAMING_SNAKE_CASE=4      ,   _SCREAMING_SNAKE_CASE="relu"      ,   _SCREAMING_SNAKE_CASE=0.1      ,   _SCREAMING_SNAKE_CASE=0.1      ,   _SCREAMING_SNAKE_CASE=0.0      ,   _SCREAMING_SNAKE_CASE=0.0      ,   _SCREAMING_SNAKE_CASE=20      ,   _SCREAMING_SNAKE_CASE=2      ,   _SCREAMING_SNAKE_CASE=1      ,   _SCREAMING_SNAKE_CASE=0      ,   )			->	Any:




									"""simple docstring"""





									_UpperCAmelCase					=      parent
									_UpperCAmelCase					=      batch_size
									_UpperCAmelCase					=      seq_length
									_UpperCAmelCase					=      is_training
									_UpperCAmelCase					=      use_labels
									_UpperCAmelCase					=      vocab_size
									_UpperCAmelCase					=      hidden_size
									_UpperCAmelCase					=      num_hidden_layers
									_UpperCAmelCase					=      num_attention_heads
									_UpperCAmelCase					=      intermediate_size
									_UpperCAmelCase					=      hidden_act
									_UpperCAmelCase					=      hidden_dropout_prob
									_UpperCAmelCase					=      attention_probs_dropout_prob
									_UpperCAmelCase					=      encoder_layerdrop
									_UpperCAmelCase					=      decoder_layerdrop
									_UpperCAmelCase					=      max_position_embeddings
									_UpperCAmelCase					=      eos_token_id
									_UpperCAmelCase					=      pad_token_id
									_UpperCAmelCase					=      bos_token_id
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase					=      ids_tensor([self.batch_size, self.seq_length]      ,   self.vocab_size						)
									_UpperCAmelCase					=      self.eos_token_id  # Eos Token
									_UpperCAmelCase					=      ids_tensor([self.batch_size, self.seq_length]      ,   self.vocab_size						)

									# we need to clamp the input ids here to avoid having pad token in between
									# this is because for M2M100 the position_ids are prepared such that
									# all pad tokens have pos id = 2 and rest are between 2..seq_length
									# and the seq_length here is seq_length - num_pad_tokens
									# but when using past, there is no way of knowing if the past input ids had
									# pad tokens in them, which results in incorrect seq_lenth and which in turn results in
									# position_ids being off by num_pad_tokens in past input
									_UpperCAmelCase					=      input_ids.clamp(self.pad_token_id + 1						)
									_UpperCAmelCase					=      decoder_input_ids.clamp(self.pad_token_id + 1						)

									_UpperCAmelCase					=      self.get_config()
									_UpperCAmelCase					=      prepare_mam_aaa_inputs_dict(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									return config, inputs_dict
		def      UpperCAmelCase__    (      self						)			->	Optional[Any]:




									"""simple docstring"""





									return MaMaaaConfig(
									    vocab_size=self.vocab_size      ,   d_model=self.hidden_size      ,   encoder_layers=self.num_hidden_layers      ,   decoder_layers=self.num_hidden_layers      ,   encoder_attention_heads=self.num_attention_heads      ,   decoder_attention_heads=self.num_attention_heads      ,   encoder_ffn_dim=self.intermediate_size      ,   decoder_ffn_dim=self.intermediate_size      ,   dropout=self.hidden_dropout_prob      ,   attention_dropout=self.attention_probs_dropout_prob      ,   encoder_layerdrop=self.encoder_layerdrop      ,   decoder_layerdrop=self.decoder_layerdrop      ,   max_position_embeddings=self.max_position_embeddings      ,   eos_token_id=self.eos_token_id      ,   bos_token_id=self.bos_token_id      ,   pad_token_id=self.pad_token_id      ,   )
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase		,	_UpperCAmelCase					=      self.prepare_config_and_inputs()
									return config, inputs_dict
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase					=      MaMaaaModel(config=_SCREAMING_SNAKE_CASE						).get_decoder().to(_SCREAMING_SNAKE_CASE						).eval()
									_UpperCAmelCase					=      inputs_dict['input_ids']
									_UpperCAmelCase					=      inputs_dict['attention_mask']
									_UpperCAmelCase					=      inputs_dict['head_mask']

									# first forward pass
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   head_mask=_SCREAMING_SNAKE_CASE      ,   use_cache=_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase		,	_UpperCAmelCase					=      outputs.to_tuple()

									# create hypothetical multiple next token and extent to next_input_ids
									_UpperCAmelCase					=      ids_tensor((self.batch_size, 3)      ,   config.vocab_size						)
									_UpperCAmelCase					=      ids_tensor((self.batch_size, 3)      ,   2						)

									# append to next input_ids and
									_UpperCAmelCase					=      torch.cat([input_ids, next_tokens]      ,   dim=-1						)
									_UpperCAmelCase					=      torch.cat([attention_mask, next_attn_mask]      ,   dim=-1						)

									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE						)['last_hidden_state']
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   past_key_values=_SCREAMING_SNAKE_CASE						)[
									    'last_hidden_state'
									]

									# select random slice
									_UpperCAmelCase					=      ids_tensor((1,)      ,   output_from_past.shape[-1]						).item()
									_UpperCAmelCase					=      output_from_no_past[:, -3:, random_slice_idx].detach()
									_UpperCAmelCase					=      output_from_past[:, :, random_slice_idx].detach()

									self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]						)

									# test that outputs are equal for slice
									self.parent.assertTrue(torch.allclose(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   atol=1e-2						)						)






		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=      MaMaaaModel(config=_SCREAMING_SNAKE_CASE						).to(_SCREAMING_SNAKE_CASE						).eval()
									_UpperCAmelCase					=      model(**_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      outputs.encoder_last_hidden_state
									_UpperCAmelCase					=      outputs.last_hidden_state

									with tempfile.TemporaryDirectory() as tmpdirname:
																_UpperCAmelCase					=      model.get_encoder()
																encoder.save_pretrained(_SCREAMING_SNAKE_CASE						)
																_UpperCAmelCase					=      MaMaaaEncoder.from_pretrained(_SCREAMING_SNAKE_CASE						).to(_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      encoder(inputs_dict['input_ids']      ,   attention_mask=inputs_dict['attention_mask']						)[
									    0
									]

									self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1e-3						)

									with tempfile.TemporaryDirectory() as tmpdirname:
																_UpperCAmelCase					=      model.get_decoder()
																decoder.save_pretrained(_SCREAMING_SNAKE_CASE						)
																_UpperCAmelCase					=      MaMaaaDecoder.from_pretrained(_SCREAMING_SNAKE_CASE						).to(_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      decoder(
									    input_ids=inputs_dict['decoder_input_ids']      ,   attention_mask=inputs_dict['decoder_attention_mask']      ,   encoder_hidden_states=_SCREAMING_SNAKE_CASE      ,   encoder_attention_mask=inputs_dict['attention_mask']      ,   )[0]

									self.parent.assertTrue((last_hidden_state_a - last_hidden_state).abs().max().item() < 1e-3						)



@require_torch
class 		__a				(    UpperCAmelCase   ,				UpperCAmelCase   ,				UpperCAmelCase   ,				unittest.TestCase      ):
		_a     :  List[Any]           =     (
		    (
		        MaMaaaModel,
		        MaMaaaForConditionalGeneration,
		    )
		    if is_torch_available()
		    else ()
		)
		_a     :  List[str]           =     (MaMaaaForConditionalGeneration,) if is_torch_available() else ()
		_a     :  int           =     (
		    {
		        'conversational': MaMaaaForConditionalGeneration,
		        'feature-extraction': MaMaaaModel,
		        'summarization': MaMaaaForConditionalGeneration,
		        'text2text-generation': MaMaaaForConditionalGeneration,
		        'translation': MaMaaaForConditionalGeneration,
		    }
		    if is_torch_available()
		    else {}
		)
		_a     :  str           =     True
		_a     :  Union[str, Any]           =     True
		_a     :  Optional[int]           =     False
		_a     :  Union[str, Any]           =     False
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	int:




									"""simple docstring"""





									if pipeline_test_casse_name == "TranslationPipelineTests":
																# Get `ValueError: Translation requires a `src_lang` and a `tgt_lang` for this model`.
																# `M2M100Config` was never used in pipeline tests: cannot create a simple tokenizer.
																return True

									return False
		def      UpperCAmelCase__    (      self						)			->	List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      MaMaaaModelTester(self						)
									_UpperCAmelCase					=      ConfigTester(self      ,   config_class=_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self						)			->	List[Any]:




									"""simple docstring"""





									self.config_tester.run_common_tests()
		def      UpperCAmelCase__    (      self						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase		,	_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									for model_class in self.all_model_classes:
																_UpperCAmelCase					=      model_class(_SCREAMING_SNAKE_CASE						)

																with tempfile.TemporaryDirectory() as tmpdirname:
																							model.save_pretrained(_SCREAMING_SNAKE_CASE						)
																							_UpperCAmelCase		,	_UpperCAmelCase					=      model_class.from_pretrained(_SCREAMING_SNAKE_CASE      ,   output_loading_info=_SCREAMING_SNAKE_CASE						)
																self.assertEqual(info['missing_keys']      ,   []						)
		def      UpperCAmelCase__    (      self						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_decoder_model_past_large_inputs(*_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs_for_common()
									self.model_tester.check_encoder_decoder_model_standalone(*_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase		,	_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in (MaMaaaModel, MaMaaaForConditionalGeneration):
																_UpperCAmelCase					=      model_class(_SCREAMING_SNAKE_CASE						)
																model.to(_SCREAMING_SNAKE_CASE						)
																model.eval()

																_UpperCAmelCase					=      copy.deepcopy(self._prepare_for_class(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)						)

																if not self.is_encoder_decoder:
																							_UpperCAmelCase					=      inputs['input_ids']
																							del inputs["input_ids"]
																else:
																							_UpperCAmelCase					=      inputs['input_ids']
																							_UpperCAmelCase					=      inputs.get('decoder_input_ids'      ,   _SCREAMING_SNAKE_CASE						)
																							del inputs["input_ids"]
																							inputs.pop('decoder_input_ids'      ,   _SCREAMING_SNAKE_CASE						)

																_UpperCAmelCase					=      model.get_input_embeddings()
																if not self.is_encoder_decoder:
																							_UpperCAmelCase					=      wte(_SCREAMING_SNAKE_CASE						)
																else:
																							_UpperCAmelCase					=      wte(_SCREAMING_SNAKE_CASE						)
																							_UpperCAmelCase					=      wte(_SCREAMING_SNAKE_CASE						)

																with torch.no_grad():
																							model(**_SCREAMING_SNAKE_CASE						)[0]






		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase		,	_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									_UpperCAmelCase					=      input_dict['input_ids']
									_UpperCAmelCase					=      input_ids.ne(1						).to(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      MaMaaaForConditionalGeneration(_SCREAMING_SNAKE_CASE						).eval().to(_SCREAMING_SNAKE_CASE						)
									if torch_device == "cuda":
																model.half()
									model.generate(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE						)
									model.generate(num_beams=4      ,   do_sample=_SCREAMING_SNAKE_CASE      ,   early_stopping=_SCREAMING_SNAKE_CASE      ,   num_return_sequences=3						)



def 	lowerCAmelCase__			(      a__:		Tuple	)  ->		Optional[int]:


							'''simple docstring'''





							return torch.tensor(a__      ,  dtype=torch.long      ,  device=a__	)


lowerCAmelCase__  :str             =							1e-4







@require_torch
@require_sentencepiece
@require_tokenizers
@slow
class 		__a				(    unittest.TestCase      ):
		@cached_property
		def      UpperCAmelCase__    (      self						)			->	List[str]:




									"""simple docstring"""





									return MaMaaaTokenizer.from_pretrained('facebook/m2m100_418M'						)
		def      UpperCAmelCase__    (      self						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      MaMaaaModel.from_pretrained('facebook/m2m100_418M'						).to(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      _long_tensor([[128028, 98, 12, 30527, 2732, 159, 7755, 61904, 39144, 38, 2]]						)
									_UpperCAmelCase					=      _long_tensor([[2, 128028, 98, 12, 30527, 2732, 159, 7755, 61904, 39144, 38]]						)
									_UpperCAmelCase					=      prepare_mam_aaa_inputs_dict(model.config      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									with torch.no_grad():
																_UpperCAmelCase					=      model(**_SCREAMING_SNAKE_CASE						)[0]
									_UpperCAmelCase					=      torch.Size((1, 11, 1024)						)
									self.assertEqual(output.shape      ,   _SCREAMING_SNAKE_CASE						)
									# change to expected output here
									_UpperCAmelCase					=      torch.tensor(
									    [[-0.7780, -0.1676, 0.1038], [-6.7556, -1.3992, 0.0567], [-7.5383, -0.5920, -0.2779]]      ,   device=_SCREAMING_SNAKE_CASE						)
									self.assertTrue(torch.allclose(output[:, :3, :3]      ,   _SCREAMING_SNAKE_CASE      ,   atol=_SCREAMING_SNAKE_CASE						)						)
		def      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





									_UpperCAmelCase					=      MaMaaaForConditionalGeneration.from_pretrained('facebook/m2m100_418M'						).to(_SCREAMING_SNAKE_CASE						)

									# change to intended input
									_UpperCAmelCase					=      _long_tensor([[128028, 98, 12, 30527, 2732, 159, 7755, 61904, 39144, 38, 2]]						)
									_UpperCAmelCase					=      _long_tensor([[2, 128028, 98, 12, 30527, 2732, 159, 7755, 61904, 39144, 38]]						)
									_UpperCAmelCase					=      prepare_mam_aaa_inputs_dict(model.config      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									with torch.no_grad():
																_UpperCAmelCase					=      model(**_SCREAMING_SNAKE_CASE						)[0]
									_UpperCAmelCase					=      torch.Size((1, 11, model.config.vocab_size)						)
									self.assertEqual(output.shape      ,   _SCREAMING_SNAKE_CASE						)
									# change to expected output here
									_UpperCAmelCase					=      torch.tensor(
									    [[-1.0448, -1.0411, 3.7992], [-3.2191, -3.2386, -1.3451], [-3.6210, -3.5993, 0.4925]]      ,   device=_SCREAMING_SNAKE_CASE						)
									self.assertTrue(torch.allclose(output[:, :3, :3]      ,   _SCREAMING_SNAKE_CASE      ,   atol=_SCREAMING_SNAKE_CASE						)						)






		def      UpperCAmelCase__    (      self						)			->	List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      MaMaaaForConditionalGeneration.from_pretrained('facebook/m2m100_418M'						).to(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      MaMaaaTokenizer.from_pretrained('facebook/m2m100_418M'      ,   src_lang='fr'      ,   tgt_lang='en'						)

									_UpperCAmelCase					=      [
									    'L\'affaire NSA souligne l\'absence totale de débat sur le renseignement',
									    'Selon moi, il y a deux niveaux de réponse de la part du gouvernement français.',
									    'Lorsque François Hollande téléphone à Barack Obama ou quand le ministre des affaires étrangères Laurent'
									    ' Fabius convoque l\'ambassadeur des Etats-Unis, ils réagissent à une vraie découverte, qui est celle de'
									    ' l\'ampleur de la surveillance américaine sur l\'ensemble des communications en France.',
									]

									# The below article tests that we don't add any hypotheses outside of the top n_beams
									_UpperCAmelCase					=      tokenizer(_SCREAMING_SNAKE_CASE      ,   padding=_SCREAMING_SNAKE_CASE      ,   return_tensors='pt'						)

									_UpperCAmelCase					=      model.generate(
									    input_ids=dct['input_ids'].to(_SCREAMING_SNAKE_CASE						)      ,   attention_mask=dct['attention_mask'].to(_SCREAMING_SNAKE_CASE						)      ,   num_beams=5      ,   forced_bos_token_id=tokenizer.get_lang_id('en'						)      ,   )

									_UpperCAmelCase					=      [
									    'The NSA case highlights the total absence of intelligence debate',
									    'I think there are two levels of response from the French government.',
									    'When François Hollande calls Barack Obama or when Foreign Minister Laurent Fabius calls the U.S.'
									    ' Ambassador, they respond to a real discovery, which is that of the scale of U.S. surveillance on all'
									    ' communications in France.',
									]

									_UpperCAmelCase					=      tokenizer.batch_decode(
									    hypotheses_batch.tolist()      ,   clean_up_tokenization_spaces=_SCREAMING_SNAKE_CASE      ,   skip_special_tokens=_SCREAMING_SNAKE_CASE						)
									assert generated == expected_en
