
from __future__ import annotations

import unittest

import numpy as np

from transformers import OPTConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import GPTaTokenizer, TFOPTForCausalLM, TFOPTModel

def a__    (       A_,     A_,     A_=None,     A_=None					):






  '''simple docstring'''





  if attention_mask is None:
    __magic_name__				     =       tf.cast(tf.math.not_equal(A_,     config.pad_token_id					),     tf.inta					)
  return {"input_ids": input_ids, "attention_mask": attention_mask}




@require_tf
class        UpperCAmelCase_			:


  '''simple docstring'''







  a__           =		OPTConfig
  a__           =		{}
  a__           =		"""gelu"""
  def __init__( self			: Dict			,							UpperCamelCase__			: Optional[int]			,							UpperCamelCase__			: List[Any]=13			,							UpperCamelCase__			: List[str]=7			,							UpperCamelCase__			: Optional[int]=True			,							UpperCamelCase__			: Any=False			,							UpperCamelCase__			: Tuple=99			,							UpperCamelCase__			: str=16			,							UpperCamelCase__			: Optional[Any]=2			,							UpperCamelCase__			: List[Any]=4			,							UpperCamelCase__			: int=4			,							UpperCamelCase__			: Optional[int]="gelu"			,							UpperCamelCase__			: Optional[Any]=0.1			,							UpperCamelCase__			: str=0.1			,							UpperCamelCase__			: List[str]=20			,							UpperCamelCase__			: int=2			,							UpperCamelCase__			: str=1			,							UpperCamelCase__			: Optional[int]=0			,							UpperCamelCase__			: List[Any]=16			,							UpperCamelCase__			: str=16			,							)			->	str:







    """simple docstring"""






    __magic_name__				     =       parent
    __magic_name__				     =       batch_size
    __magic_name__				     =       seq_length
    __magic_name__				     =       is_training
    __magic_name__				     =       use_labels
    __magic_name__				     =       vocab_size
    __magic_name__				     =       hidden_size
    __magic_name__				     =       num_hidden_layers
    __magic_name__				     =       num_attention_heads
    __magic_name__				     =       intermediate_size
    __magic_name__				     =       hidden_act
    __magic_name__				     =       hidden_dropout_prob
    __magic_name__				     =       attention_probs_dropout_prob
    __magic_name__				     =       max_position_embeddings
    __magic_name__				     =       eos_token_id
    __magic_name__				     =       pad_token_id
    __magic_name__				     =       bos_token_id
    __magic_name__				     =       embed_dim
    __magic_name__				     =       word_embed_proj_dim
    __magic_name__				     =       False
  def 							_lowercase							( self			: int    )			->	Optional[Any]:







    """simple docstring"""






    __magic_name__				     =       ids_tensor([self.batch_size, self.seq_length - 1]			,							self.vocab_size    )
    __magic_name__				     =       tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size    )			,							1    )
    __magic_name__				     =       tf.concat([input_ids, eos_tensor]			,							axis=1    )

    __magic_name__				     =       self.config_cls(
        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			,							word_embed_proj_dim=self.word_embed_proj_dim			,							is_encoder_decoder=UpperCamelCase__			,							**self.config_updates			,							)
    __magic_name__				     =       prepare_opt_inputs_dict(UpperCamelCase__			,							UpperCamelCase__    )
    return config, inputs_dict





  def 							_lowercase							( self			: Dict			,							UpperCamelCase__			: Dict			,							UpperCamelCase__			: List[Any]    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       TFOPTModel(config=UpperCamelCase__    )
    __magic_name__				     =       inputs_dict["""input_ids"""]

    __magic_name__				     =       input_ids[:1, :]
    __magic_name__				     =       inputs_dict["""attention_mask"""][:1, :]
    __magic_name__				     =       1

    # first forward pass
    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__			,							use_cache=UpperCamelCase__    )

    __magic_name__     ,    __magic_name__				     =       outputs.to_tuple()

    # create hypothetical next token and extent to next_input_ids
    __magic_name__				     =       ids_tensor((self.batch_size, 3)			,							config.vocab_size    )
    __magic_name__				     =       tf.cast(ids_tensor((self.batch_size, 3)			,							2    )			,							tf.inta    )

    # append to next input_ids and
    __magic_name__				     =       tf.concat([input_ids, next_tokens]			,							axis=-1    )
    __magic_name__				     =       tf.concat([attention_mask, next_attn_mask]			,							axis=-1    )

    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__    )[0]
    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__			,							past_key_values=UpperCamelCase__    )[0]

    self.parent.assertEqual(next_tokens.shape[1]			,							output_from_past.shape[1]    )

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

    # test that outputs are equal for slice
    tf.debugging.assert_near(UpperCamelCase__			,							UpperCamelCase__			,							rtol=1E-3    )







