




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  UpperCAmelCase_							(				__UpperCAmelCase						:     Dict      ,      __UpperCAmelCase						:     Optional[int]      ,      __UpperCAmelCase						:     int=None      ,      __UpperCAmelCase						:     Tuple=None							) ->			Optional[int]:
						if attention_mask is None:
												SCREAMING_SNAKE_CASE_     						=      tf.cast(tf.math.not_equal(__UpperCAmelCase      ,      config.pad_token_id							)      ,      tf.inta							)
						return {"input_ids": input_ids, "attention_mask": attention_mask}





@require_tf
class 	lowerCamelCase_							:


	'''simple docstring'''







	lowercase_       =   OPTConfig
	lowercase_       =   {}
	lowercase_       =   "gelu"



	def __init__(							self  :		Any    ,    _lowerCAmelCase  :		Dict    ,    _lowerCAmelCase  :		Tuple=13    ,    _lowerCAmelCase  :		int=7    ,    _lowerCAmelCase  :		Optional[Any]=True    ,    _lowerCAmelCase  :		Any=False    ,    _lowerCAmelCase  :		List[str]=99    ,    _lowerCAmelCase  :		Tuple=16    ,    _lowerCAmelCase  :		Dict=2    ,    _lowerCAmelCase  :		str=4    ,    _lowerCAmelCase  :		Union[str, Any]=4    ,    _lowerCAmelCase  :		Any="gelu"    ,    _lowerCAmelCase  :		Tuple=0.1    ,    _lowerCAmelCase  :		List[str]=0.1    ,    _lowerCAmelCase  :		Optional[int]=20    ,    _lowerCAmelCase  :		Union[str, Any]=2    ,    _lowerCAmelCase  :		Dict=1    ,    _lowerCAmelCase  :		List[str]=0    ,    _lowerCAmelCase  :		int=16    ,    _lowerCAmelCase  :		List[str]=16    ,    ):
							SCREAMING_SNAKE_CASE_     						=      parent
							SCREAMING_SNAKE_CASE_     						=      batch_size
							SCREAMING_SNAKE_CASE_     						=      seq_length
							SCREAMING_SNAKE_CASE_     						=      is_training
							SCREAMING_SNAKE_CASE_     						=      use_labels
							SCREAMING_SNAKE_CASE_     						=      vocab_size
							SCREAMING_SNAKE_CASE_     						=      hidden_size
							SCREAMING_SNAKE_CASE_     						=      num_hidden_layers
							SCREAMING_SNAKE_CASE_     						=      num_attention_heads
							SCREAMING_SNAKE_CASE_     						=      intermediate_size
							SCREAMING_SNAKE_CASE_     						=      hidden_act
							SCREAMING_SNAKE_CASE_     						=      hidden_dropout_prob
							SCREAMING_SNAKE_CASE_     						=      attention_probs_dropout_prob
							SCREAMING_SNAKE_CASE_     						=      max_position_embeddings
							SCREAMING_SNAKE_CASE_     						=      eos_token_id
							SCREAMING_SNAKE_CASE_     						=      pad_token_id
							SCREAMING_SNAKE_CASE_     						=      bos_token_id
							SCREAMING_SNAKE_CASE_     						=      embed_dim
							SCREAMING_SNAKE_CASE_     						=      word_embed_proj_dim
							SCREAMING_SNAKE_CASE_     						=      False



	def 				lowerCAmelCase_  (							self  :		List[str]						):
							SCREAMING_SNAKE_CASE_     						=      ids_tensor([self.batch_size, self.seq_length - 1]    ,    self.vocab_size						)
							SCREAMING_SNAKE_CASE_     						=      tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size						)    ,    1						)
							SCREAMING_SNAKE_CASE_     						=      tf.concat([input_ids, eos_tensor]    ,    axis=1						)

							SCREAMING_SNAKE_CASE_     						=      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=_lowerCAmelCase    ,    **self.config_updates    ,    )
							SCREAMING_SNAKE_CASE_     						=      prepare_opt_inputs_dict(_lowerCAmelCase    ,    _lowerCAmelCase						)
							return config, inputs_dict




	def 				lowerCAmelCase_  (							self  :		Union[str, Any]    ,    _lowerCAmelCase  :		Union[str, Any]    ,    _lowerCAmelCase  :		List[str]						):
							SCREAMING_SNAKE_CASE_     						=      TFOPTModel(config=_lowerCAmelCase						)
							SCREAMING_SNAKE_CASE_     						=      inputs_dict['input_ids']

							SCREAMING_SNAKE_CASE_     						=      input_ids[:1, :]
							SCREAMING_SNAKE_CASE_     						=      inputs_dict['attention_mask'][:1, :]
							SCREAMING_SNAKE_CASE_     						=      1

							# first forward pass
							SCREAMING_SNAKE_CASE_     						=      model(_lowerCAmelCase    ,    attention_mask=_lowerCAmelCase    ,    use_cache=_lowerCAmelCase						)

							SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_     						=      outputs.to_tuple()

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

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

							SCREAMING_SNAKE_CASE_     						=      model(_lowerCAmelCase    ,    attention_mask=_lowerCAmelCase						)[0]
							SCREAMING_SNAKE_CASE_     						=      model(_lowerCAmelCase    ,    attention_mask=_lowerCAmelCase    ,    past_key_values=_lowerCAmelCase						)[0]

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

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

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





