


import unittest

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
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 (
       CpmAntConfig,
       CpmAntForCausalLM,
       CpmAntModel,
       CpmAntTokenizer,
   )


@require_torch
class     __lowerCAmelCase     :






  """simple docstring"""





  def __init__(     self       :       Dict    ,				_lowerCAmelCase       :       int    ,				_lowerCAmelCase       :       Tuple=2    ,				_lowerCAmelCase       :       str=8    ,				_lowerCAmelCase       :       int=True    ,				_lowerCAmelCase       :       int=False    ,				_lowerCAmelCase       :       List[str]=False    ,				_lowerCAmelCase       :       Tuple=False    ,				_lowerCAmelCase       :       List[str]=False    ,				_lowerCAmelCase       :       str=9_9    ,				_lowerCAmelCase       :       Any=3_2    ,				_lowerCAmelCase       :       List[str]=3    ,				_lowerCAmelCase       :       Any=4    ,				_lowerCAmelCase       :       str=3_7    ,				_lowerCAmelCase       :       List[str]=3_2    ,				_lowerCAmelCase       :       Dict=1_2_8    ,				_lowerCAmelCase       :       List[Any]=8    ,				_lowerCAmelCase       :       Optional[int]=8    ,				_lowerCAmelCase       :       int=8    ,				_lowerCAmelCase       :       Optional[int]=1.0    ,				_lowerCAmelCase       :       Tuple=True    ,				)   -> str:




     """simple docstring"""
     snake_case_					    =					parent
     snake_case_					    =					batch_size
     snake_case_					    =					seq_length
     snake_case_					    =					is_training
     snake_case_					    =					use_token_type_ids
     snake_case_					    =					use_input_mask
     snake_case_					    =					use_labels
     snake_case_					    =					use_mc_token_ids
     snake_case_					    =					vocab_size
     snake_case_					    =					hidden_size
     snake_case_					    =					num_hidden_layers
     snake_case_					    =					num_attention_heads
     snake_case_					    =					intermediate_size
     snake_case_					    =					num_buckets
     snake_case_					    =					max_distance
     snake_case_					    =					prompt_length
     snake_case_					    =					prompt_types
     snake_case_					    =					segment_types
     snake_case_					    =					init_std
     snake_case_					    =					return_dict



  def 					lowerCAmelCase__  (     self       :       int  )   -> Optional[Any]:




     """simple docstring"""
     snake_case_					    =					{}
     snake_case_					    =					ids_tensor([self.batch_size, self.seq_length]    ,				self.vocab_size  ).type(torch.intaa  )
     snake_case_					    =					False

     snake_case_					    =					self.get_config()

     return (config, input_ids)



  def 					lowerCAmelCase__  (     self       :       Optional[Any]  )   -> int:




     """simple docstring"""
     return CpmAntConfig(
         vocab_size=self.vocab_size    ,				hidden_size=self.hidden_size    ,				num_hidden_layers=self.num_hidden_layers    ,				num_attention_heads=self.num_attention_heads    ,				dim_ff=self.intermediate_size    ,				position_bias_num_buckets=self.num_buckets    ,				position_bias_max_distance=self.max_distance    ,				prompt_types=self.prompt_types    ,				prompt_length=self.prompt_length    ,				segment_types=self.segment_types    ,				use_cache=_lowerCAmelCase    ,				init_std=self.init_std    ,				return_dict=self.return_dict    ,				)



  def 					lowerCAmelCase__  (     self       :       List[str]    ,				_lowerCAmelCase       :       str    ,				_lowerCAmelCase       :       Union[str, Any]    ,				*_lowerCAmelCase       :       Dict  )   -> List[Any]:




     """simple docstring"""
     snake_case_					    =					CpmAntModel(config=_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     model.eval()

     snake_case_					    =					model(**_lowerCAmelCase  ).last_hidden_state

     self.parent.assertEqual(hidden_states.shape    ,				(self.batch_size, self.seq_length, config.hidden_size)  )



  def 					lowerCAmelCase__  (     self       :       Union[str, Any]    ,				_lowerCAmelCase       :       List[str]    ,				_lowerCAmelCase       :       Any    ,				*_lowerCAmelCase       :       Union[str, Any]  )   -> Union[str, Any]:




     """simple docstring"""
     snake_case_					    =					CpmAntForCausalLM(_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     snake_case_					    =					input_ids["input_ids"].to(_lowerCAmelCase  )
     model.eval()

     snake_case_					    =					model(**_lowerCAmelCase  )
     self.parent.assertEqual(
         model_output.logits.shape    ,				(self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)    ,				)



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




     """simple docstring"""
     snake_case_    ,    snake_case_					    =					self.prepare_config_and_inputs()
     return config, inputs_dict






@require_torch
class     __lowerCAmelCase     ( a ,  a ,  unittest.TestCase					):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		    =       (CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
  _SCREAMING_SNAKE_CASE		    =       (
      {'feature-extraction': CpmAntModel, 'text-generation': CpmAntForCausalLM} if is_torch_available() else {}
  )

  _SCREAMING_SNAKE_CASE		    =       False
  _SCREAMING_SNAKE_CASE		    =       False
  _SCREAMING_SNAKE_CASE		    =       False
  _SCREAMING_SNAKE_CASE		    =       False
  _SCREAMING_SNAKE_CASE		    =       False



  def 					lowerCAmelCase__  (     self       :       Optional[int]  )   -> List[Any]:




     """simple docstring"""
     snake_case_					    =					CpmAntModelTester(self  )
     snake_case_					    =					ConfigTester(self    ,				config_class=_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       Optional[Any]  )   -> List[Any]:




     """simple docstring"""
     self.config_tester.create_and_test_config_common_properties()
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()



  def 					lowerCAmelCase__  (     self       :       str  )   -> Optional[int]:




     """simple docstring"""
     unittest.skip("CPMAnt doesn't support input_embeds."  )(self.test_inputs_embeds  )



  def 					lowerCAmelCase__  (     self       :       Optional[Any]  )   -> Optional[int]:




     """simple docstring"""
     unittest.skip(
         "CPMAnt doesn't support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model."  )(self.test_retain_grad_hidden_states_attentions  )



  def 					lowerCAmelCase__  (     self       :       Optional[int]  )   -> int:




     """simple docstring"""
     snake_case_    ,    snake_case_					    =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_cpmant_model(_lowerCAmelCase    ,				_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       Dict  )   -> Dict:




     """simple docstring"""
     snake_case_    ,    snake_case_					    =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_lm_head_model(_lowerCAmelCase    ,				_lowerCAmelCase  )






@require_torch
class     __lowerCAmelCase     ( unittest.TestCase					):






  """simple docstring"""





  @tooslow
  def 					lowerCAmelCase__  (     self       :       Tuple  )   -> Any:




     """simple docstring"""
     snake_case_					    =					"今天天气真好！"
     snake_case_					    =					"openbmb/cpm-ant-10b"
     snake_case_					    =					CpmAntModel.from_pretrained(_lowerCAmelCase  )
     snake_case_					    =					CpmAntTokenizer.from_pretrained(_lowerCAmelCase  )
     snake_case_					    =					tokenizer(_lowerCAmelCase    ,				return_tensors="pt"  )
     snake_case_					    =					model(**_lowerCAmelCase  ).last_hidden_state

     snake_case_					    =					torch.tensor(
         [[[6.1_708, 5.9_244, 1.0_835], [6.5_207, 6.2_893, -11.3_324], [-1.0_107, -0.0_576, -5.9_577]]]    ,				)
     self.assertTrue(torch.allclose(hidden_states[:, :3, :3]    ,				_lowerCAmelCase    ,				atol=1e-2  )  )






@require_torch
class     __lowerCAmelCase     ( unittest.TestCase					):






  """simple docstring"""





  @tooslow
  def 					lowerCAmelCase__  (     self       :       Optional[int]  )   -> Any:




     """simple docstring"""
     snake_case_					    =					"今天天气真好！"
     snake_case_					    =					"openbmb/cpm-ant-10b"
     snake_case_					    =					CpmAntForCausalLM.from_pretrained(_lowerCAmelCase  )
     snake_case_					    =					CpmAntTokenizer.from_pretrained(_lowerCAmelCase  )
     snake_case_					    =					tokenizer(_lowerCAmelCase    ,				return_tensors="pt"  )
     snake_case_					    =					model(**_lowerCAmelCase  ).logits

     snake_case_					    =					torch.tensor(
         [[[-6.4_267, -6.4_083, -6.3_958], [-5.8_802, -5.9_447, -5.7_811], [-5.3_896, -5.4_820, -5.4_295]]]    ,				)
     self.assertTrue(torch.allclose(hidden_states[:, :3, :3]    ,				_lowerCAmelCase    ,				atol=1e-2  )  )



  @tooslow
  def 					lowerCAmelCase__  (     self       :       Union[str, Any]  )   -> Optional[int]:




     """simple docstring"""
     snake_case_					    =					"openbmb/cpm-ant-10b"
     snake_case_					    =					CpmAntForCausalLM.from_pretrained(_lowerCAmelCase  )
     snake_case_					    =					CpmAntTokenizer.from_pretrained(_lowerCAmelCase  )
     snake_case_					    =					"今天天气不错，"
     snake_case_					    =					"今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的"
     snake_case_					    =					tokenizer(_lowerCAmelCase    ,				return_tensors="pt"  )
     snake_case_					    =					model.generate(**_lowerCAmelCase  )
     snake_case_					    =					tokenizer.batch_decode(_lowerCAmelCase  )
     self.assertEqual(_lowerCAmelCase    ,				_lowerCAmelCase  )



  @tooslow
  def 					lowerCAmelCase__  (     self       :       Union[str, Any]  )   -> int:




     """simple docstring"""
     snake_case_					    =					"openbmb/cpm-ant-10b"
     snake_case_					    =					CpmAntForCausalLM.from_pretrained(_lowerCAmelCase  )
     snake_case_					    =					CpmAntTokenizer.from_pretrained(_lowerCAmelCase  )
     snake_case_					    =					["今天天气不错，", "新年快乐，万事如意！"]
     snake_case_					    =					[
         "今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的",
         "新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的",
     ]
     snake_case_					    =					tokenizer(_lowerCAmelCase    ,				return_tensors="pt"    ,				padding=_lowerCAmelCase  )
     snake_case_					    =					model.generate(**_lowerCAmelCase  )
     snake_case_					    =					tokenizer.batch_decode(_lowerCAmelCase  )
     self.assertEqual(_lowerCAmelCase    ,				_lowerCAmelCase  )








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
   SCREAMING_SNAKE_CASE	: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 			_lowerCAmelCase						(							lowerCAmelCase_							:Union[str, Any]		,						lowerCAmelCase_							:str		,						lowerCAmelCase_							:Any=None		,						lowerCAmelCase_							:Optional[Any]=None		,						lowerCAmelCase_							:Any=None		,						lowerCAmelCase_							:List[Any]=None		,						lowerCAmelCase_							:Optional[int]=None		,						lowerCAmelCase_							:Union[str, Any]=None		,						)->Optional[Any]:
   '''simple docstring'''






   if attention_mask is None:
      snake_case_					    =					np.where(input_ids != config.pad_token_id		,						1		,						0							)
   if decoder_attention_mask is None:
      snake_case_					    =					np.where(decoder_input_ids != config.pad_token_id		,						1		,						0							)
   if head_mask is None:
      snake_case_					    =					np.ones((config.encoder_layers, config.encoder_attention_heads)							)
   if decoder_head_mask is None:
      snake_case_					    =					np.ones((config.decoder_layers, config.decoder_attention_heads)							)
   if cross_attn_head_mask is None:
      snake_case_					    =					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					):






  """simple docstring"""





  def __init__(     self       :       List[Any]    ,				_lowerCAmelCase       :       str    ,				_lowerCAmelCase       :       Optional[int]=1_3    ,				_lowerCAmelCase       :       List[str]=7    ,				_lowerCAmelCase       :       Dict=True    ,				_lowerCAmelCase       :       Optional[Any]=False    ,				_lowerCAmelCase       :       Dict=9_9    ,				_lowerCAmelCase       :       int=1_6    ,				_lowerCAmelCase       :       Union[str, Any]=2    ,				_lowerCAmelCase       :       Tuple=4    ,				_lowerCAmelCase       :       List[Any]=4    ,				_lowerCAmelCase       :       Any="gelu"    ,				_lowerCAmelCase       :       int=0.1    ,				_lowerCAmelCase       :       Union[str, Any]=0.1    ,				_lowerCAmelCase       :       Dict=3_2    ,				_lowerCAmelCase       :       Optional[int]=2    ,				_lowerCAmelCase       :       Any=1    ,				_lowerCAmelCase       :       Union[str, Any]=0    ,				_lowerCAmelCase       :       Any=2    ,				_lowerCAmelCase       :       Any=0.02    ,				)   -> Optional[int]:




     """simple docstring"""
     snake_case_					    =					parent
     snake_case_					    =					batch_size
     snake_case_					    =					seq_length
     snake_case_					    =					is_training
     snake_case_					    =					use_labels
     snake_case_					    =					vocab_size
     snake_case_					    =					hidden_size
     snake_case_					    =					num_hidden_layers
     snake_case_					    =					num_attention_heads
     snake_case_					    =					intermediate_size
     snake_case_					    =					hidden_act
     snake_case_					    =					hidden_dropout_prob
     snake_case_					    =					attention_probs_dropout_prob
     snake_case_					    =					max_position_embeddings
     snake_case_					    =					eos_token_id
     snake_case_					    =					pad_token_id
     snake_case_					    =					bos_token_id
     snake_case_					    =					decoder_start_token_id
     snake_case_					    =					initializer_range



  def 					lowerCAmelCase__  (     self       :       Union[str, Any]  )   -> List[Any]:




     """simple docstring"""
     snake_case_					    =					np.clip(ids_tensor([self.batch_size, self.seq_length - 1]    ,				self.vocab_size  )    ,				3    ,				self.vocab_size  )
     snake_case_					    =					np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)    ,				dtype=np.intaa  ))    ,				-1  )

     snake_case_					    =					shift_tokens_right(_lowerCAmelCase    ,				1  )

     snake_case_					    =					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=_lowerCAmelCase    ,				)
     snake_case_					    =					prepare_mbart_inputs_dict(_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase  )
     return config, inputs_dict



  def 					lowerCAmelCase__  (     self       :       int  )   -> Union[str, Any]:




     """simple docstring"""
     snake_case_    ,    snake_case_					    =					self.prepare_config_and_inputs()
     return config, inputs_dict



  def 					lowerCAmelCase__  (     self       :       str    ,				_lowerCAmelCase       :       List[str]    ,				_lowerCAmelCase       :       Union[str, Any]    ,				_lowerCAmelCase       :       int  )   -> Optional[Any]:




     """simple docstring"""
     snake_case_					    =					2_0
     snake_case_					    =					model_class_name(_lowerCAmelCase  )

     snake_case_					    =					model.encode(inputs_dict["input_ids"]  )

     snake_case_    ,    snake_case_					    =					(
         inputs_dict["decoder_input_ids"],
         inputs_dict["decoder_attention_mask"],
     )

     snake_case_					    =					model.init_cache(decoder_input_ids.shape[0]    ,				_lowerCAmelCase    ,				_lowerCAmelCase  )
     snake_case_					    =					jnp.ones((decoder_input_ids.shape[0], max_decoder_length)    ,				dtype="i4"  )

     snake_case_					    =					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_					    =					model.decode(
         decoder_input_ids[:, :-1]    ,				_lowerCAmelCase    ,				decoder_attention_mask=_lowerCAmelCase    ,				past_key_values=_lowerCAmelCase    ,				decoder_position_ids=_lowerCAmelCase    ,				)

     snake_case_					    =					jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]    ,				dtype="i4"  )
     snake_case_					    =					model.decode(
         decoder_input_ids[:, -1:]    ,				_lowerCAmelCase    ,				decoder_attention_mask=_lowerCAmelCase    ,				past_key_values=outputs_cache.past_key_values    ,				decoder_position_ids=_lowerCAmelCase    ,				)

     snake_case_					    =					model.decode(_lowerCAmelCase    ,				_lowerCAmelCase  )

     snake_case_					    =					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 					lowerCAmelCase__  (     self       :       Union[str, Any]    ,				_lowerCAmelCase       :       int    ,				_lowerCAmelCase       :       Any    ,				_lowerCAmelCase       :       Dict  )   -> str:




     """simple docstring"""
     snake_case_					    =					2_0
     snake_case_					    =					model_class_name(_lowerCAmelCase  )

     snake_case_					    =					model.encode(inputs_dict["input_ids"]  )

     snake_case_    ,    snake_case_					    =					(
         inputs_dict["decoder_input_ids"],
         inputs_dict["decoder_attention_mask"],
     )

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

     snake_case_					    =					model.init_cache(decoder_input_ids.shape[0]    ,				_lowerCAmelCase    ,				_lowerCAmelCase  )
     snake_case_					    =					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_					    =					model.decode(
         decoder_input_ids[:, :-1]    ,				_lowerCAmelCase    ,				decoder_attention_mask=_lowerCAmelCase    ,				past_key_values=_lowerCAmelCase    ,				decoder_position_ids=_lowerCAmelCase    ,				)
     snake_case_					    =					jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]    ,				dtype="i4"  )
     snake_case_					    =					model.decode(
         decoder_input_ids[:, -1:]    ,				_lowerCAmelCase    ,				past_key_values=outputs_cache.past_key_values    ,				decoder_attention_mask=_lowerCAmelCase    ,				decoder_position_ids=_lowerCAmelCase    ,				)

     snake_case_					    =					model.decode(_lowerCAmelCase    ,				_lowerCAmelCase    ,				decoder_attention_mask=_lowerCAmelCase  )

     snake_case_					    =					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					):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		    =       99



  def 					lowerCAmelCase__  (     self       :       Tuple  )   -> Tuple:




     """simple docstring"""
     snake_case_					    =					np.array(
         [
             [7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2],
             [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2],
             [5, 9_7, 1_7, 3_9, 9_4, 4_0, 2],
             [7_6, 8_3, 9_4, 2_5, 7_0, 7_8, 2],
             [8_7, 5_9, 4_1, 3_5, 4_8, 6_6, 2],
             [5_5, 1_3, 1_6, 5_8, 5, 2, 1],  # note padding
             [6_4, 2_7, 3_1, 5_1, 1_2, 7_5, 2],
             [5_2, 6_4, 8_6, 1_7, 8_3, 3_9, 2],
             [4_8, 6_1, 9, 2_4, 7_1, 8_2, 2],
             [2_6, 1, 6_0, 4_8, 2_2, 1_3, 2],
             [2_1, 5, 6_2, 2_8, 1_4, 7_6, 2],
             [4_5, 9_8, 3_7, 8_6, 5_9, 4_8, 2],
             [7_0, 7_0, 5_0, 9, 2_8, 0, 2],
         ]    ,				dtype=np.intaa    ,				)

     snake_case_					    =					input_ids.shape[0]
     snake_case_					    =					MBartConfig(
         vocab_size=self.vocab_size    ,				d_model=2_4    ,				encoder_layers=2    ,				decoder_layers=2    ,				encoder_attention_heads=2    ,				decoder_attention_heads=2    ,				encoder_ffn_dim=3_2    ,				decoder_ffn_dim=3_2    ,				max_position_embeddings=4_8    ,				eos_token_id=2    ,				pad_token_id=1    ,				bos_token_id=0    ,				)
     return config, input_ids, batch_size



  def 					lowerCAmelCase__  (     self       :       str  )   -> Dict:




     """simple docstring"""
     snake_case_    ,    snake_case_    ,    snake_case_					    =					self._get_config_and_data()
     snake_case_					    =					FlaxMBartForSequenceClassification(_lowerCAmelCase  )
     snake_case_					    =					model(input_ids=_lowerCAmelCase    ,				decoder_input_ids=_lowerCAmelCase  )
     snake_case_					    =					(batch_size, config.num_labels)
     self.assertEqual(outputs["logits"].shape    ,				_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       int  )   -> Optional[Any]:




     """simple docstring"""
     snake_case_    ,    snake_case_    ,    snake_case_					    =					self._get_config_and_data()
     snake_case_					    =					FlaxMBartForQuestionAnswering(_lowerCAmelCase  )
     snake_case_					    =					model(input_ids=_lowerCAmelCase  )

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



  def 					lowerCAmelCase__  (     self       :       Optional[int]  )   -> Dict:




     """simple docstring"""
     snake_case_    ,    snake_case_    ,    snake_case_					    =					self._get_config_and_data()
     snake_case_					    =					FlaxMBartForConditionalGeneration(_lowerCAmelCase  )
     snake_case_					    =					lm_model(input_ids=_lowerCAmelCase  )
     snake_case_					    =					(batch_size, input_ids.shape[1], config.vocab_size)
     self.assertEqual(outputs["logits"].shape    ,				_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       Any  )   -> Tuple:




     """simple docstring"""
     snake_case_					    =					MBartConfig(
         vocab_size=self.vocab_size    ,				d_model=1_4    ,				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=4_8    ,				)
     snake_case_					    =					FlaxMBartForConditionalGeneration(_lowerCAmelCase  )
     snake_case_					    =					np.array([[7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2], [6_8, 3_4, 2_6, 5_8, 3_0, 2, 1]]    ,				dtype=np.intaa  )
     snake_case_					    =					np.array([[8_2, 7_1, 8_2, 1_8, 2], [5_8, 6_8, 2, 1, 1]]    ,				dtype=np.intaa  )
     snake_case_					    =					lm_model(input_ids=_lowerCAmelCase    ,				decoder_input_ids=_lowerCAmelCase  )
     snake_case_					    =					(*summary.shape, config.vocab_size)
     self.assertEqual(outputs["logits"].shape    ,				_lowerCAmelCase  )



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




     """simple docstring"""
     snake_case_					    =					np.array([[7_1, 8_2, 1_8, 3_3, 2, 1, 1], [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2]]    ,				dtype=np.intaa  )
     snake_case_					    =					shift_tokens_right(_lowerCAmelCase    ,				1  )
     snake_case_					    =					np.equal(_lowerCAmelCase    ,				1  ).astype(np.floataa  ).sum()
     snake_case_					    =					np.equal(_lowerCAmelCase    ,				1  ).astype(np.floataa  ).sum()
     self.assertEqual(shifted.shape    ,				input_ids.shape  )
     self.assertEqual(_lowerCAmelCase    ,				n_pad_before - 1  )
     self.assertTrue(np.equal(shifted[:, 0]    ,				2  ).all()  )






