






"""simple docstring"""


import json
import os
import re
import warnings
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowercase_								=	logging.get_logger(__name__)

lowercase_								=	{
    'source_spm': 'source.spm',
    'target_spm': 'target.spm',
    'vocab': 'vocab.json',
    'target_vocab_file': 'target_vocab.json',
    'tokenizer_config_file': 'tokenizer_config.json',
}

lowercase_								=	{
    'source_spm': {
        'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/source.spm'
    },
    'target_spm': {
        'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/target.spm'
    },
    'vocab': {
        'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/vocab.json'
    },
    'tokenizer_config_file': {
        'Helsinki-NLP/opus-mt-en-de': (
            'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/tokenizer_config.json'
        )
    },
}

lowercase_								=	{'Helsinki-NLP/opus-mt-en-de': 512}
lowercase_								=	{}


class    snake_case						(				__UpperCAmelCase		):


   '''simple docstring'''

   A_		:      int							=					VOCAB_FILES_NAMES
   A_		:      Tuple							=					PRETRAINED_VOCAB_FILES_MAP
   A_		:      Union[str, Any]							=					PRETRAINED_INIT_CONFIGURATION
   A_		:      List[Any]							=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   A_		:      List[str]							=					["input_ids", "attention_mask"]
   A_		:      Dict							=					re.compile(">>.+<<"		)  # type: re.Pattern


   def __init__(      self					:	Optional[Any],			_lowerCamelCase					:	Tuple,			_lowerCamelCase					:	Optional[int],			_lowerCamelCase					:	Optional[Any],			_lowerCamelCase					:	List[Any]=None,			_lowerCamelCase					:	int=None,			_lowerCamelCase					:	Optional[Any]=None,			_lowerCamelCase					:	Optional[int]="<unk>",			_lowerCamelCase					:	str="</s>",			_lowerCamelCase					:	Tuple="<pad>",			_lowerCamelCase					:	List[str]=5_12,			_lowerCamelCase					:	Optional[Dict[str, Any]] = None,			_lowerCamelCase					:	Union[str, Any]=False,			**_lowerCamelCase					:	int,			):

     '''simple docstring'''
     __A		     =					{} if sp_model_kwargs is None else sp_model_kwargs

     super().__init__(
         # bos_token=bos_token,  unused. Start decoding with config.decoder_start_token_id
         source_lang=__snake_case,			target_lang=__snake_case,			unk_token=__snake_case,			eos_token=__snake_case,			pad_token=__snake_case,			model_max_length=__snake_case,			sp_model_kwargs=self.sp_model_kwargs,			target_vocab_file=__snake_case,			separate_vocabs=__snake_case,			**__snake_case,			)
     assert Path(__snake_case    ).exists(), f'cannot find spm source {source_spm}'

     __A		     =					separate_vocabs
     __A		     =					load_json(__snake_case    )
     if self.unk_token not in self.encoder:
       raise KeyError('''<unk> token must be in vocab'''    )
     assert self.pad_token in self.encoder

     if separate_vocabs:
       __A		     =					load_json(__snake_case    )
       __A		     =					{v: k for k, v in self.target_encoder.items()}
       __A		     =					[]
     else:
       __A		     =					{v: k for k, v in self.encoder.items()}
       __A		     =					[k for k in self.encoder if k.startswith('''>>'''    ) and k.endswith('''<<'''    )]

     __A		     =					source_lang
     __A		     =					target_lang
     __A		     =					[source_spm, target_spm]

     # load SentencePiece model for pre-processing
     __A		     =					load_spm(__snake_case,			self.sp_model_kwargs    )
     __A		     =					load_spm(__snake_case,			self.sp_model_kwargs    )
     __A		     =					self.spm_source
     __A		     =					self.encoder

     # Multilingual target side: default to using first supported language code.

     self._setup_normalizer()


   def 							_SCREAMING_SNAKE_CASE							(      self					:	str    ):

     '''simple docstring'''





     try:
       from sacremoses import MosesPunctNormalizer

       __A		     =					MosesPunctNormalizer(self.source_lang    ).normalize
     except (ImportError, FileNotFoundError):
       warnings.warn('''Recommended: pip install sacremoses.'''    )
       __A		     =					lambda _lowerCamelCase    : x


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[int],			_lowerCamelCase					:	str    ):

     '''simple docstring'''
     return self.punc_normalizer(__snake_case    ) if x else ""


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[int],			_lowerCamelCase					:	Tuple    ):

     '''simple docstring'''
     return self.current_encoder.get(__snake_case,			self.current_encoder[self.unk_token]    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[str],			_lowerCamelCase					:	str    ):

     '''simple docstring'''
     __A		     =					self.language_code_re.match(__snake_case    )
     __A		     =					[match.group(0    )] if match else []
     return code, self.language_code_re.sub('''''',			__snake_case    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[int],			_lowerCamelCase					:	str    ):

     '''simple docstring'''
     __A		     =					self.remove_language_code(__snake_case    )
     __A		     =					self.current_spm.encode(__snake_case,			out_type=__snake_case    )
     return code + pieces


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Dict,			_lowerCamelCase					:	int    ):

     '''simple docstring'''
     return self.decoder.get(__snake_case,			self.unk_token    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	int,			_lowerCamelCase					:	Tuple,			**_lowerCamelCase					:	int    ):

     '''simple docstring'''
     return super().batch_decode(__snake_case,			**__snake_case    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Any,			_lowerCamelCase					:	Any,			**_lowerCamelCase					:	List[str]    ):

     '''simple docstring'''
     return super().decode(__snake_case,			**__snake_case    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	int,			_lowerCamelCase					:	List[str]    ):

     '''simple docstring'''
     __A		     =					self.spm_source if self._decode_use_source_tokenizer else self.spm_target
     __A		     =					[]
     __A		     =					''
     for token in tokens:
       # make sure that special tokens are not decoded using sentencepiece model
       if token in self.all_special_tokens:
         out_string += sp_model.decode_pieces(__snake_case    ) + token + " "
         __A		     =					[]
       else:
         current_sub_tokens.append(__snake_case    )
     out_string += sp_model.decode_pieces(__snake_case    )
     return out_string.strip()


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[Any],			_lowerCamelCase					:	Tuple,			_lowerCamelCase					:	Union[str, Any]=None    ):

     '''simple docstring'''





     if token_ids_a is None:
       return token_ids_a + [self.eos_token_id]
     # We don't expect to process pairs, but leave the pair logic for API consistency
     return token_ids_a + token_ids_a + [self.eos_token_id]


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[Any]    ):

     '''simple docstring'''
     __A		     =					self.spm_source
     __A		     =					self.encoder


   def 							_SCREAMING_SNAKE_CASE							(      self					:	int    ):

     '''simple docstring'''
     __A		     =					self.spm_target
     if self.separate_vocabs:
       __A		     =					self.target_encoder


   @property
   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[int]    ):

     '''simple docstring'''
     return len(self.encoder    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Tuple,			_lowerCamelCase					:	str,			_lowerCamelCase					:	Optional[str] = None    ):

     '''simple docstring'''





     if not os.path.isdir(__snake_case    ):
       logger.error(f'Vocabulary path ({save_directory}) should be a directory'    )
       return
     __A		     =					[]

     if self.separate_vocabs:
       __A		     =					os.path.join(
           __snake_case,			(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab'''],			)
       __A		     =					os.path.join(
           __snake_case,			(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''target_vocab_file'''],			)
       save_json(self.encoder,			__snake_case    )
       save_json(self.target_encoder,			__snake_case    )
       saved_files.append(__snake_case    )
       saved_files.append(__snake_case    )
     else:
       __A		     =					os.path.join(
           __snake_case,			(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab''']    )
       save_json(self.encoder,			__snake_case    )
       saved_files.append(__snake_case    )

     for spm_save_filename, spm_orig_path, spm_model in zip(
         [VOCAB_FILES_NAMES['''source_spm'''], VOCAB_FILES_NAMES['''target_spm''']],			self.spm_files,			[self.spm_source, self.spm_target],			):
       __A		     =					os.path.join(
           __snake_case,			(filename_prefix + '''-''' if filename_prefix else '''''') + spm_save_filename    )
       if os.path.abspath(__snake_case    ) != os.path.abspath(__snake_case    ) and os.path.isfile(__snake_case    ):
         copyfile(__snake_case,			__snake_case    )
         saved_files.append(__snake_case    )
       elif not os.path.isfile(__snake_case    ):
         with open(__snake_case,			'''wb'''    ) as fi:
           __A		     =					spm_model.serialized_model_proto()
           fi.write(__snake_case    )
         saved_files.append(__snake_case    )

     return tuple(__snake_case    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[Any]    ):

     '''simple docstring'''
     return self.get_src_vocab()


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Any    ):

     '''simple docstring'''
     return dict(self.encoder,			**self.added_tokens_encoder    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Dict    ):

     '''simple docstring'''
     return dict(self.target_encoder,			**self.added_tokens_decoder    )


   def __getstate__(      self					:	Union[str, Any]    ):

     '''simple docstring'''
     __A		     =					self.__dict__.copy()
     state.update(
         {k: None for k in ['''spm_source''', '''spm_target''', '''current_spm''', '''punc_normalizer''', '''target_vocab_file''']}    )
     return state


   def __setstate__(      self					:	str,			_lowerCamelCase					:	Dict    ):

     '''simple docstring'''
     __A		     =					d

     # for backward compatibility
     if not hasattr(self,			'''sp_model_kwargs'''    ):
       __A		     =					{}

     __A		     =					(load_spm(__snake_case,			self.sp_model_kwargs    ) for f in self.spm_files)
     __A		     =					self.spm_source
     self._setup_normalizer()


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[int],			*_lowerCamelCase					:	Tuple,			**_lowerCamelCase					:	Tuple    ):

     '''simple docstring'''
     return 1


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Any,			_lowerCamelCase					:	Optional[int]    ):

     '''simple docstring'''
     __A		     =					set(self.all_special_ids    )  # call it once instead of inside list comp
     all_special_ids.remove(self.unk_token_id    )  # <unk> is only sometimes special
     return [1 if x in all_special_ids else 0 for x in seq]






   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[Any],			_lowerCamelCase					:	List,			_lowerCamelCase					:	Optional[List] = None,			_lowerCamelCase					:	bool = False    ):

     '''simple docstring'''





     if already_has_special_tokens:
       return self._special_token_mask(__snake_case    )
     elif token_ids_a is None:
       return self._special_token_mask(__snake_case    ) + [1]
     else:
       return self._special_token_mask(token_ids_a + token_ids_a    ) + [1]





def    lowerCAmelCase							(       __UpperCamelCase						,			__UpperCamelCase					):




  """simple docstring"""







  __A		     =					sentencepiece.SentencePieceProcessor(**lowercase__					)
  spm.Load(lowercase__					)
  return spm



def    lowerCAmelCase							(       __UpperCamelCase						,			__UpperCamelCase					):




  """simple docstring"""







  with open(lowercase__						,			'''w'''					) as f:
    json.dump(lowercase__						,			lowercase__						,			indent=2					)





def    lowerCAmelCase							(       __UpperCamelCase					):




  """simple docstring"""







  with open(lowercase__						,			'''r'''					) as f:
    return json.load(lowercase__					)











'''simple docstring'''






import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import MBartConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
		import os

		# The slow tests are often failing with OOM error on GPU
		# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
		# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
		A :			Dict							   =      'platform'

		import jax
		import jax.numpy as jnp

		from transformers import AutoTokenizer
		from transformers.models.mbart.modeling_flax_mbart import (
		    FlaxMBartForConditionalGeneration,
		    FlaxMBartForQuestionAnswering,
		    FlaxMBartForSequenceClassification,
		    FlaxMBartModel,
		    shift_tokens_right,
		)





def    lowercase_   (		lowercase__					,	lowercase__					,	lowercase__=None					,	lowercase__=None					,	lowercase__=None					,	lowercase__=None					,	lowercase__=None					,	lowercase__=None					,	) ->List[Any]:
						if attention_mask is None:
												_snake_case: Optional[int]	       =		np.where(input_ids != config.pad_token_id					,	1					,	0    )
						if decoder_attention_mask is None:
												_snake_case: Optional[Any]	       =		np.where(decoder_input_ids != config.pad_token_id					,	1					,	0    )
						if head_mask is None:
												_snake_case: Union[str, Any]	       =		np.ones((config.encoder_layers, config.encoder_attention_heads)    )
						if decoder_head_mask is None:
												_snake_case: List[str]	       =		np.ones((config.decoder_layers, config.decoder_attention_heads)    )
						if cross_attn_head_mask is None:
												_snake_case: Union[str, Any]	       =		np.ones((config.decoder_layers, config.decoder_attention_heads)    )
						return {
						    "input_ids": input_ids,
						    "decoder_input_ids": decoder_input_ids,
						    "attention_mask": attention_mask,
						    "decoder_attention_mask": decoder_attention_mask,
						}







class       lowerCamelCase       (    unittest.TestCase	):


					def __init__(       self						:	Optional[Any]							,  __snake_case						:	Optional[Any]							,  __snake_case						:	List[Any]=13							,  __snake_case						:	Dict=7							,  __snake_case						:	int=True							,  __snake_case						:	Tuple=False							,  __snake_case						:	Optional[int]=99							,  __snake_case						:	Optional[Any]=16							,  __snake_case						:	Any=2							,  __snake_case						:	Dict=4							,  __snake_case						:	int=4							,  __snake_case						:	Any="gelu"							,  __snake_case						:	List[Any]=0.1							,  __snake_case						:	List[str]=0.1							,  __snake_case						:	int=32							,  __snake_case						:	Optional[Any]=2							,  __snake_case						:	int=1							,  __snake_case						:	str=0							,  __snake_case						:	Any=2							,  __snake_case						:	Optional[Any]=0.02							,  ):





											'''simple docstring'''





											_snake_case: Dict	       =		parent
											_snake_case: List[Any]	       =		batch_size
											_snake_case: Union[str, Any]	       =		seq_length
											_snake_case: List[str]	       =		is_training
											_snake_case: Optional[int]	       =		use_labels
											_snake_case: List[str]	       =		vocab_size
											_snake_case: Union[str, Any]	       =		hidden_size
											_snake_case: Optional[Any]	       =		num_hidden_layers
											_snake_case: int	       =		num_attention_heads
											_snake_case: Tuple	       =		intermediate_size
											_snake_case: Optional[int]	       =		hidden_act
											_snake_case: Optional[int]	       =		hidden_dropout_prob
											_snake_case: Union[str, Any]	       =		attention_probs_dropout_prob
											_snake_case: Tuple	       =		max_position_embeddings
											_snake_case: List[Any]	       =		eos_token_id
											_snake_case: int	       =		pad_token_id
											_snake_case: Optional[Any]	       =		bos_token_id
											_snake_case: Any	       =		decoder_start_token_id
											_snake_case: Dict	       =		initializer_range


					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							):





											'''simple docstring'''





											_snake_case: Optional[int]	       =		np.clip(ids_tensor([self.batch_size, self.seq_length - 1]							,  self.vocab_size							)							,  3							,  self.vocab_size							)
											_snake_case: Dict	       =		np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)							,  dtype=np.intaa							))							,  -1							)

											_snake_case: List[Any]	       =		shift_tokens_right(__snake_case							,  1							)

											_snake_case: Dict	       =		MBartConfig(
											    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							,  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							,  decoder_start_token_id=self.decoder_start_token_id							,  initializer_range=self.initializer_range							,  use_cache=__snake_case							,  )
											_snake_case: Optional[Any]	       =		prepare_mbart_inputs_dict(__snake_case							,  __snake_case							,  __snake_case							)
											return config, inputs_dict


					def  SCREAMING_SNAKE_CASE_		(       self						:	Any							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: int	       =		self.prepare_config_and_inputs()
											return config, inputs_dict


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[str]							,  __snake_case						:	str							,  __snake_case						:	List[Any]							,  __snake_case						:	Optional[int]							):





											'''simple docstring'''





											_snake_case: List[str]	       =		20
											_snake_case: Optional[Any]	       =		model_class_name(__snake_case							)

											_snake_case: List[Any]	       =		model.encode(inputs_dict['input_ids']							)

											_snake_case       ,    _snake_case: str	       =		(
											    inputs_dict['decoder_input_ids'],
											    inputs_dict['decoder_attention_mask'],
											)

											_snake_case: str	       =		model.init_cache(decoder_input_ids.shape[0]							,  __snake_case							,  __snake_case							)
											_snake_case: Tuple	       =		jnp.ones((decoder_input_ids.shape[0], max_decoder_length)							,  dtype='i4'							)

											_snake_case: Union[str, Any]	       =		jnp.broadcast_to(
											    jnp.arange(decoder_input_ids.shape[-1] - 1							)[None, :]							,  (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,  )
											_snake_case: Dict	       =		model.decode(
											    decoder_input_ids[:, :-1]							,  __snake_case							,  decoder_attention_mask=__snake_case							,  past_key_values=__snake_case							,  decoder_position_ids=__snake_case							,  )

											_snake_case: Optional[int]	       =		jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,  dtype='i4'							)
											_snake_case: List[str]	       =		model.decode(
											    decoder_input_ids[:, -1:]							,  __snake_case							,  decoder_attention_mask=__snake_case							,  past_key_values=outputs_cache.past_key_values							,  decoder_position_ids=__snake_case							,  )

											_snake_case: Any	       =		model.decode(__snake_case							,  __snake_case							)

											_snake_case: List[Any]	       =		np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])							)							)
											self.parent.assertTrue(diff < 1e-3							,  msg=f'''Max diff is {diff}'''							)




					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[int]							,  __snake_case						:	int							,  __snake_case						:	int							,  __snake_case						:	int							):





											'''simple docstring'''





											_snake_case: Dict	       =		20
											_snake_case: Any	       =		model_class_name(__snake_case							)

											_snake_case: Dict	       =		model.encode(inputs_dict['input_ids']							)

											_snake_case       ,    _snake_case: Tuple	       =		(
											    inputs_dict['decoder_input_ids'],
											    inputs_dict['decoder_attention_mask'],
											)

											_snake_case: Tuple	       =		jnp.concatenate(
											    [
											        decoder_attention_mask,
											        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])							),
											    ]							,  axis=-1							,  )

											_snake_case: str	       =		model.init_cache(decoder_input_ids.shape[0]							,  __snake_case							,  __snake_case							)
											_snake_case: Dict	       =		jnp.broadcast_to(
											    jnp.arange(decoder_input_ids.shape[-1] - 1							)[None, :]							,  (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,  )

											_snake_case: int	       =		model.decode(
											    decoder_input_ids[:, :-1]							,  __snake_case							,  decoder_attention_mask=__snake_case							,  past_key_values=__snake_case							,  decoder_position_ids=__snake_case							,  )
											_snake_case: Any	       =		jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,  dtype='i4'							)
											_snake_case: str	       =		model.decode(
											    decoder_input_ids[:, -1:]							,  __snake_case							,  past_key_values=outputs_cache.past_key_values							,  decoder_attention_mask=__snake_case							,  decoder_position_ids=__snake_case							,  )

											_snake_case: Any	       =		model.decode(__snake_case							,  __snake_case							,  decoder_attention_mask=__snake_case							)

											_snake_case: Optional[int]	       =		np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])							)							)
											self.parent.assertTrue(diff < 1e-3							,  msg=f'''Max diff is {diff}'''							)







@require_flax
class       lowerCamelCase       (    unittest.TestCase	):
					_SCREAMING_SNAKE_CASE				 =     99


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[str]							):





											'''simple docstring'''





											_snake_case: int	       =		np.array(
											    [
											        [71, 82, 18, 33, 46, 91, 2],
											        [68, 34, 26, 58, 30, 82, 2],
											        [5, 97, 17, 39, 94, 40, 2],
											        [76, 83, 94, 25, 70, 78, 2],
											        [87, 59, 41, 35, 48, 66, 2],
											        [55, 13, 16, 58, 5, 2, 1],  # note padding
											        [64, 27, 31, 51, 12, 75, 2],
											        [52, 64, 86, 17, 83, 39, 2],
											        [48, 61, 9, 24, 71, 82, 2],
											        [26, 1, 60, 48, 22, 13, 2],
											        [21, 5, 62, 28, 14, 76, 2],
											        [45, 98, 37, 86, 59, 48, 2],
											        [70, 70, 50, 9, 28, 0, 2],
											    ]							,  dtype=np.intaa							,  )

											_snake_case: Any	       =		input_ids.shape[0]
											_snake_case: str	       =		MBartConfig(
											    vocab_size=self.vocab_size							,  d_model=24							,  encoder_layers=2							,  decoder_layers=2							,  encoder_attention_heads=2							,  decoder_attention_heads=2							,  encoder_ffn_dim=32							,  decoder_ffn_dim=32							,  max_position_embeddings=48							,  eos_token_id=2							,  pad_token_id=1							,  bos_token_id=0							,  )
											return config, input_ids, batch_size


					def  SCREAMING_SNAKE_CASE_		(       self						:	Union[str, Any]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case       ,    _snake_case: Tuple	       =		self._get_config_and_data()
											_snake_case: Any	       =		FlaxMBartForSequenceClassification(__snake_case							)
											_snake_case: List[Any]	       =		model(input_ids=__snake_case							,  decoder_input_ids=__snake_case							)
											_snake_case: List[str]	       =		(batch_size, config.num_labels)
											self.assertEqual(outputs['logits'].shape							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[Any]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case       ,    _snake_case: Union[str, Any]	       =		self._get_config_and_data()
											_snake_case: List[str]	       =		FlaxMBartForQuestionAnswering(__snake_case							)
											_snake_case: Union[str, Any]	       =		model(input_ids=__snake_case							)

											self.assertEqual(outputs['start_logits'].shape							,  input_ids.shape							)
											self.assertEqual(outputs['end_logits'].shape							,  input_ids.shape							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''simple docstring'''





											_snake_case       ,    _snake_case       ,    _snake_case: Any	       =		self._get_config_and_data()
											_snake_case: Optional[int]	       =		FlaxMBartForConditionalGeneration(__snake_case							)
											_snake_case: List[str]	       =		lm_model(input_ids=__snake_case							)
											_snake_case: List[Any]	       =		(batch_size, input_ids.shape[1], config.vocab_size)
											self.assertEqual(outputs['logits'].shape							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[Any]							):





											'''simple docstring'''





											_snake_case: Optional[Any]	       =		MBartConfig(
											    vocab_size=self.vocab_size							,  d_model=14							,  encoder_layers=2							,  decoder_layers=2							,  encoder_attention_heads=2							,  decoder_attention_heads=2							,  encoder_ffn_dim=8							,  decoder_ffn_dim=8							,  max_position_embeddings=48							,  )
											_snake_case: Any	       =		FlaxMBartForConditionalGeneration(__snake_case							)
											_snake_case: Optional[Any]	       =		np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]							,  dtype=np.intaa							)
											_snake_case: Optional[int]	       =		np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]							,  dtype=np.intaa							)
											_snake_case: Optional[int]	       =		lm_model(input_ids=__snake_case							,  decoder_input_ids=__snake_case							)
											_snake_case: int	       =		(*summary.shape, config.vocab_size)
											self.assertEqual(outputs['logits'].shape							,  __snake_case							)




					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[int]							):





											'''simple docstring'''





											_snake_case: Any	       =		np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]							,  dtype=np.intaa							)
											_snake_case: Any	       =		shift_tokens_right(__snake_case							,  1							)
											_snake_case: List[str]	       =		np.equal(__snake_case							,  1							).astype(np.floataa							).sum()
											_snake_case: List[str]	       =		np.equal(__snake_case							,  1							).astype(np.floataa							).sum()
											self.assertEqual(shifted.shape							,  input_ids.shape							)
											self.assertEqual(__snake_case							,  n_pad_before - 1							)
											self.assertTrue(np.equal(shifted[:, 0]							,  2							).all()							)







