


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     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_      ,       lowercase_=None      ,       lowercase_=None      )		->     Optional[Any]:




						'''simple docstring'''







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




@require_tf
class 							lowerCamelCase   :
	_lowerCAmelCase     :		Union[str, Any]										=						OPTConfig
	_lowerCAmelCase     :		List[Any]										=						{}
	_lowerCAmelCase     :		int										=						'''gelu'''


	def __init__(    self				,     lowercase__				,     lowercase__=1_3				,     lowercase__=7				,     lowercase__=True				,     lowercase__=False				,     lowercase__=9_9				,     lowercase__=1_6				,     lowercase__=2				,     lowercase__=4				,     lowercase__=4				,     lowercase__="gelu"				,     lowercase__=0.1				,     lowercase__=0.1				,     lowercase__=2_0				,     lowercase__=2				,     lowercase__=1				,     lowercase__=0				,     lowercase__=1_6				,     lowercase__=1_6				,     ):
							__UpperCAmelCase			:							Any	    =	parent
							__UpperCAmelCase			:							Any	    =	batch_size
							__UpperCAmelCase			:							str	    =	seq_length
							__UpperCAmelCase			:							Dict	    =	is_training
							__UpperCAmelCase			:							Optional[Any]	    =	use_labels
							__UpperCAmelCase			:							Union[str, Any]	    =	vocab_size
							__UpperCAmelCase			:							Optional[int]	    =	hidden_size
							__UpperCAmelCase			:							Optional[Any]	    =	num_hidden_layers
							__UpperCAmelCase			:							Tuple	    =	num_attention_heads
							__UpperCAmelCase			:							Optional[Any]	    =	intermediate_size
							__UpperCAmelCase			:							Dict	    =	hidden_act
							__UpperCAmelCase			:							Any	    =	hidden_dropout_prob
							__UpperCAmelCase			:							Tuple	    =	attention_probs_dropout_prob
							__UpperCAmelCase			:							Optional[int]	    =	max_position_embeddings
							__UpperCAmelCase			:							Tuple	    =	eos_token_id
							__UpperCAmelCase			:							Tuple	    =	pad_token_id
							__UpperCAmelCase			:							Optional[int]	    =	bos_token_id
							__UpperCAmelCase			:							Optional[int]	    =	embed_dim
							__UpperCAmelCase			:							List[Any]	    =	word_embed_proj_dim
							__UpperCAmelCase			:							int	    =	False


	def  A(    self):
							__UpperCAmelCase			:							Union[str, Any]	    =	ids_tensor([self.batch_size, self.seq_length - 1]				,     self.vocab_size)
							__UpperCAmelCase			:							Dict	    =	tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size)				,     1)
							__UpperCAmelCase			:							List[Any]	    =	tf.concat([input_ids, eos_tensor]				,     axis=1)

							__UpperCAmelCase			:							str	    =	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=lowercase__				,     **self.config_updates				,     )
							__UpperCAmelCase			:							Dict	    =	prepare_opt_inputs_dict(lowercase__				,     lowercase__)
							return config, inputs_dict


	def  A(    self				,     lowercase__				,     lowercase__):
							__UpperCAmelCase			:							List[str]	    =	TFOPTModel(config=lowercase__)
							__UpperCAmelCase			:							List[str]	    =	inputs_dict['''input_ids''']

							__UpperCAmelCase			:							List[Any]	    =	input_ids[:1, :]
							__UpperCAmelCase			:							Dict	    =	inputs_dict['''attention_mask'''][:1, :]
							__UpperCAmelCase			:							Optional[Any]	    =	1

							# first forward pass
							__UpperCAmelCase			:							List[Any]	    =	model(lowercase__				,     attention_mask=lowercase__				,     use_cache=lowercase__)

							__UpperCAmelCase    ,      __UpperCAmelCase			:							int	    =	outputs.to_tuple()

							# create hypothetical next token and extent to next_input_ids
							__UpperCAmelCase			:							Tuple	    =	ids_tensor((self.batch_size, 3)				,     config.vocab_size)
							__UpperCAmelCase			:							List[Any]	    =	tf.cast(ids_tensor((self.batch_size, 3)				,     2)				,     tf.inta)

							# append to next input_ids and
							__UpperCAmelCase			:							List[str]	    =	tf.concat([input_ids, next_tokens]				,     axis=-1)
							__UpperCAmelCase			:							Optional[Any]	    =	tf.concat([attention_mask, next_attn_mask]				,     axis=-1)

							__UpperCAmelCase			:							Any	    =	model(lowercase__				,     attention_mask=lowercase__)[0]
							__UpperCAmelCase			:							Tuple	    =	model(lowercase__				,     attention_mask=lowercase__				,     past_key_values=lowercase__)[0]

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

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

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