@require_tf
class        UpperCAmelCase_			( _A  ,       _A  ,       unittest.TestCase     ):


  '''simple docstring'''







  a__           =		(TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
  a__           =		(TFOPTForCausalLM,) if is_tf_available() else ()
  a__           =		(
      {"""feature-extraction""": TFOPTModel, """text-generation""": TFOPTForCausalLM} if is_tf_available() else {}
  )
  a__           =		False
  a__           =		False
  a__           =		False
  a__           =		10
  def 							_lowercase							( self			: str    )			->	Union[str, Any]:







    """simple docstring"""






    __magic_name__				     =       TFOPTModelTester(self    )
    __magic_name__				     =       ConfigTester(self			,							config_class=UpperCamelCase__    )
  def 							_lowercase							( self			: List[str]    )			->	Optional[Any]:







    """simple docstring"""






    self.config_tester.run_common_tests()
  def 							_lowercase							( self			: Any    )			->	Any:







    """simple docstring"""






    __magic_name__				     =       self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.check_decoder_model_past_large_inputs(*UpperCamelCase__    )





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







    """simple docstring"""






    __magic_name__     ,    __magic_name__				     =       self.model_tester.prepare_config_and_inputs_for_common()

    def _get_word_embedding_weight(UpperCamelCase__			: Any			,							UpperCamelCase__			: str    ):
      if hasattr(UpperCamelCase__			,							"""weight"""    ):
        return embedding_layer.weight
      else:
        # Here we build the word embeddings weights if not exists.
        # And then we retry to get the attribute once built.
        model.build()
        if hasattr(UpperCamelCase__			,							"""weight"""    ):
          return embedding_layer.weight
        else:
          return None

    for model_class in self.all_model_classes:
      for size in [config.vocab_size - 10, config.vocab_size + 10]:
        # build the embeddings
        __magic_name__				     =       model_class(config=UpperCamelCase__    )
        __magic_name__				     =       _get_word_embedding_weight(UpperCamelCase__			,							model.get_input_embeddings()    )
        __magic_name__				     =       _get_word_embedding_weight(UpperCamelCase__			,							model.get_output_embeddings()    )

        # reshape the embeddings
        model.resize_token_embeddings(UpperCamelCase__    )
        __magic_name__				     =       _get_word_embedding_weight(UpperCamelCase__			,							model.get_input_embeddings()    )
        __magic_name__				     =       _get_word_embedding_weight(UpperCamelCase__			,							model.get_output_embeddings()    )

        # check that the resized embeddings size matches the desired size.
        __magic_name__				     =       size if size is not None else config.vocab_size

        self.assertEqual(new_input_embeddings.shape[0]			,							UpperCamelCase__    )

        # check that weights remain the same after resizing
        __magic_name__				     =       True
        for pa, pa in zip(old_input_embeddings.value()			,							new_input_embeddings.value()    ):
          if tf.math.reduce_sum(tf.math.abs(pa - pa    )    ) > 0:
            __magic_name__				     =       False
        self.assertTrue(UpperCamelCase__    )

        if old_output_embeddings is not None and new_output_embeddings is not None:
          self.assertEqual(new_output_embeddings.shape[0]			,							UpperCamelCase__    )

          __magic_name__				     =       True
          for pa, pa in zip(old_output_embeddings.value()			,							new_output_embeddings.value()    ):
            if tf.math.reduce_sum(tf.math.abs(pa - pa    )    ) > 0:
              __magic_name__				     =       False
          self.assertTrue(UpperCamelCase__    )







def a__    (       A_					):






  '''simple docstring'''





  return tf.constant(A_,     dtype=tf.intaa					)




@require_tf
class        UpperCAmelCase_			( unittest.TestCase     ):


  '''simple docstring'''







  a__           =		99
  def 							_lowercase							( self			: str    )			->	Optional[Any]:







    """simple docstring"""






    __magic_name__				     =       tf.ones((4, 1)			,							dtype=tf.intaa    ) * 2
    __magic_name__				     =       tf.concat([ids_tensor((4, 6)			,							self.vocab_size - 3    ) + 3, eos_column_vector]			,							axis=1    )
    __magic_name__				     =       input_ids.shape[0]
    __magic_name__				     =       OPTConfig(
        vocab_size=self.vocab_size			,							hidden_size=24			,							num_hidden_layers=2			,							num_attention_heads=2			,							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







@require_sentencepiece
@require_tf
class        UpperCAmelCase_			( unittest.TestCase     ):


  '''simple docstring'''







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







    """simple docstring"""






    __magic_name__				     =       TFOPTModel.from_pretrained("""facebook/opt-350m"""    )
    __magic_name__				     =       _long_tensor([[0, 3_1414, 232, 328, 740, 1140, 1_2695, 69, 4_6078, 1588, 2]]    )
    __magic_name__				     =       tf.not_equal(UpperCamelCase__			,							model.config.pad_token_id    )
    with tf.GradientTape():
      __magic_name__				     =       model(input_ids=UpperCamelCase__			,							attention_mask=UpperCamelCase__    ).last_hidden_state
    __magic_name__				     =       (1, 11, 512)
    self.assertEqual(output.shape			,							UpperCamelCase__    )
    __magic_name__				     =       tf.constant(
        [[-0.2873, -1.9218, -0.3033], [-1.2710, -0.1338, -0.1902], [0.4095, 0.1214, -1.3121]]    )
    self.assertTrue(np.allclose(output[:, :3, :3]			,							UpperCamelCase__			,							atol=4E-3    )    )

    __magic_name__				     =       tf.function(UpperCamelCase__			,							jit_compile=UpperCamelCase__    )
    __magic_name__				     =       xla_generate(UpperCamelCase__			,							UpperCamelCase__    )[0]
    self.assertTrue(np.allclose(output[:, :3, :3]			,							UpperCamelCase__			,							atol=4E-2    )    )







@require_tf
@slow
class        UpperCAmelCase_			( unittest.TestCase     ):


  '''simple docstring'''







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







    """simple docstring"""






    super().setUp()
    __magic_name__				     =       """facebook/opt-350m"""





  def 							_lowercase							( self			: str    )			->	Optional[Any]:







    """simple docstring"""






    __magic_name__				     =       TFOPTForCausalLM.from_pretrained(self.path_model    )
    __magic_name__				     =       GPTaTokenizer.from_pretrained(self.path_model    )

    __magic_name__				     =       [
        """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
    __magic_name__				     =       tokenizer(UpperCamelCase__			,							return_tensors="""tf"""			,							padding=UpperCamelCase__			,							add_special_tokens=UpperCamelCase__    )
    __magic_name__				     =       tf.math.reduce_mean(model(inputs.input_ids			,							attention_mask=inputs.attention_mask    )[0]			,							axis=-1    )
    __magic_name__				     =       tf.constant(
        [
            [1.3851, -13.8923, -10.5229, -10.7533, -0.2309, -10.2384, -0.5365, -9.0947, -5.1670],
            [-4.7073, -10.6276, -3.9415, -21.5242, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
            [0.6247, -3.4229, -8.9179, -1.4297, -14.1650, 1.4146, -9.0218, -0.2703, -0.2703],
            [6.4783, -1.9913, -10.7926, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
        ]    )
    self.assertTrue(np.allclose(UpperCamelCase__			,							UpperCamelCase__			,							atol=1E-4    )    )

    __magic_name__				     =       tf.function(UpperCamelCase__			,							jit_compile=UpperCamelCase__    )
    __magic_name__				     =       tf.math.reduce_mean(xla_generate(inputs.input_ids			,							attention_mask=inputs.attention_mask    )[0]			,							axis=-1    )
    self.assertTrue(np.allclose(UpperCamelCase__			,							UpperCamelCase__			,							atol=1E-4    )    )







@require_tf
@slow
class        UpperCAmelCase_			( unittest.TestCase     ):


  '''simple docstring'''







  @property
  def 							_lowercase							( self			: Any    )			->	int:







    """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 							_lowercase							( self			: Dict    )			->	List[Any]:







    """simple docstring"""






    __magic_name__				     =       """facebook/opt-125m"""

    __magic_name__				     =       [
        """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""",
    ]

    __magic_name__				     =       []
    __magic_name__				     =       GPTaTokenizer.from_pretrained(UpperCamelCase__    )
    __magic_name__				     =       TFOPTForCausalLM.from_pretrained(UpperCamelCase__    )

    for prompt in self.prompts:
      __magic_name__				     =       tokenizer(UpperCamelCase__			,							return_tensors="""tf"""    ).input_ids

      __magic_name__				     =       model.generate(UpperCamelCase__			,							max_length=10    )

      __magic_name__				     =       tokenizer.batch_decode(UpperCamelCase__			,							skip_special_tokens=UpperCamelCase__    )
      predicted_outputs += generated_string

    self.assertListEqual(UpperCamelCase__			,							UpperCamelCase__    )
  def 							_lowercase							( self			: Dict    )			->	Any:







    """simple docstring"""






    __magic_name__				     =       """facebook/opt-350m"""

    __magic_name__				     =       GPTaTokenizer.from_pretrained(UpperCamelCase__    )
    __magic_name__				     =       TFOPTForCausalLM.from_pretrained(UpperCamelCase__    )

    __magic_name__				     =       """left"""

    # use different length sentences to test batching
    __magic_name__				     =       [
        """Hello, my dog is a little""",
        """Today, I""",
    ]

    __magic_name__				     =       tokenizer(UpperCamelCase__			,							return_tensors="""tf"""			,							padding=UpperCamelCase__    )
    __magic_name__				     =       inputs["""input_ids"""]

    __magic_name__				     =       model.generate(input_ids=UpperCamelCase__			,							attention_mask=inputs["""attention_mask"""]    )

    __magic_name__				     =       tokenizer(sentences[0]			,							return_tensors="""tf"""    ).input_ids
    __magic_name__				     =       model.generate(input_ids=UpperCamelCase__    )

    __magic_name__				     =       inputs_non_padded.shape[-1] - tf.math.reduce_sum(
        tf.cast(inputs["""attention_mask"""][-1]			,							tf.intaa    )    )
    __magic_name__				     =       tokenizer(sentences[1]			,							return_tensors="""tf"""    ).input_ids
    __magic_name__				     =       model.generate(input_ids=UpperCamelCase__			,							max_length=model.config.max_length - num_paddings    )

    __magic_name__				     =       tokenizer.batch_decode(UpperCamelCase__			,							skip_special_tokens=UpperCamelCase__    )
    __magic_name__				     =       tokenizer.decode(output_non_padded[0]			,							skip_special_tokens=UpperCamelCase__    )
    __magic_name__				     =       tokenizer.decode(output_padded[0]			,							skip_special_tokens=UpperCamelCase__    )

    __magic_name__				     =       [
        """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(UpperCamelCase__			,							UpperCamelCase__    )
    self.assertListEqual(UpperCamelCase__			,							[non_padded_sentence, padded_sentence]    )





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







    """simple docstring"""






    __magic_name__				     =       """facebook/opt-350m"""

    __magic_name__				     =       [
        """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""",
    ]

    __magic_name__				     =       []
    __magic_name__				     =       GPTaTokenizer.from_pretrained(UpperCamelCase__    )
    __magic_name__				     =       TFOPTForCausalLM.from_pretrained(UpperCamelCase__    )

    for prompt in self.prompts:
      __magic_name__				     =       tokenizer(UpperCamelCase__			,							return_tensors="""tf"""    ).input_ids

      __magic_name__				     =       model.generate(UpperCamelCase__			,							max_length=10    )

      __magic_name__				     =       tokenizer.batch_decode(UpperCamelCase__			,							skip_special_tokens=UpperCamelCase__    )
      predicted_outputs += generated_string

    self.assertListEqual(UpperCamelCase__			,							UpperCamelCase__    )



import pickle
import unittest

import torch

from accelerate import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils import require_cpu




@require_cpu
class        UpperCAmelCase_			( unittest.TestCase     ):


  '''simple docstring'''







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







    """simple docstring"""






    __magic_name__				     =       torch.nn.Linear(10			,							10    )
    __magic_name__				     =       torch.optim.SGD(model.parameters()			,							0.1    )
    __magic_name__				     =       Accelerator()
    __magic_name__				     =       accelerator.prepare(UpperCamelCase__    )
    try:
      pickle.loads(pickle.dumps(UpperCamelCase__    )    )
    except Exception as e:
      self.fail(F'''Accelerated optimizer pickling failed with {e}'''    )
    AcceleratorState._reset_state()

