
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			(	a_		,   a_		,   a_=None		,   a_=None						)		->	Any:







				"""simple docstring"""



				if attention_mask is None:
								__A   			= 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'''





	snake_case_					    =     OPTConfig
	snake_case_					    =     {}
	snake_case_					    =     "gelu"






	def __init__(       self					: Tuple							,A					: Tuple							,A					: Optional[Any]=13							,A					: Union[str, Any]=7							,A					: Tuple=True							,A					: Optional[int]=False							,A					: List[str]=99							,A					: str=16							,A					: Optional[Any]=2							,A					: List[str]=4							,A					: Optional[int]=4							,A					: str="gelu"							,A					: Any=0.1							,A					: int=0.1							,A					: int=20							,A					: Tuple=2							,A					: Optional[int]=1							,A					: Union[str, Any]=0							,A					: str=16							,A					: Dict=16							,):
					__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   			= embed_dim
					__A   			= word_embed_proj_dim
					__A   			= False






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

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







	def 				UpperCamelCase_  (       self					: Dict							,A					: int							,A					: Optional[int] ):
					__A   			= TFOPTModel(config=A )
					__A   			= inputs_dict["input_ids"]

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

					# first forward pass
					__A   			= model(A							,attention_mask=A							,use_cache=A )

					__A			,							__A   			= outputs.to_tuple()

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

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

					__A   			= model(A							,attention_mask=A )[0]
					__A   			= model(A							,attention_mask=A							,past_key_values=A )[0]

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

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

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





@require_tf
class 							UpperCAmelCase      (  __SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			unittest.TestCase  ):




	'''simple docstring'''





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






	def 				UpperCamelCase_  (       self					: Union[str, Any] ):
					__A   			= TFOPTModelTester(self )
					__A   			= ConfigTester(self							,config_class=A )






	def 				UpperCamelCase_  (       self					: Dict ):
					self.config_tester.run_common_tests()






	def 				UpperCamelCase_  (       self					: Optional[int] ):
					__A   			= self.model_tester.prepare_config_and_inputs_for_common()
					self.model_tester.check_decoder_model_past_large_inputs(*A )







	def 				UpperCamelCase_  (       self					: Any ):
					__A			,							__A   			= self.model_tester.prepare_config_and_inputs_for_common()

					def _get_word_embedding_weight(A					: Optional[int]							,A					: Optional[Any] ):
									if hasattr(A							,"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(A							,"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
													__A   			= model_class(config=A )
													__A   			= _get_word_embedding_weight(A							,model.get_input_embeddings() )
													__A   			= _get_word_embedding_weight(A							,model.get_output_embeddings() )

													# reshape the embeddings
													model.resize_token_embeddings(A )
													__A   			= _get_word_embedding_weight(A							,model.get_input_embeddings() )
													__A   			= _get_word_embedding_weight(A							,model.get_output_embeddings() )

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

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

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

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

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





def 						UpperCAmelCase			(	a_						)		->	List[str]:







				"""simple docstring"""



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





@require_tf
class 							UpperCAmelCase      (  unittest.TestCase  ):




	'''simple docstring'''





	snake_case_					    =     99






	def 				UpperCamelCase_  (       self					: Tuple ):
					__A   			= tf.ones((4, 1)							,dtype=tf.intaa ) * 2
					__A   			= tf.concat([ids_tensor((4, 6)							,self.vocab_size - 3 ) + 3, eos_column_vector]							,axis=1 )
					__A   			= input_ids.shape[0]
					__A   			= 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 				UpperCamelCase_  (       self					: List[str] ):
					__A   			= TFOPTModel.from_pretrained("facebook/opt-350m" )
					__A   			= _long_tensor([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]] )
					__A   			= tf.not_equal(A							,model.config.pad_token_id )
					with tf.GradientTape():
									__A   			= model(input_ids=A							,attention_mask=A ).last_hidden_state
					__A   			= (1, 11, 5_12)
					self.assertEqual(output.shape							,A )
					__A   			= tf.constant(
					    [[-0.28_73, -1.92_18, -0.30_33], [-1.27_10, -0.13_38, -0.19_02], [0.40_95, 0.12_14, -1.31_21]] )
					self.assertTrue(np.allclose(output[:, :3, :3]							,A							,atol=4E-3 ) )

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





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




	'''simple docstring'''





	def 				UpperCamelCase_  (       self					: Optional[Any] ):
					super().setUp()
					__A   			= "facebook/opt-350m"







	def 				UpperCamelCase_  (       self					: Optional[int] ):
					__A   			= TFOPTForCausalLM.from_pretrained(self.path_model )
					__A   			= GPTaTokenizer.from_pretrained(self.path_model )

					__A   			= [
					    "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
					__A   			= tokenizer(A							,return_tensors="tf"							,padding=A							,add_special_tokens=A )
					__A   			= tf.math.reduce_mean(model(inputs.input_ids							,attention_mask=inputs.attention_mask )[0]							,axis=-1 )
					__A   			= tf.constant(
					    [
					        [1.38_51, -13.89_23, -10.52_29, -10.75_33, -0.23_09, -10.23_84, -0.53_65, -9.09_47, -5.16_70],
					        [-4.70_73, -10.62_76, -3.94_15, -21.52_42, -0.28_22, -0.28_22, -0.28_22, -0.28_22, -0.28_22],
					        [0.62_47, -3.42_29, -8.91_79, -1.42_97, -14.16_50, 1.41_46, -9.02_18, -0.27_03, -0.27_03],
					        [6.47_83, -1.99_13, -10.79_26, -2.33_36, 1.50_92, -0.99_74, -6.82_13, 1.34_77, 1.34_77],
					    ] )
					self.assertTrue(np.allclose(A							,A							,atol=1E-4 ) )

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





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




	'''simple docstring'''





	@property
	def 				UpperCamelCase_  (       self					: int ):
					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					: Optional[int] ):
					__A   			= "facebook/opt-125m"

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

					__A   			= []
					__A   			= GPTaTokenizer.from_pretrained(A )
					__A   			= TFOPTForCausalLM.from_pretrained(A )

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

									__A   			= model.generate(A							,max_length=10 )

									__A   			= tokenizer.batch_decode(A							,skip_special_tokens=A )
									predicted_outputs += generated_string

					self.assertListEqual(A							,A )






	def 				UpperCamelCase_  (       self					: int ):
					__A   			= "facebook/opt-350m"

					__A   			= GPTaTokenizer.from_pretrained(A )
					__A   			= TFOPTForCausalLM.from_pretrained(A )

					__A   			= "left"

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

					__A   			= tokenizer(A							,return_tensors="tf"							,padding=A )
					__A   			= inputs["input_ids"]

					__A   			= model.generate(input_ids=A							,attention_mask=inputs["attention_mask"] )

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

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

					__A   			= tokenizer.batch_decode(A							,skip_special_tokens=A )
					__A   			= tokenizer.decode(output_non_padded[0]							,skip_special_tokens=A )
					__A   			= tokenizer.decode(output_padded[0]							,skip_special_tokens=A )

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







	def 				UpperCamelCase_  (       self					: Dict ):
					__A   			= "facebook/opt-350m"

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

					__A   			= []
					__A   			= GPTaTokenizer.from_pretrained(A )
					__A   			= TFOPTForCausalLM.from_pretrained(A )

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

									__A   			= model.generate(A							,max_length=10 )

									__A   			= tokenizer.batch_decode(A							,skip_special_tokens=A )
									predicted_outputs += generated_string

					self.assertListEqual(A							,A )



import argparse

import torch

from transformers import OpenAIGPTConfig, OpenAIGPTModel, load_tf_weights_in_openai_gpt
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()






def        UpperCamelCase_(							_snake_case   :  Dict       ,     _snake_case   :  Optional[int]       ,     _snake_case   :  str	):




			"""simple docstring"""






			if openai_config_file == "":
						__a        =OpenAIGPTConfig()
			else:
						__a        =OpenAIGPTConfig.from_json_file(_snake_case	)
			__a        =OpenAIGPTModel(_snake_case	)

			# Load weights from numpy
			load_tf_weights_in_openai_gpt(_snake_case       ,     _snake_case       ,     _snake_case	)

			# Save pytorch-model
			__a        =pytorch_dump_folder_path + '/' + WEIGHTS_NAME
			__a        =pytorch_dump_folder_path + '/' + CONFIG_NAME
			print(F'Save PyTorch model to {pytorch_weights_dump_path}'	)
			torch.save(model.state_dict()       ,     _snake_case	)
			print(F'Save configuration file to {pytorch_config_dump_path}'	)
			with open(_snake_case       ,     'w'       ,     encoding='utf-8'	) as f:
						f.write(config.to_json_string()	)


if __name__ == "__main__":
						_lowerCAmelCase					:			Tuple  				=			argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    "--openai_checkpoint_folder_path",
						    default=None,
						    type=str,
						    required=True,
						    help="Path to the TensorFlow checkpoint path.",
						)
						parser.add_argument(
						    "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
						)
						parser.add_argument(
						    "--openai_config_file",
						    default="",
						    type=str,
						    help=(
						        "An optional config json file corresponding to the pre-trained OpenAI model. \n"
						        "This specifies the model architecture."
						    ),
						)
						_lowerCAmelCase					:			int  				=			parser.parse_args()
						convert_openai_checkpoint_to_pytorch(
						    args.openai_checkpoint_folder_path, args.openai_config_file, args.pytorch_dump_folder_path
						)