@require_tf
class 							lowerCamelCase   (  _UpperCamelCase							,						_UpperCamelCase							,						unittest.TestCase ):
	_lowerCAmelCase     :		List[Any]										=						(TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
	_lowerCAmelCase     :		List[Any]										=						(TFOPTForCausalLM,) if is_tf_available() else ()
	_lowerCAmelCase     :		Any										=						(
	    {'''feature-extraction''': TFOPTModel, '''text-generation''': TFOPTForCausalLM} if is_tf_available() else {}
	)
	_lowerCAmelCase     :		int										=						False
	_lowerCAmelCase     :		List[Any]										=						False
	_lowerCAmelCase     :		Tuple										=						False
	_lowerCAmelCase     :		Tuple										=						1_0


	def  A(    self):
							__UpperCAmelCase			:							Any	    =	TFOPTModelTester(self)
							__UpperCAmelCase			:							Any	    =	ConfigTester(self				,     config_class=lowercase__)


	def  A(    self):
							self.config_tester.run_common_tests()


	def  A(    self):
							__UpperCAmelCase			:							Union[str, Any]	    =	self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.check_decoder_model_past_large_inputs(*lowercase__)


	def  A(    self):
							__UpperCAmelCase    ,      __UpperCAmelCase			:							Tuple	    =	self.model_tester.prepare_config_and_inputs_for_common()

							def _get_word_embedding_weight(lowercase__				,     lowercase__):
													if hasattr(lowercase__				,     '''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(lowercase__				,     '''weight'''):
																									return embedding_layer.weight
																			else:
																									return None

							for model_class in self.all_model_classes:
													for size in [config.vocab_size - 1_0, config.vocab_size + 1_0]:
																			# build the embeddings
																			__UpperCAmelCase			:							int	    =	model_class(config=lowercase__)
																			__UpperCAmelCase			:							List[Any]	    =	_get_word_embedding_weight(lowercase__				,     model.get_input_embeddings())
																			__UpperCAmelCase			:							Dict	    =	_get_word_embedding_weight(lowercase__				,     model.get_output_embeddings())

																			# reshape the embeddings
																			model.resize_token_embeddings(lowercase__)
																			__UpperCAmelCase			:							Tuple	    =	_get_word_embedding_weight(lowercase__				,     model.get_input_embeddings())
																			__UpperCAmelCase			:							List[str]	    =	_get_word_embedding_weight(lowercase__				,     model.get_output_embeddings())

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

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

																			# check that weights remain the same after resizing
																			__UpperCAmelCase			:							List[str]	    =	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:
																															__UpperCAmelCase			:							int	    =	False
																			self.assertTrue(lowercase__)

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

																									__UpperCAmelCase			:							Union[str, Any]	    =	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:
																																					__UpperCAmelCase			:							Optional[int]	    =	False
																									self.assertTrue(lowercase__)






def     __SCREAMING_SNAKE_CASE							(       lowercase_      )		->     Union[str, Any]:




						'''simple docstring'''







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




@require_tf
class 							lowerCamelCase   (  unittest.TestCase ):
	_lowerCAmelCase     :		List[str]										=						9_9


	def  A(    self):
							__UpperCAmelCase			:							Any	    =	tf.ones((4, 1)				,     dtype=tf.intaa) * 2
							__UpperCAmelCase			:							Optional[Any]	    =	tf.concat([ids_tensor((4, 6)				,     self.vocab_size - 3) + 3, eos_column_vector]				,     axis=1)
							__UpperCAmelCase			:							Union[str, Any]	    =	input_ids.shape[0]
							__UpperCAmelCase			:							Union[str, Any]	    =	OPTConfig(
							    vocab_size=self.vocab_size				,     hidden_size=2_4				,     num_hidden_layers=2				,     num_attention_heads=2				,     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






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


	@slow
	def  A(    self):
							__UpperCAmelCase			:							Dict	    =	TFOPTModel.from_pretrained('''facebook/opt-350m''')
							__UpperCAmelCase			:							Optional[int]	    =	_long_tensor([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]])
							__UpperCAmelCase			:							Union[str, Any]	    =	tf.not_equal(lowercase__				,     model.config.pad_token_id)
							with tf.GradientTape():
													__UpperCAmelCase			:							int	    =	model(input_ids=lowercase__				,     attention_mask=lowercase__).last_hidden_state
							__UpperCAmelCase			:							str	    =	(1, 1_1, 5_1_2)
							self.assertEqual(output.shape				,     lowercase__)
							__UpperCAmelCase			:							Tuple	    =	tf.constant(
							    [[-0.2_8_7_3, -1.9_2_1_8, -0.3_0_3_3], [-1.2_7_1_0, -0.1_3_3_8, -0.1_9_0_2], [0.4_0_9_5, 0.1_2_1_4, -1.3_1_2_1]])
							self.assertTrue(np.allclose(output[:, :3, :3]				,     lowercase__				,     atol=4e-3))

							__UpperCAmelCase			:							str	    =	tf.function(lowercase__				,     jit_compile=lowercase__)
							__UpperCAmelCase			:							Any	    =	xla_generate(lowercase__				,     lowercase__)[0]
							self.assertTrue(np.allclose(output[:, :3, :3]				,     lowercase__				,     atol=4e-2))






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


	def  A(    self):
							super().setUp()
							__UpperCAmelCase			:							Union[str, Any]	    =	'''facebook/opt-350m'''


	def  A(    self):
							__UpperCAmelCase			:							Tuple	    =	TFOPTForCausalLM.from_pretrained(self.path_model)
							__UpperCAmelCase			:							Dict	    =	GPTaTokenizer.from_pretrained(self.path_model)

							__UpperCAmelCase			:							Dict	    =	[
							    '''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
							__UpperCAmelCase			:							List[Any]	    =	tokenizer(lowercase__				,     return_tensors='''tf'''				,     padding=lowercase__				,     add_special_tokens=lowercase__)
							__UpperCAmelCase			:							Any	    =	tf.math.reduce_mean(model(inputs.input_ids				,     attention_mask=inputs.attention_mask)[0]				,     axis=-1)
							__UpperCAmelCase			:							List[Any]	    =	tf.constant(
							    [
							        [1.3_8_5_1, -1_3.8_9_2_3, -1_0.5_2_2_9, -1_0.7_5_3_3, -0.2_3_0_9, -1_0.2_3_8_4, -0.5_3_6_5, -9.0_9_4_7, -5.1_6_7_0],
							        [-4.7_0_7_3, -1_0.6_2_7_6, -3.9_4_1_5, -2_1.5_2_4_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2],
							        [0.6_2_4_7, -3.4_2_2_9, -8.9_1_7_9, -1.4_2_9_7, -1_4.1_6_5_0, 1.4_1_4_6, -9.0_2_1_8, -0.2_7_0_3, -0.2_7_0_3],
							        [6.4_7_8_3, -1.9_9_1_3, -1_0.7_9_2_6, -2.3_3_3_6, 1.5_0_9_2, -0.9_9_7_4, -6.8_2_1_3, 1.3_4_7_7, 1.3_4_7_7],
							    ])
							self.assertTrue(np.allclose(lowercase__				,     lowercase__				,     atol=1e-4))

							__UpperCAmelCase			:							str	    =	tf.function(lowercase__				,     jit_compile=lowercase__)
							__UpperCAmelCase			:							Optional[int]	    =	tf.math.reduce_mean(xla_generate(inputs.input_ids				,     attention_mask=inputs.attention_mask)[0]				,     axis=-1)
							self.assertTrue(np.allclose(lowercase__				,     lowercase__				,     atol=1e-4))






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


	@property
	def  A(    self):
							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  A(    self):
							__UpperCAmelCase			:							Any	    =	'''facebook/opt-125m'''

							__UpperCAmelCase			:							List[Any]	    =	[
							    '''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''',
							]

							__UpperCAmelCase			:							Dict	    =	[]
							__UpperCAmelCase			:							List[str]	    =	GPTaTokenizer.from_pretrained(lowercase__)
							__UpperCAmelCase			:							List[Any]	    =	TFOPTForCausalLM.from_pretrained(lowercase__)

							for prompt in self.prompts:
													__UpperCAmelCase			:							Dict	    =	tokenizer(lowercase__				,     return_tensors='''tf''').input_ids

													__UpperCAmelCase			:							Tuple	    =	model.generate(lowercase__				,     max_length=1_0)

													__UpperCAmelCase			:							List[str]	    =	tokenizer.batch_decode(lowercase__				,     skip_special_tokens=lowercase__)
													predicted_outputs += generated_string

							self.assertListEqual(lowercase__				,     lowercase__)


	def  A(    self):
							__UpperCAmelCase			:							int	    =	'''facebook/opt-350m'''

							__UpperCAmelCase			:							List[str]	    =	GPTaTokenizer.from_pretrained(lowercase__)
							__UpperCAmelCase			:							List[str]	    =	TFOPTForCausalLM.from_pretrained(lowercase__)

							__UpperCAmelCase			:							List[str]	    =	'''left'''

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

							__UpperCAmelCase			:							str	    =	tokenizer(lowercase__				,     return_tensors='''tf'''				,     padding=lowercase__)
							__UpperCAmelCase			:							Union[str, Any]	    =	inputs['''input_ids''']

							__UpperCAmelCase			:							Union[str, Any]	    =	model.generate(input_ids=lowercase__				,     attention_mask=inputs['''attention_mask'''])

							__UpperCAmelCase			:							Optional[int]	    =	tokenizer(sentences[0]				,     return_tensors='''tf''').input_ids
							__UpperCAmelCase			:							Union[str, Any]	    =	model.generate(input_ids=lowercase__)

							__UpperCAmelCase			:							Tuple	    =	inputs_non_padded.shape[-1] - tf.math.reduce_sum(
							    tf.cast(inputs['''attention_mask'''][-1]				,     tf.intaa))
							__UpperCAmelCase			:							Dict	    =	tokenizer(sentences[1]				,     return_tensors='''tf''').input_ids
							__UpperCAmelCase			:							Any	    =	model.generate(input_ids=lowercase__				,     max_length=model.config.max_length - num_paddings)

							__UpperCAmelCase			:							str	    =	tokenizer.batch_decode(lowercase__				,     skip_special_tokens=lowercase__)
							__UpperCAmelCase			:							Optional[int]	    =	tokenizer.decode(output_non_padded[0]				,     skip_special_tokens=lowercase__)
							__UpperCAmelCase			:							Any	    =	tokenizer.decode(output_padded[0]				,     skip_special_tokens=lowercase__)

							__UpperCAmelCase			:							Dict	    =	[
							    '''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(lowercase__				,     lowercase__)
							self.assertListEqual(lowercase__				,     [non_padded_sentence, padded_sentence])


	def  A(    self):
							__UpperCAmelCase			:							Optional[int]	    =	'''facebook/opt-350m'''

							__UpperCAmelCase			:							Union[str, Any]	    =	[
							    '''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''',
							]

							__UpperCAmelCase			:							Optional[int]	    =	[]
							__UpperCAmelCase			:							str	    =	GPTaTokenizer.from_pretrained(lowercase__)
							__UpperCAmelCase			:							List[str]	    =	TFOPTForCausalLM.from_pretrained(lowercase__)

							for prompt in self.prompts:
													__UpperCAmelCase			:							str	    =	tokenizer(lowercase__				,     return_tensors='''tf''').input_ids

													__UpperCAmelCase			:							str	    =	model.generate(lowercase__				,     max_length=1_0)

													__UpperCAmelCase			:							str	    =	tokenizer.batch_decode(lowercase__				,     skip_special_tokens=lowercase__)
													predicted_outputs += generated_string

							self.assertListEqual(lowercase__				,     lowercase__)





from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase       		=    logging.get_logger(__name__)

lowerCAmelCase       		=    {
    """abeja/gpt-neox-japanese-2.7b""": """https://huggingface.co/abeja/gpt-neox-japanese-2.7b/resolve/main/config.json""",
}




class 							lowerCamelCase   (  _UpperCamelCase ):
	_lowerCAmelCase     :		Tuple										=						'''gpt_neox_japanese'''


	def __init__(    self				,     lowercase__=3_2_0_0_0				,     lowercase__=2_5_6_0				,     lowercase__=3_2				,     lowercase__=3_2				,     lowercase__=4				,     lowercase__="gelu"				,     lowercase__=1.0_0				,     lowercase__=1_0_0_0_0				,     lowercase__=2_0_4_8				,     lowercase__=0.0_2				,     lowercase__=1e-5				,     lowercase__=True				,     lowercase__=3_1_9_9_6				,     lowercase__=3_1_9_9_9				,     lowercase__=0.1				,     lowercase__=0.0				,     **lowercase__				,     ):
							super().__init__(bos_token_id=lowercase__				,     eos_token_id=lowercase__				,     **lowercase__)
							__UpperCAmelCase			:							Optional[Any]	    =	vocab_size
							__UpperCAmelCase			:							Dict	    =	max_position_embeddings
							__UpperCAmelCase			:							int	    =	hidden_size
							__UpperCAmelCase			:							List[Any]	    =	num_hidden_layers
							__UpperCAmelCase			:							Union[str, Any]	    =	num_attention_heads
							__UpperCAmelCase			:							Optional[int]	    =	intermediate_multiple_size
							__UpperCAmelCase			:							Any	    =	hidden_act
							__UpperCAmelCase			:							str	    =	rotary_pct
							__UpperCAmelCase			:							Any	    =	rotary_emb_base
							__UpperCAmelCase			:							Dict	    =	initializer_range
							__UpperCAmelCase			:							Tuple	    =	layer_norm_eps
							__UpperCAmelCase			:							Optional[int]	    =	use_cache
							__UpperCAmelCase			:							str	    =	attention_dropout
							__UpperCAmelCase			:							Any	    =	hidden_dropout

