code: 
  
  
  
  
  
  import unittest
  
  import numpy as np
  import timeout_decorator  # noqa
  
  from transformers import OPTConfig, is_flax_available
  from transformers.testing_utils import require_flax, require_sentencepiece, slow
  
  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 FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    Union[str, Any]   ,     lowercase       :    Tuple   ,     lowercase       :    Tuple=None   ,     lowercase       :    List[str]=None				) ->    Dict:
    if attention_mask is None:
      __snake_case   :				Dict        =       np.where(input_ids != config.pad_token_id   ,     1   ,     0				)
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
    }
  
  
  
  
  @require_flax
  class 				_lowerCamelCase :
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        def __init__( self				,						UpperCAmelCase				,						UpperCAmelCase=13				,						UpperCAmelCase=7				,						UpperCAmelCase=True				,						UpperCAmelCase=False				,						UpperCAmelCase=99				,						UpperCAmelCase=16				,						UpperCAmelCase=2				,						UpperCAmelCase=4				,						UpperCAmelCase=4				,						UpperCAmelCase="gelu"				,						UpperCAmelCase=0.1				,						UpperCAmelCase=0.1				,						UpperCAmelCase=20				,						UpperCAmelCase=2				,						UpperCAmelCase=1				,						UpperCAmelCase=0				,						UpperCAmelCase=16				,						UpperCAmelCase=16				,						UpperCAmelCase=0.02				,						)       ->   Any:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Tuple        =       parent
          __snake_case   :				int        =       batch_size
          __snake_case   :				Dict        =       seq_length
          __snake_case   :				Optional[int]        =       is_training
          __snake_case   :				List[str]        =       use_labels
          __snake_case   :				Optional[int]        =       vocab_size
          __snake_case   :				Optional[Any]        =       hidden_size
          __snake_case   :				str        =       num_hidden_layers
          __snake_case   :				List[Any]        =       num_attention_heads
          __snake_case   :				Any        =       intermediate_size
          __snake_case   :				Any        =       hidden_act
          __snake_case   :				List[Any]        =       hidden_dropout_prob
          __snake_case   :				Optional[int]        =       attention_probs_dropout_prob
          __snake_case   :				List[Any]        =       max_position_embeddings
          __snake_case   :				Tuple        =       eos_token_id
          __snake_case   :				Optional[int]        =       pad_token_id
          __snake_case   :				Tuple        =       bos_token_id
          __snake_case   :				int        =       embed_dim
          __snake_case   :				Any        =       word_embed_proj_dim
          __snake_case   :				str        =       initializer_range
          __snake_case   :				Optional[Any]        =       False
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Dict:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Optional[Any]        =       np.clip(ids_tensor([self.batch_size, self.seq_length - 1]				,						self.vocab_size   )				,						3				,						self.vocab_size   )
          __snake_case   :				int        =       np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)				,						dtype=np.intaa   ))				,						-1   )
  
          __snake_case   :				Union[str, Any]        =       OPTConfig(
              vocab_size=self.vocab_size				,						hidden_size=self.hidden_size				,						num_hidden_layers=self.num_hidden_layers				,						num_attention_heads=self.num_attention_heads				,						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				,						embed_dim=self.embed_dim				,						is_encoder_decoder=lowerCamelCase__				,						word_embed_proj_dim=self.word_embed_proj_dim				,						initializer_range=self.initializer_range				,						use_cache=lowerCamelCase__				,						)
          __snake_case   :				Any        =       prepare_opt_inputs_dict(lowerCamelCase__				,						lowerCamelCase__   )
          return config, inputs_dict
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   str:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Optional[int]        =       self.prepare_config_and_inputs()
          return config, inputs_dict
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )       ->   Optional[Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Tuple        =       20
          __snake_case   :				Any        =       model_class_name(lowerCamelCase__   )
  
          __snake_case   :				str        =       inputs_dict["input_ids"]
          __snake_case   :				str        =       inputs_dict["attention_mask"]
  
          __snake_case   :				List[str]        =       model.init_cache(input_ids.shape[0]				,						lowerCamelCase__   )
          __snake_case   :				Optional[int]        =       jnp.ones((input_ids.shape[0], max_length)				,						dtype="i4"   )
  
          __snake_case   :				Dict        =       jnp.broadcast_to(
              jnp.arange(input_ids.shape[-1] - 1   )[None, :]				,						(input_ids.shape[0], input_ids.shape[-1] - 1)				,						)
          __snake_case   :				List[Any]        =       model(
              input_ids[:, :-1]				,						attention_mask=lowerCamelCase__				,						past_key_values=lowerCamelCase__				,						position_ids=lowerCamelCase__				,						)
  
          __snake_case   :				Tuple        =       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,						dtype="i4"   )
          __snake_case   :				Any        =       model(
              input_ids[:, -1:]				,						attention_mask=lowerCamelCase__				,						past_key_values=outputs_cache.past_key_values				,						position_ids=lowerCamelCase__				,						)
  
          __snake_case   :				str        =       model(lowerCamelCase__   )
  
          __snake_case   :				List[str]        =       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 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )       ->   Dict:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Optional[Any]        =       20
          __snake_case   :				Tuple        =       model_class_name(lowerCamelCase__   )
  
          __snake_case   :				Any        =       (
              inputs_dict["input_ids"],
              inputs_dict["attention_mask"],
          )
  
          __snake_case   :				int        =       jnp.concatenate(
              [
                  attention_mask,
                  jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])   ),
              ]				,						axis=-1				,						)
  
          __snake_case   :				Tuple        =       model.init_cache(input_ids.shape[0]				,						lowerCamelCase__   )
          __snake_case   :				str        =       jnp.broadcast_to(
              jnp.arange(input_ids.shape[-1] - 1   )[None, :]				,						(input_ids.shape[0], input_ids.shape[-1] - 1)				,						)
  
          __snake_case   :				Optional[int]        =       model(
              input_ids[:, :-1]				,						attention_mask=lowerCamelCase__				,						past_key_values=lowerCamelCase__				,						position_ids=lowerCamelCase__				,						)
          __snake_case   :				Tuple        =       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,						dtype="i4"   )
          __snake_case   :				Dict        =       model(
              input_ids[:, -1:]				,						past_key_values=outputs_cache.past_key_values				,						attention_mask=lowerCamelCase__				,						position_ids=lowerCamelCase__				,						)
  
          __snake_case   :				int        =       model(lowerCamelCase__				,						attention_mask=lowerCamelCase__   )
  
          __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}"""   )
  
  
  
  
  @require_flax
  class 				_lowerCamelCase (    _a ,   unittest.TestCase ,   _a    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        UpperCAmelCase_						:      List[str]       =(FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
        UpperCAmelCase_						:      Optional[Any]       =() if is_flax_available() else ()
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   List[Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				List[Any]        =       FlaxOPTModelTester(self   )
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   int:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __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(lowerCamelCase__				,						lowerCamelCase__				,						lowerCamelCase__   )
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Optional[int]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __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_with_attn_mask(lowerCamelCase__				,						lowerCamelCase__				,						lowerCamelCase__   )
  
  
  
  
  
  
  
        @slow
        def 						UpperCAmelCase		( self   )       ->   Optional[Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          for model_class_name in self.all_model_classes:
            __snake_case   :				Union[str, Any]        =       model_class_name.from_pretrained("facebook/opt-125m"   )
            __snake_case   :				int        =       np.ones((1, 1)   ) * model.config.eos_token_id
            __snake_case   :				Optional[int]        =       model(lowerCamelCase__   )
            self.assertIsNotNone(lowerCamelCase__   )
  
  
  
  
  @require_sentencepiece
  @require_flax
  class 				_lowerCamelCase (    unittest.TestCase    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        @slow
        def 						UpperCAmelCase		( self   )       ->   List[str]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				List[Any]        =       FlaxOPTModel.from_pretrained("facebook/opt-350m"   )
          __snake_case   :				Tuple        =       jnp.array([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]   )
          __snake_case   :				int        =       model(input_ids=lowerCamelCase__   ).last_hidden_state
          __snake_case   :				Optional[int]        =       (1, 11, 512)
          self.assertEqual(output.shape				,						lowerCamelCase__   )
          __snake_case   :				List[Any]        =       jnp.array(
              [[-0.2_867, -1.9_256, -0.3_062], [-1.2_711, -0.1_337, -0.1_897], [0.4_109, 0.1_187, -1.3_142]]   )
          self.assertTrue(jnp.allclose(output[:, :3, :3]				,						lowerCamelCase__				,						atol=4E-2   )   )
  
  
  
  
  @require_flax
  @slow
  class 				_lowerCamelCase (    unittest.TestCase    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   int:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          super().setUp()
          __snake_case   :				Optional[Any]        =       "facebook/opt-350m"
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Dict:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Union[str, Any]        =       FlaxOPTForCausalLM.from_pretrained(self.path_model   )
          __snake_case   :				Dict        =       GPTaTokenizer.from_pretrained(self.path_model   )
  
          __snake_case   :				List[Any]        =       [
              "Today is a beautiful day and I want to",
              "In the city of",
              "Paris is the capital of France and",
              "Computers and mobile phones have taken",
          ]
          # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
          __snake_case   :				List[Any]        =       tokenizer(lowerCamelCase__				,						return_tensors="jax"				,						padding=lowerCamelCase__				,						add_special_tokens=lowerCamelCase__   )
          __snake_case   :				List[Any]        =       model(inputs.input_ids				,						attention_mask=inputs.attention_mask   )[0].mean(axis=-1   )
          __snake_case   :				Dict        =       jnp.array(
              [
                  [1.3_851, -13.8_923, -10.5_229, -10.7_533, -0.2_309, -10.2_384, -0.5_365, -9.0_947, -5.1_670],
                  [-4.7_073, -10.6_276, -3.9_415, -21.5_242, -0.2_822, -0.2_822, -0.2_822, -0.2_822, -0.2_822],
                  [0.6_247, -3.4_229, -8.9_179, -1.4_297, -14.1_650, 1.4_146, -9.0_218, -0.2_703, -0.2_703],
                  [6.4_783, -1.9_913, -10.7_926, -2.3_336, 1.5_092, -0.9_974, -6.8_213, 1.3_477, 1.3_477],
              ]   )
          self.assertTrue(jnp.allclose(lowerCamelCase__				,						lowerCamelCase__				,						atol=4E-2   )   )
  
          __snake_case   :				Optional[Any]        =       jax.jit(lowerCamelCase__   )
          __snake_case   :				List[Any]        =       model(inputs.input_ids				,						attention_mask=inputs.attention_mask   )[0].mean(axis=-1   )
          self.assertTrue(jnp.allclose(lowerCamelCase__				,						lowerCamelCase__				,						atol=4E-2   )   )
  
  
  
  
  @require_flax
  @slow
  class 				_lowerCamelCase (    unittest.TestCase    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        @property
        def 						UpperCAmelCase		( self   )       ->   Tuple:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          return [
              "Today is a beautiful day and I want",
              "In the city of",
              "Paris is the capital of France and",
              "Computers and mobile phones have taken",
          ]
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Dict:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				List[Any]        =       "facebook/opt-125m"
  
          __snake_case   :				Tuple        =       [
              "Today is a beautiful day and I want to",
              "In the city of New York, the city",
              "Paris is the capital of France and the capital",
              "Computers and mobile phones have taken over the",
          ]
  
          __snake_case   :				Tuple        =       []
  
          __snake_case   :				int        =       FlaxOPTForCausalLM.from_pretrained(lowerCamelCase__   )
          __snake_case   :				int        =       GPTaTokenizer.from_pretrained(lowerCamelCase__   )
  
          for prompt in self.prompts:
            __snake_case   :				Tuple        =       tokenizer(lowerCamelCase__				,						return_tensors="jax"   ).input_ids
  
            __snake_case   :				Optional[Any]        =       model.generate(lowerCamelCase__				,						max_length=10   )
            __snake_case   :				Optional[int]        =       generated_ids[0]
  
            __snake_case   :				List[str]        =       tokenizer.batch_decode(lowerCamelCase__				,						skip_special_tokens=lowerCamelCase__   )
            predicted_outputs += generated_string
  
          self.assertListEqual(lowerCamelCase__				,						lowerCamelCase__   )
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   int:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Union[str, Any]        =       "facebook/opt-350m"
  
          __snake_case   :				List[str]        =       [
              "Today is a beautiful day and I want to",
              "In the city of San Francisco, the city",
              "Paris is the capital of France and the capital",
              "Computers and mobile phones have taken over the",
          ]
  
          __snake_case   :				Tuple        =       []
          __snake_case   :				List[Any]        =       FlaxOPTForCausalLM.from_pretrained(lowerCamelCase__   )
          __snake_case   :				Any        =       GPTaTokenizer.from_pretrained(lowerCamelCase__   )
  
          for prompt in self.prompts:
            __snake_case   :				Dict        =       tokenizer(lowerCamelCase__				,						return_tensors="jax"   ).input_ids
  
            __snake_case   :				Dict        =       model.generate(lowerCamelCase__				,						max_length=10   )
            __snake_case   :				str        =       generated_ids[0]
  
            __snake_case   :				str        =       tokenizer.batch_decode(lowerCamelCase__				,						skip_special_tokens=lowerCamelCase__   )
            predicted_outputs += generated_string
  
          self.assertListEqual(lowerCamelCase__				,						lowerCamelCase__   )
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   str:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Dict        =       "facebook/opt-125m"
          __snake_case   :				Optional[Any]        =       [
              "Today is a beautiful day and I want to thank",
              "In the city of Rome Canaver Canaver Canaver Canaver",
          ]
          __snake_case   :				int        =       FlaxOPTForCausalLM.from_pretrained(lowerCamelCase__   )
          __snake_case   :				List[Any]        =       GPTaTokenizer.from_pretrained(lowerCamelCase__   )
          __snake_case   :				List[str]        =       tokenizer(
              [
                  "Today is a beautiful day and I want to",
                  "In the city of",
              ]				,						return_tensors="jax"				,						padding=lowerCamelCase__				,						)
  
          __snake_case   :				int        =       jax.jit(model.generate   )
  
          __snake_case   :				str        =       jit_generate(inputs["input_ids"]				,						attention_mask=inputs["attention_mask"]   ).sequences
  
          __snake_case   :				Optional[int]        =       tokenizer.batch_decode(lowerCamelCase__				,						skip_special_tokens=lowerCamelCase__   )
  
          self.assertIsNotNone(lowerCamelCase__				,						lowerCamelCase__   )
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Any:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Tuple        =       "facebook/opt-350m"
  
          __snake_case   :				List[Any]        =       GPTaTokenizer.from_pretrained(lowerCamelCase__   )
          __snake_case   :				Union[str, Any]        =       FlaxOPTForCausalLM.from_pretrained(lowerCamelCase__   )
  
          __snake_case   :				List[str]        =       "left"
  
          # use different length sentences to test batching
          __snake_case   :				Tuple        =       [
              "Hello, my dog is a little",
              "Today, I",
          ]
  
          __snake_case   :				int        =       tokenizer(lowerCamelCase__				,						return_tensors="jax"				,						padding=lowerCamelCase__   )
          __snake_case   :				List[str]        =       inputs["input_ids"]
  
          __snake_case   :				Tuple        =       model.generate(input_ids=lowerCamelCase__				,						attention_mask=inputs["attention_mask"]				,						trace=lowerCamelCase__   )
  
          __snake_case   :				List[Any]        =       tokenizer(sentences[0]				,						return_tensors="jax"   ).input_ids
          __snake_case   :				Union[str, Any]        =       model.generate(input_ids=lowerCamelCase__   )
  
          __snake_case   :				Tuple        =       inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].sum()
          __snake_case   :				Dict        =       tokenizer(sentences[1]				,						return_tensors="jax"   ).input_ids
          __snake_case   :				int        =       model.generate(input_ids=lowerCamelCase__				,						max_length=model.config.max_length - num_paddings   )
  
          __snake_case   :				List[str]        =       tokenizer.batch_decode(outputs[0]				,						skip_special_tokens=lowerCamelCase__   )
          __snake_case   :				Tuple        =       tokenizer.decode(output_non_padded[0][0]				,						skip_special_tokens=lowerCamelCase__   )
          __snake_case   :				Tuple        =       tokenizer.decode(output_padded[0][0]				,						skip_special_tokens=lowerCamelCase__   )
  
          __snake_case   :				Union[str, Any]        =       [
              "Hello, my dog is a little bit of a dork.\nI'm a little bit",
              "Today, I was in the middle of a conversation with a friend about the",
          ]
          self.assertListEqual(lowerCamelCase__				,						lowerCamelCase__   )
          self.assertListEqual(lowerCamelCase__				,						[non_padded_sentence, padded_sentence]   )
  
  
style_context: 
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    str				) ->    int:
    __snake_case   :				List[str]        =       str(lowercase				).strip()
    if not oct_string:
      raise ValueError("Empty string was passed to the function"				)
    __snake_case   :				int        =       oct_string[0] == "-"
    if is_negative:
      __snake_case   :				Union[str, Any]        =       oct_string[1:]
    if not oct_string.isdigit() or not all(0 <= int(lowercase				) <= 7 for char in oct_string				):
      raise ValueError("Non-octal value was passed to the function"				)
    __snake_case   :				List[Any]        =       0
    for char in oct_string:
      __snake_case   :				Optional[Any]        =       8 * decimal_number + int(lowercase				)
    if is_negative:
      __snake_case   :				Any        =       -decimal_number
    return decimal_number
  
  
  if __name__ == "__main__":
       from doctest import testmod
  
       testmod()
  
  