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
      UpperCAmelCase_    		=				'''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 				UpperCamelCase					(     lowerCAmelCase_							,       lowerCAmelCase_							,       lowerCAmelCase_=None							,       lowerCAmelCase_=None							,       lowerCAmelCase_=None							,       lowerCAmelCase_=None							,       lowerCAmelCase_=None							,       lowerCAmelCase_=None							,       )		->    int:




       '''simple docstring'''






       if attention_mask is None:
              _A=						np.where(input_ids != config.pad_token_id							,       1							,       0 )
       if decoder_attention_mask is None:
              _A=						np.where(decoder_input_ids != config.pad_token_id							,       1							,       0 )
       if head_mask is None:
              _A=						np.ones((config.encoder_layers, config.encoder_attention_heads) )
       if decoder_head_mask is None:
              _A=						np.ones((config.decoder_layers, config.decoder_attention_heads) )
       if cross_attn_head_mask is None:
              _A=						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					,   lowerCAmelCase__					,   lowerCAmelCase__=13					,   lowerCAmelCase__=7					,   lowerCAmelCase__=True					,   lowerCAmelCase__=False					,   lowerCAmelCase__=99					,   lowerCAmelCase__=16					,   lowerCAmelCase__=2					,   lowerCAmelCase__=4					,   lowerCAmelCase__=4					,   lowerCAmelCase__="gelu"					,   lowerCAmelCase__=0.1					,   lowerCAmelCase__=0.1					,   lowerCAmelCase__=32					,   lowerCAmelCase__=2					,   lowerCAmelCase__=1					,   lowerCAmelCase__=0					,   lowerCAmelCase__=2					,   lowerCAmelCase__=0.02					,   ):
            _A=						parent
            _A=						batch_size
            _A=						seq_length
            _A=						is_training
            _A=						use_labels
            _A=						vocab_size
            _A=						hidden_size
            _A=						num_hidden_layers
            _A=						num_attention_heads
            _A=						intermediate_size
            _A=						hidden_act
            _A=						hidden_dropout_prob
            _A=						attention_probs_dropout_prob
            _A=						max_position_embeddings
            _A=						eos_token_id
            _A=						pad_token_id
            _A=						bos_token_id
            _A=						decoder_start_token_id
            _A=						initializer_range



     def       a__	(       self							):
            _A=						np.clip(ids_tensor([self.batch_size, self.seq_length - 1]					,   self.vocab_size							)					,   3					,   self.vocab_size							)
            _A=						np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)					,   dtype=np.intaa							))					,   -1							)

            _A=						shift_tokens_right(lowerCAmelCase__					,   1							)

            _A=						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__					,   )
            _A=						prepare_mbart_inputs_dict(lowerCAmelCase__					,   lowerCAmelCase__					,   lowerCAmelCase__							)
            return config, inputs_dict



     def       a__	(       self							):
            _A,      _A=						self.prepare_config_and_inputs()
            return config, inputs_dict



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            _A=						20
            _A=						model_class_name(lowerCAmelCase__							)

            _A=						model.encode(inputs_dict['input_ids']							)

            _A,      _A=						(
                inputs_dict['decoder_input_ids'],
                inputs_dict['decoder_attention_mask'],
            )

            _A=						model.init_cache(decoder_input_ids.shape[0]					,   lowerCAmelCase__					,   lowerCAmelCase__							)
            _A=						jnp.ones((decoder_input_ids.shape[0], max_decoder_length)					,   dtype='i4'							)

            _A=						jnp.broadcast_to(
                jnp.arange(decoder_input_ids.shape[-1] - 1							)[None, :]					,   (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)					,   )
            _A=						model.decode(
                decoder_input_ids[:, :-1]					,   lowerCAmelCase__					,   decoder_attention_mask=lowerCAmelCase__					,   past_key_values=lowerCAmelCase__					,   decoder_position_ids=lowerCAmelCase__					,   )

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

            _A=						model.decode(lowerCAmelCase__					,   lowerCAmelCase__							)

            _A=						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       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            _A=						20
            _A=						model_class_name(lowerCAmelCase__							)

            _A=						model.encode(inputs_dict['input_ids']							)

            _A,      _A=						(
                inputs_dict['decoder_input_ids'],
                inputs_dict['decoder_attention_mask'],
            )

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

            _A=						model.init_cache(decoder_input_ids.shape[0]					,   lowerCAmelCase__					,   lowerCAmelCase__							)
            _A=						jnp.broadcast_to(
                jnp.arange(decoder_input_ids.shape[-1] - 1							)[None, :]					,   (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)					,   )

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

            _A=						model.decode(lowerCAmelCase__					,   lowerCAmelCase__					,   decoder_attention_mask=lowerCAmelCase__							)

            _A=						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  :  Optional[Any]		=99



     def       a__	(       self							):
            _A=						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					,   )

            _A=						input_ids.shape[0]
            _A=						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       a__	(       self							):
            _A,      _A,      _A=						self._get_config_and_data()
            _A=						FlaxMBartForSequenceClassification(lowerCAmelCase__							)
            _A=						model(input_ids=lowerCAmelCase__					,   decoder_input_ids=lowerCAmelCase__							)
            _A=						(batch_size, config.num_labels)
            self.assertEqual(outputs['logits'].shape					,   lowerCAmelCase__							)



     def       a__	(       self							):
            _A,      _A,      _A=						self._get_config_and_data()
            _A=						FlaxMBartForQuestionAnswering(lowerCAmelCase__							)
            _A=						model(input_ids=lowerCAmelCase__							)

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



     def       a__	(       self							):
            _A,      _A,      _A=						self._get_config_and_data()
            _A=						FlaxMBartForConditionalGeneration(lowerCAmelCase__							)
            _A=						lm_model(input_ids=lowerCAmelCase__							)
            _A=						(batch_size, input_ids.shape[1], config.vocab_size)
            self.assertEqual(outputs['logits'].shape					,   lowerCAmelCase__							)



     def       a__	(       self							):
            _A=						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					,   )
            _A=						FlaxMBartForConditionalGeneration(lowerCAmelCase__							)
            _A=						np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]					,   dtype=np.intaa							)
            _A=						np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]					,   dtype=np.intaa							)
            _A=						lm_model(input_ids=lowerCAmelCase__					,   decoder_input_ids=lowerCAmelCase__							)
            _A=						(*summary.shape, config.vocab_size)
            self.assertEqual(outputs['logits'].shape					,   lowerCAmelCase__							)





     def       a__	(       self							):
            _A=						np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]					,   dtype=np.intaa							)
            _A=						shift_tokens_right(lowerCAmelCase__					,   1							)
            _A=						np.equal(lowerCAmelCase__					,   1							).astype(np.floataa							).sum()
            _A=						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      ):
     _SCREAMING_SNAKE_CASE  :  Optional[int]		=True
     _SCREAMING_SNAKE_CASE  :  str		=(
         (
             FlaxMBartModel,
             FlaxMBartForConditionalGeneration,
             FlaxMBartForSequenceClassification,
             FlaxMBartForQuestionAnswering,
         )
         if is_flax_available()
         else ()
     )
     _SCREAMING_SNAKE_CASE  :  Optional[int]		=(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()



     def       a__	(       self							):
            _A=						FlaxMBartModelTester(self							)



     def       a__	(       self							):
            _A,      _A=						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       a__	(       self							):
            _A,      _A=						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       a__	(       self							):
            _A,      _A=						self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                   with self.subTest(model_class.__name__							):
                          _A=						self._prepare_for_class(lowerCAmelCase__					,   lowerCAmelCase__							)
                          _A=						model_class(lowerCAmelCase__							)

                          @jax.jit
                          def encode_jitted(lowerCAmelCase__					,   lowerCAmelCase__=None					,   **lowerCAmelCase__							):
                                 return model.encode(input_ids=lowerCAmelCase__					,   attention_mask=lowerCAmelCase__							)

                          with self.subTest('JIT Enabled'							):
                                 _A=						encode_jitted(**lowerCAmelCase__							).to_tuple()

                          with self.subTest('JIT Disabled'							):
                                 with jax.disable_jit():
                                        _A=						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       a__	(       self							):
            _A,      _A=						self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                   with self.subTest(model_class.__name__							):
                          _A=						model_class(lowerCAmelCase__							)
                          _A=						model.encode(inputs_dict['input_ids']					,   inputs_dict['attention_mask']							)

                          _A=						{
                              '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__					,   lowerCAmelCase__					,   lowerCAmelCase__							):
                                 return model.decode(
                                     decoder_input_ids=lowerCAmelCase__					,   decoder_attention_mask=lowerCAmelCase__					,   encoder_outputs=lowerCAmelCase__					,   )

                          with self.subTest('JIT Enabled'							):
                                 _A=						decode_jitted(**lowerCAmelCase__							).to_tuple()

                          with self.subTest('JIT Disabled'							):
                                 with jax.disable_jit():
                                        _A=						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       a__	(       self							):
            for model_class_name in self.all_model_classes:
                   _A=						model_class_name.from_pretrained('facebook/mbart-large-cc25'					,   from_pt=lowerCAmelCase__							)
                   # FlaxMBartForSequenceClassification expects eos token in input_ids
                   _A=						np.ones((1, 1)							) * model.config.eos_token_id
                   _A=						model(lowerCAmelCase__							)
                   self.assertIsNotNone(lowerCAmelCase__							)





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



     @cached_property
     def       a__	(       self							):
            return AutoTokenizer.from_pretrained(self.model_name							)



     @cached_property
     def       a__	(       self							):
            _A=						FlaxMBartForConditionalGeneration.from_pretrained(self.model_name					,   from_pt=lowerCAmelCase__							)
            return model



     def       a__	(       self					,   **lowerCAmelCase__							):
            _A=						self.translate_src_text(**lowerCAmelCase__							)
            self.assertListEqual(self.expected_text					,   lowerCAmelCase__							)



     def       a__	(       self					,   **lowerCAmelCase__							):
            _A=						self.tokenizer(self.src_text					,   **lowerCAmelCase__					,   return_tensors='np'							)
            _A=						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
            _A=						self.tokenizer.batch_decode(lowerCAmelCase__					,   skip_special_tokens=lowerCAmelCase__							)
            return generated_words





     @slow
     def       a__	(       self							):
            self._assert_generated_batch_equal_expected()
import os
import shutil
import unittest
from unittest.mock import patch

from transformers.testing_utils import CaptureStd, is_pt_tf_cross_test, require_torch





class   lowerCAmelCase    (       unittest.TestCase      ):



     @patch('sys.argv'					,   ['fakeprogrampath', 'env']							)
     def       a__	(       self							):
            # test transformers-cli env
            import transformers.commands.transformers_cli

            with CaptureStd() as cs:
                   transformers.commands.transformers_cli.main()
            self.assertIn('Python version'					,   cs.out							)
            self.assertIn('Platform'					,   cs.out							)
            self.assertIn('Using distributed or parallel set-up in script?'					,   cs.out							)



     @is_pt_tf_cross_test
     @patch(
         'sys.argv'					,   ['fakeprogrampath', 'pt-to-tf', '--model-name', 'hf-internal-testing/tiny-random-gptj', '--no-pr']							)
     def       a__	(       self							):
            import transformers.commands.transformers_cli

            shutil.rmtree('/tmp/hf-internal-testing/tiny-random-gptj'					,   ignore_errors=lowerCAmelCase__							)  # cleans potential past runs
            transformers.commands.transformers_cli.main()

            # The original repo has no TF weights -- if they exist, they were created by the CLI
            self.assertTrue(os.path.exists('/tmp/hf-internal-testing/tiny-random-gptj/tf_model.h5'							)							)



     @require_torch
     @patch('sys.argv'					,   ['fakeprogrampath', 'download', 'hf-internal-testing/tiny-random-gptj', '--cache-dir', '/tmp']							)
     def       a__	(       self							):
            import transformers.commands.transformers_cli

            # # remove any previously downloaded model to start clean
            shutil.rmtree('/tmp/models--hf-internal-testing--tiny-random-gptj'					,   ignore_errors=lowerCAmelCase__							)

            # run the command
            transformers.commands.transformers_cli.main()

            # check if the model files are downloaded correctly on /tmp/models--hf-internal-testing--tiny-random-gptj
            self.assertTrue(os.path.exists('/tmp/models--hf-internal-testing--tiny-random-gptj/blobs'							)							)
            self.assertTrue(os.path.exists('/tmp/models--hf-internal-testing--tiny-random-gptj/refs'							)							)
            self.assertTrue(os.path.exists('/tmp/models--hf-internal-testing--tiny-random-gptj/snapshots'							)							)





     @require_torch
     @patch(
         'sys.argv'					,   [
             'fakeprogrampath',
             'download',
             'hf-internal-testing/test_dynamic_model_with_tokenizer',
             '--trust-remote-code',
             '--cache-dir',
             '/tmp',
         ]					,   )
     def       a__	(       self							):
            import transformers.commands.transformers_cli

            # # remove any previously downloaded model to start clean
            shutil.rmtree('/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer'					,   ignore_errors=lowerCAmelCase__							)

            # run the command
            transformers.commands.transformers_cli.main()

            # check if the model files are downloaded correctly on /tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer
            self.assertTrue(os.path.exists('/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer/blobs'							)							)
            self.assertTrue(os.path.exists('/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer/refs'							)							)
            self.assertTrue(
                os.path.exists('/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer/snapshots'							)							)