@require_flax
class     __lowerCAmelCase     ( a ,  unittest.TestCase ,  a					):






  """simple docstring"""





  _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 					lowerCAmelCase__  (     self       :       Tuple  )   -> Dict:




     """simple docstring"""
     snake_case_					    =					FlaxMBartModelTester(self  )



  def 					lowerCAmelCase__  (     self       :       List[str]  )   -> List[Any]:




     """simple docstring"""
     snake_case_    ,    snake_case_					    =					self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
        self.model_tester.check_use_cache_forward(_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       List[str]  )   -> Tuple:




     """simple docstring"""
     snake_case_    ,    snake_case_					    =					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(_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase  )



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




     """simple docstring"""
     snake_case_    ,    snake_case_					    =					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_					    =					self._prepare_for_class(_lowerCAmelCase    ,				_lowerCAmelCase  )
           snake_case_					    =					model_class(_lowerCAmelCase  )

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

           with self.subTest("JIT Enabled"  ):
              snake_case_					    =					encode_jitted(**_lowerCAmelCase  ).to_tuple()

           with self.subTest("JIT Disabled"  ):
              with jax.disable_jit():
                 snake_case_					    =					encode_jitted(**_lowerCAmelCase  ).to_tuple()

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



  def 					lowerCAmelCase__  (     self       :       Optional[Any]  )   -> List[Any]:




     """simple docstring"""
     snake_case_    ,    snake_case_					    =					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_					    =					model_class(_lowerCAmelCase  )
           snake_case_					    =					model.encode(inputs_dict["input_ids"]    ,				inputs_dict["attention_mask"]  )

           snake_case_					    =					{
               "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(_lowerCAmelCase       :       List[str]    ,				_lowerCAmelCase       :       List[str]    ,				_lowerCAmelCase       :       str  ):
              return model.decode(
                  decoder_input_ids=_lowerCAmelCase    ,				decoder_attention_mask=_lowerCAmelCase    ,				encoder_outputs=_lowerCAmelCase    ,				)

           with self.subTest("JIT Enabled"  ):
              snake_case_					    =					decode_jitted(**_lowerCAmelCase  ).to_tuple()

           with self.subTest("JIT Disabled"  ):
              with jax.disable_jit():
                 snake_case_					    =					decode_jitted(**_lowerCAmelCase  ).to_tuple()

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



  @slow
  def 					lowerCAmelCase__  (     self       :       Tuple  )   -> Optional[Any]:




     """simple docstring"""
     for model_class_name in self.all_model_classes:
        snake_case_					    =					model_class_name.from_pretrained("facebook/mbart-large-cc25"    ,				from_pt=_lowerCAmelCase  )
        # FlaxMBartForSequenceClassification expects eos token in input_ids
        snake_case_					    =					np.ones((1, 1)  ) * model.config.eos_token_id
        snake_case_					    =					model(_lowerCAmelCase  )
        self.assertIsNotNone(_lowerCAmelCase  )






@require_flax
@require_sentencepiece
@require_tokenizers
class     __lowerCAmelCase     ( unittest.TestCase					):






  """simple docstring"""





  _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 					lowerCAmelCase__  (     self       :       Union[str, Any]  )   -> Any:




     """simple docstring"""
     return AutoTokenizer.from_pretrained(self.model_name  )



  @cached_property
  def 					lowerCAmelCase__  (     self       :       List[str]  )   -> Optional[int]:




     """simple docstring"""
     snake_case_					    =					FlaxMBartForConditionalGeneration.from_pretrained(self.model_name    ,				from_pt=_lowerCAmelCase  )
     return model



  def 					lowerCAmelCase__  (     self       :       Optional[Any]    ,				**_lowerCAmelCase       :       Tuple  )   -> str:




     """simple docstring"""
     snake_case_					    =					self.translate_src_text(**_lowerCAmelCase  )
     self.assertListEqual(self.expected_text    ,				_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       Optional[Any]    ,				**_lowerCAmelCase       :       Optional[int]  )   -> List[str]:




     """simple docstring"""
     snake_case_					    =					self.tokenizer(self.src_text    ,				**_lowerCAmelCase    ,				return_tensors="np"  )
     snake_case_					    =					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=_lowerCAmelCase    ,				num_beams=2    ,				).sequences
     snake_case_					    =					self.tokenizer.batch_decode(_lowerCAmelCase    ,				skip_special_tokens=_lowerCAmelCase  )
     return generated_words



  @slow
  def 					lowerCAmelCase__  (     self       :       Any  )   -> Dict:




     """simple docstring"""
     self._assert_generated_batch_equal_expected()