@require_tf
class 	lowerCamelCase_							(				_SCREAMING_SNAKE_CASE	,       _SCREAMING_SNAKE_CASE	,       unittest.TestCase ):


	'''simple docstring'''







	lowercase_       =   (TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
	lowercase_       =   (TFOPTForCausalLM,) if is_tf_available() else ()
	lowercase_       =   (
	    {"feature-extraction": TFOPTModel, "text-generation": TFOPTForCausalLM} if is_tf_available() else {}
	)
	lowercase_       =   False
	lowercase_       =   False
	lowercase_       =   False
	lowercase_       =   10



	def 				lowerCAmelCase_  (							self  :		str						):
							SCREAMING_SNAKE_CASE_     						=      TFOPTModelTester(self						)
							SCREAMING_SNAKE_CASE_     						=      ConfigTester(self    ,    config_class=_lowerCAmelCase						)



	def 				lowerCAmelCase_  (							self  :		Tuple						):
							self.config_tester.run_common_tests()



	def 				lowerCAmelCase_  (							self  :		int						):
							SCREAMING_SNAKE_CASE_     						=      self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.check_decoder_model_past_large_inputs(*_lowerCAmelCase						)




	def 				lowerCAmelCase_  (							self  :		List[str]						):
							SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_     						=      self.model_tester.prepare_config_and_inputs_for_common()

							def _get_word_embedding_weight(_lowerCAmelCase  :		Tuple    ,    _lowerCAmelCase  :		Tuple						):
													if hasattr(_lowerCAmelCase    ,    '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(_lowerCAmelCase    ,    '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
																			SCREAMING_SNAKE_CASE_     						=      model_class(config=_lowerCAmelCase						)
																			SCREAMING_SNAKE_CASE_     						=      _get_word_embedding_weight(_lowerCAmelCase    ,    model.get_input_embeddings()						)
																			SCREAMING_SNAKE_CASE_     						=      _get_word_embedding_weight(_lowerCAmelCase    ,    model.get_output_embeddings()						)

																			# reshape the embeddings
																			model.resize_token_embeddings(_lowerCAmelCase						)
																			SCREAMING_SNAKE_CASE_     						=      _get_word_embedding_weight(_lowerCAmelCase    ,    model.get_input_embeddings()						)
																			SCREAMING_SNAKE_CASE_     						=      _get_word_embedding_weight(_lowerCAmelCase    ,    model.get_output_embeddings()						)

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

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

																			# check that weights remain the same after resizing
																			SCREAMING_SNAKE_CASE_     						=      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:
																															SCREAMING_SNAKE_CASE_     						=      False
																			self.assertTrue(_lowerCAmelCase						)

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

																									SCREAMING_SNAKE_CASE_     						=      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:
																																					SCREAMING_SNAKE_CASE_     						=      False
																									self.assertTrue(_lowerCAmelCase						)


def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[int]							) ->			List[Any]:
						return tf.constant(__UpperCAmelCase      ,      dtype=tf.intaa							)





@require_tf
class 	lowerCamelCase_							(				unittest.TestCase ):


	'''simple docstring'''







	lowercase_       =   99



	def 				lowerCAmelCase_  (							self  :		List[str]						):
							SCREAMING_SNAKE_CASE_     						=      tf.ones((4, 1)    ,    dtype=tf.intaa						) * 2
							SCREAMING_SNAKE_CASE_     						=      tf.concat([ids_tensor((4, 6)    ,    self.vocab_size - 3						) + 3, eos_column_vector]    ,    axis=1						)
							SCREAMING_SNAKE_CASE_     						=      input_ids.shape[0]
							SCREAMING_SNAKE_CASE_     						=      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 	lowerCamelCase_							(				unittest.TestCase ):


	'''simple docstring'''







	@slow
	def 				lowerCAmelCase_  (							self  :		Any						):
							SCREAMING_SNAKE_CASE_     						=      TFOPTModel.from_pretrained('facebook/opt-350m'						)
							SCREAMING_SNAKE_CASE_     						=      _long_tensor([[0, 31_414, 232, 328, 740, 1_140, 12_695, 69, 46_078, 1_588, 2]]						)
							SCREAMING_SNAKE_CASE_     						=      tf.not_equal(_lowerCAmelCase    ,    model.config.pad_token_id						)
							with tf.GradientTape():
													SCREAMING_SNAKE_CASE_     						=      model(input_ids=_lowerCAmelCase    ,    attention_mask=_lowerCAmelCase						).last_hidden_state
							SCREAMING_SNAKE_CASE_     						=      (1, 11, 512)
							self.assertEqual(output.shape    ,    _lowerCAmelCase						)
							SCREAMING_SNAKE_CASE_     						=      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]    ,    _lowerCAmelCase    ,    atol=4E-3						)						)

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





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


	'''simple docstring'''







	def 				lowerCAmelCase_  (							self  :		List[str]						):
							super().setUp()
							SCREAMING_SNAKE_CASE_     						=      'facebook/opt-350m'




	def 				lowerCAmelCase_  (							self  :		Any						):
							SCREAMING_SNAKE_CASE_     						=      TFOPTForCausalLM.from_pretrained(self.path_model						)
							SCREAMING_SNAKE_CASE_     						=      GPTaTokenizer.from_pretrained(self.path_model						)

							SCREAMING_SNAKE_CASE_     						=      [
							    '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
							SCREAMING_SNAKE_CASE_     						=      tokenizer(_lowerCAmelCase    ,    return_tensors='tf'    ,    padding=_lowerCAmelCase    ,    add_special_tokens=_lowerCAmelCase						)
							SCREAMING_SNAKE_CASE_     						=      tf.math.reduce_mean(model(inputs.input_ids    ,    attention_mask=inputs.attention_mask						)[0]    ,    axis=-1						)
							SCREAMING_SNAKE_CASE_     						=      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(_lowerCAmelCase    ,    _lowerCAmelCase    ,    atol=1E-4						)						)

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





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


	'''simple docstring'''







	@property
	def 				lowerCAmelCase_  (							self  :		Tuple						):
							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 				lowerCAmelCase_  (							self  :		Dict						):
							SCREAMING_SNAKE_CASE_     						=      'facebook/opt-125m'

							SCREAMING_SNAKE_CASE_     						=      [
							    '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',
							]

							SCREAMING_SNAKE_CASE_     						=      []
							SCREAMING_SNAKE_CASE_     						=      GPTaTokenizer.from_pretrained(_lowerCAmelCase						)
							SCREAMING_SNAKE_CASE_     						=      TFOPTForCausalLM.from_pretrained(_lowerCAmelCase						)

							for prompt in self.prompts:
													SCREAMING_SNAKE_CASE_     						=      tokenizer(_lowerCAmelCase    ,    return_tensors='tf'						).input_ids

													SCREAMING_SNAKE_CASE_     						=      model.generate(_lowerCAmelCase    ,    max_length=10						)

													SCREAMING_SNAKE_CASE_     						=      tokenizer.batch_decode(_lowerCAmelCase    ,    skip_special_tokens=_lowerCAmelCase						)
													predicted_outputs += generated_string

							self.assertListEqual(_lowerCAmelCase    ,    _lowerCAmelCase						)



	def 				lowerCAmelCase_  (							self  :		List[Any]						):
							SCREAMING_SNAKE_CASE_     						=      'facebook/opt-350m'

							SCREAMING_SNAKE_CASE_     						=      GPTaTokenizer.from_pretrained(_lowerCAmelCase						)
							SCREAMING_SNAKE_CASE_     						=      TFOPTForCausalLM.from_pretrained(_lowerCAmelCase						)

							SCREAMING_SNAKE_CASE_     						=      'left'

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

							SCREAMING_SNAKE_CASE_     						=      tokenizer(_lowerCAmelCase    ,    return_tensors='tf'    ,    padding=_lowerCAmelCase						)
							SCREAMING_SNAKE_CASE_     						=      inputs['input_ids']

							SCREAMING_SNAKE_CASE_     						=      model.generate(input_ids=_lowerCAmelCase    ,    attention_mask=inputs['attention_mask']						)

							SCREAMING_SNAKE_CASE_     						=      tokenizer(sentences[0]    ,    return_tensors='tf'						).input_ids
							SCREAMING_SNAKE_CASE_     						=      model.generate(input_ids=_lowerCAmelCase						)

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

							SCREAMING_SNAKE_CASE_     						=      tokenizer.batch_decode(_lowerCAmelCase    ,    skip_special_tokens=_lowerCAmelCase						)
							SCREAMING_SNAKE_CASE_     						=      tokenizer.decode(output_non_padded[0]    ,    skip_special_tokens=_lowerCAmelCase						)
							SCREAMING_SNAKE_CASE_     						=      tokenizer.decode(output_padded[0]    ,    skip_special_tokens=_lowerCAmelCase						)

							SCREAMING_SNAKE_CASE_     						=      [
							    '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]						)




	def 				lowerCAmelCase_  (							self  :		Tuple						):
							SCREAMING_SNAKE_CASE_     						=      'facebook/opt-350m'

							SCREAMING_SNAKE_CASE_     						=      [
							    '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',
							]

							SCREAMING_SNAKE_CASE_     						=      []
							SCREAMING_SNAKE_CASE_     						=      GPTaTokenizer.from_pretrained(_lowerCAmelCase						)
							SCREAMING_SNAKE_CASE_     						=      TFOPTForCausalLM.from_pretrained(_lowerCAmelCase						)

							for prompt in self.prompts:
													SCREAMING_SNAKE_CASE_     						=      tokenizer(_lowerCAmelCase    ,    return_tensors='tf'						).input_ids

													SCREAMING_SNAKE_CASE_     						=      model.generate(_lowerCAmelCase    ,    max_length=10						)

													SCREAMING_SNAKE_CASE_     						=      tokenizer.batch_decode(_lowerCAmelCase    ,    skip_special_tokens=_lowerCAmelCase						)
													predicted_outputs += generated_string

							self.assertListEqual(_lowerCAmelCase    ,    _lowerCAmelCase						)





from abc import ABC, abstractmethod
from argparse import ArgumentParser





class 	lowerCamelCase_							(				_SCREAMING_SNAKE_CASE ):


 '''simple docstring'''







 @staticmethod
 @abstractmethod
 def 				lowerCAmelCase_  (							_lowerCAmelCase  :		ArgumentParser						):
       raise NotImplementedError()




 @abstractmethod
 def 				lowerCAmelCase_  (							self  :		Dict						):
       raise NotImplementedError()