@require_flax
class       lowerCamelCase       (    __UpperCAmelCase  ,				unittest.TestCase  ,				__UpperCAmelCase	):
					_SCREAMING_SNAKE_CASE				 =     True
					_SCREAMING_SNAKE_CASE				 =     (
					    (
					        FlaxMBartModel,
					        FlaxMBartForConditionalGeneration,
					        FlaxMBartForSequenceClassification,
					        FlaxMBartForQuestionAnswering,
					    )
					    if is_flax_available()
					    else ()
					)
					_SCREAMING_SNAKE_CASE				 =     (FlaxMBartForConditionalGeneration,) if is_flax_available() else ()


					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''simple docstring'''





											_snake_case: Union[str, Any]	       =		FlaxMBartModelTester(self							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	int							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: int	       =		self.model_tester.prepare_config_and_inputs()
											for model_class in self.all_model_classes:
																	self.model_tester.check_use_cache_forward(__snake_case							,  __snake_case							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Union[str, Any]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Tuple	       =		self.model_tester.prepare_config_and_inputs()
											for model_class in self.all_model_classes:
																	self.model_tester.check_use_cache_forward_with_attn_mask(__snake_case							,  __snake_case							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[str]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Any	       =		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	with self.subTest(model_class.__name__							):
																							_snake_case: Dict	       =		self._prepare_for_class(__snake_case							,  __snake_case							)
																							_snake_case: List[Any]	       =		model_class(__snake_case							)

																							@jax.jit
																							def encode_jitted(__snake_case						:	str							,  __snake_case						:	Tuple=None							,  **__snake_case						:	Optional[int]							):
																													return model.encode(input_ids=__snake_case							,  attention_mask=__snake_case							)

																							with self.subTest('JIT Enabled'							):
																													_snake_case: List[str]	       =		encode_jitted(**__snake_case							).to_tuple()

																							with self.subTest('JIT Disabled'							):
																													with jax.disable_jit():
																																			_snake_case: List[Any]	       =		encode_jitted(**__snake_case							).to_tuple()

																							self.assertEqual(len(__snake_case							)							,  len(__snake_case							)							)
																							for jitted_output, output in zip(__snake_case							,  __snake_case							):
																													self.assertEqual(jitted_output.shape							,  output.shape							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	str							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Optional[int]	       =		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	with self.subTest(model_class.__name__							):
																							_snake_case: Tuple	       =		model_class(__snake_case							)
																							_snake_case: Dict	       =		model.encode(inputs_dict['input_ids']							,  inputs_dict['attention_mask']							)

																							_snake_case: Dict	       =		{
																							    'decoder_input_ids': inputs_dict['decoder_input_ids'],
																							    'decoder_attention_mask': inputs_dict['decoder_attention_mask'],
																							    'encoder_outputs': encoder_outputs,
																							}

																							@jax.jit
																							def decode_jitted(__snake_case						:	Dict							,  __snake_case						:	Dict							,  __snake_case						:	Union[str, Any]							):
																													return model.decode(
																													    decoder_input_ids=__snake_case							,  decoder_attention_mask=__snake_case							,  encoder_outputs=__snake_case							,  )

																							with self.subTest('JIT Enabled'							):
																													_snake_case: str	       =		decode_jitted(**__snake_case							).to_tuple()

																							with self.subTest('JIT Disabled'							):
																													with jax.disable_jit():
																																			_snake_case: Dict	       =		decode_jitted(**__snake_case							).to_tuple()

																							self.assertEqual(len(__snake_case							)							,  len(__snake_case							)							)
																							for jitted_output, output in zip(__snake_case							,  __snake_case							):
																													self.assertEqual(jitted_output.shape							,  output.shape							)




					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	int							):





											'''simple docstring'''





											for model_class_name in self.all_model_classes:
																	_snake_case: List[str]	       =		model_class_name.from_pretrained('facebook/mbart-large-cc25'							,  from_pt=__snake_case							)
																	# FlaxMBartForSequenceClassification expects eos token in input_ids
																	_snake_case: List[str]	       =		np.ones((1, 1)							) * model.config.eos_token_id
																	_snake_case: Optional[Any]	       =		model(__snake_case							)
																	self.assertIsNotNone(__snake_case							)







@require_flax
@require_sentencepiece
@require_tokenizers
class       lowerCamelCase       (    unittest.TestCase	):
					_SCREAMING_SNAKE_CASE				 =     [
					    " UN Chief Says There Is No Military Solution in Syria",
					]
					_SCREAMING_SNAKE_CASE				 =     [
					    "Şeful ONU declară că nu există o soluţie militară în Siria",
					]
					_SCREAMING_SNAKE_CASE				 =     "facebook/mbart-large-en-ro"


					@cached_property
					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[int]							):





											'''simple docstring'''





											return AutoTokenizer.from_pretrained(self.model_name							)


					@cached_property
					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							):





											'''simple docstring'''





											_snake_case: Dict	       =		FlaxMBartForConditionalGeneration.from_pretrained(self.model_name							,  from_pt=__snake_case							)
											return model


					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[int]							,  **__snake_case						:	int							):





											'''simple docstring'''





											_snake_case: List[Any]	       =		self.translate_src_text(**__snake_case							)
											self.assertListEqual(self.expected_text							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	str							,  **__snake_case						:	Tuple							):





											'''simple docstring'''





											_snake_case: int	       =		self.tokenizer(self.src_text							,  **__snake_case							,  return_tensors='np'							)
											_snake_case: Optional[Any]	       =		self.model.generate(
											    model_inputs.input_ids							,  attention_mask=model_inputs.attention_mask							,  decoder_start_token_id=self.tokenizer.lang_code_to_id['ro_RO']							,  early_stopping=__snake_case							,  num_beams=2							,  ).sequences
											_snake_case: Dict	       =		self.tokenizer.batch_decode(__snake_case							,  skip_special_tokens=__snake_case							)
											return generated_words




					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[Any]							):





											'''simple docstring'''





											self._assert_generated_batch_equal_expected()





