





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__  (      __a				,	__a				,	__a=None				,	__a=None				):
			if attention_mask is None:
						snake_case_						:     List[Any]     =		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 SCREAMING_SNAKE_CASE_  :
						__magic_name__:     List[str]	 =    OPTConfig
						__magic_name__:     Optional[Any]	 =    {}
						__magic_name__:     Union[str, Any]	 =    "gelu"






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







									"""simple docstring"""


									snake_case_						:     int     =		parent
									snake_case_						:     List[str]     =		batch_size
									snake_case_						:     str     =		seq_length
									snake_case_						:     Any     =		is_training
									snake_case_						:     List[str]     =		use_labels
									snake_case_						:     Any     =		vocab_size
									snake_case_						:     str     =		hidden_size
									snake_case_						:     Union[str, Any]     =		num_hidden_layers
									snake_case_						:     Optional[int]     =		num_attention_heads
									snake_case_						:     Tuple     =		intermediate_size
									snake_case_						:     str     =		hidden_act
									snake_case_						:     Optional[Any]     =		hidden_dropout_prob
									snake_case_						:     Tuple     =		attention_probs_dropout_prob
									snake_case_						:     Any     =		max_position_embeddings
									snake_case_						:     Optional[int]     =		eos_token_id
									snake_case_						:     List[Any]     =		pad_token_id
									snake_case_						:     int     =		bos_token_id
									snake_case_						:     List[Any]     =		embed_dim
									snake_case_						:     int     =		word_embed_proj_dim
									snake_case_						:     List[str]     =		False






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







									"""simple docstring"""


									snake_case_						:     int     =		ids_tensor([self.batch_size, self.seq_length - 1]						,				self.vocab_size		)
									snake_case_						:     Optional[Any]     =		tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size		)						,				1		)
									snake_case_						:     List[str]     =		tf.concat([input_ids, eos_tensor]						,				axis=1		)

									snake_case_						:     Optional[Any]     =		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						,				)
									snake_case_						:     List[str]     =		prepare_opt_inputs_dict(_A						,				_A		)
									return config, inputs_dict






						def UpperCAmelCase_			( self : str						,				_A : List[Any]						,				_A : List[str]		)							->      Optional[Any]:







									"""simple docstring"""


									snake_case_						:     str     =		TFOPTModel(config=_A		)
									snake_case_						:     Any     =		inputs_dict['input_ids']

									snake_case_						:     Tuple     =		input_ids[:1, :]
									snake_case_						:     List[Any]     =		inputs_dict['attention_mask'][:1, :]
									snake_case_						:     str     =		1

									# first forward pass
									snake_case_						:     Optional[int]     =		model(_A						,				attention_mask=_A						,				use_cache=_A		)

									snake_case_						:     int     =		outputs.to_tuple()

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

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

									snake_case_						:     Dict     =		model(_A						,				attention_mask=_A		)[0]
									snake_case_						:     Dict     =		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
									snake_case_						:     str     =		int(ids_tensor((1,)						,				output_from_past.shape[-1]		)		)
									snake_case_						:     int     =		output_from_no_past[:, -3:, random_slice_idx]
									snake_case_						:     List[str]     =		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 SCREAMING_SNAKE_CASE_  (   snake_case_    ,			snake_case_    ,			unittest.TestCase   ):
						__magic_name__:     List[str]	 =    (TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
						__magic_name__:     Optional[Any]	 =    (TFOPTForCausalLM,) if is_tf_available() else ()
						__magic_name__:     List[str]	 =    (
						    {"feature-extraction": TFOPTModel, "text-generation": TFOPTForCausalLM} if is_tf_available() else {}
						)
						__magic_name__:     List[str]	 =    False
						__magic_name__:     List[Any]	 =    False
						__magic_name__:     List[Any]	 =    False
						__magic_name__:     Optional[Any]	 =    10






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







									"""simple docstring"""


									snake_case_						:     int     =		TFOPTModelTester(self		)
									snake_case_						:     Any     =		ConfigTester(self						,				config_class=_A		)






						def UpperCAmelCase_			( self : str		)							->      Tuple:







									"""simple docstring"""


									self.config_tester.run_common_tests()






						def UpperCAmelCase_			( self : Dict		)							->      List[Any]:







									"""simple docstring"""


									snake_case_						:     List[Any]     =		self.model_tester.prepare_config_and_inputs_for_common()
									self.model_tester.check_decoder_model_past_large_inputs(*_A		)






						def UpperCAmelCase_			( self : Tuple		)							->      List[str]:







									"""simple docstring"""


									snake_case_						:     str     =		self.model_tester.prepare_config_and_inputs_for_common()

									def _get_word_embedding_weight(_A : List[str]						,				_A : List[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
															snake_case_						:     Union[str, Any]     =		model_class(config=_A		)
															snake_case_						:     Tuple     =		_get_word_embedding_weight(_A						,				model.get_input_embeddings()		)
															snake_case_						:     Optional[int]     =		_get_word_embedding_weight(_A						,				model.get_output_embeddings()		)

															# reshape the embeddings
															model.resize_token_embeddings(_A		)
															snake_case_						:     Optional[int]     =		_get_word_embedding_weight(_A						,				model.get_input_embeddings()		)
															snake_case_						:     int     =		_get_word_embedding_weight(_A						,				model.get_output_embeddings()		)

															# check that the resized embeddings size matches the desired size.
															snake_case_						:     List[Any]     =		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
															snake_case_						:     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:
																					snake_case_						:     Optional[Any]     =		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		)

																		snake_case_						:     Optional[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:
																								snake_case_						:     List[str]     =		False
																		self.assertTrue(_A		)


def 							SCREAMING_SNAKE_CASE__  (      __a				):
			return tf.constant(__a				,	dtype=tf.intaa				)
@require_tf
class SCREAMING_SNAKE_CASE_  (   unittest.TestCase   ):
						__magic_name__:     List[str]	 =    99






						def UpperCAmelCase_			( self : List[str]		)							->      int:







									"""simple docstring"""


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






						@slow
						def UpperCAmelCase_			( self : List[str]		)							->      str:







									"""simple docstring"""


									snake_case_						:     Any     =		TFOPTModel.from_pretrained('facebook/opt-350m'		)
									snake_case_						:     int     =		_long_tensor([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]		)
									snake_case_						:     Optional[Any]     =		tf.not_equal(_A						,				model.config.pad_token_id		)
									with tf.GradientTape():
												snake_case_						:     List[Any]     =		model(input_ids=_A						,				attention_mask=_A		).last_hidden_state
									snake_case_						:     int     =		(1, 11, 512)
									self.assertEqual(output.shape						,				_A		)
									snake_case_						:     Optional[int]     =		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]						,				_A						,				atol=4E-3		)		)

									snake_case_						:     Dict     =		tf.function(_A						,				jit_compile=_A		)
									snake_case_						:     Dict     =		xla_generate(_A						,				_A		)[0]
									self.assertTrue(np.allclose(output[:, :3, :3]						,				_A						,				atol=4E-2		)		)
@require_tf
@slow
class SCREAMING_SNAKE_CASE_  (   unittest.TestCase   ):






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







									"""simple docstring"""


									super().setUp()
									snake_case_						:     str     =		'facebook/opt-350m'






						def UpperCAmelCase_			( self : Optional[Any]		)							->      Union[str, Any]:







									"""simple docstring"""


									snake_case_						:     Any     =		TFOPTForCausalLM.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_						:     int     =		tokenizer(_A						,				return_tensors='tf'						,				padding=_A						,				add_special_tokens=_A		)
									snake_case_						:     List[str]     =		tf.math.reduce_mean(model(inputs.input_ids						,				attention_mask=inputs.attention_mask		)[0]						,				axis=-1		)
									snake_case_						:     Optional[Any]     =		tf.constant(
									    [
									        [1.3_8_5_1, -13.8923, -10.5229, -10.7533, -0.2_3_0_9, -10.2384, -0.5_3_6_5, -9.0_9_4_7, -5.1_6_7_0],
									        [-4.7_0_7_3, -10.6276, -3.9_4_1_5, -21.5242, -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, -14.1650, 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, -10.7926, -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(_A						,				_A						,				atol=1E-4		)		)

									snake_case_						:     Dict     =		tf.function(_A						,				jit_compile=_A		)
									snake_case_						:     Union[str, Any]     =		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 SCREAMING_SNAKE_CASE_  (   unittest.TestCase   ):






						@property
						def UpperCAmelCase_			( self : Optional[int]		)							->      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 : Any		)							->      Optional[int]:







									"""simple docstring"""


									snake_case_						:     Optional[int]     =		'facebook/opt-125m'

									snake_case_						:     List[str]     =		[
									    '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_						:     Dict     =		[]
									snake_case_						:     str     =		GPTaTokenizer.from_pretrained(_A		)
									snake_case_						:     Union[str, Any]     =		TFOPTForCausalLM.from_pretrained(_A		)

									for prompt in self.prompts:
												snake_case_						:     Optional[int]     =		tokenizer(_A						,				return_tensors='tf'		).input_ids

												snake_case_						:     Any     =		model.generate(_A						,				max_length=10		)

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

									self.assertListEqual(_A						,				_A		)






						def UpperCAmelCase_			( self : List[Any]		)							->      Dict:







									"""simple docstring"""


									snake_case_						:     Dict     =		'facebook/opt-350m'

									snake_case_						:     Tuple     =		GPTaTokenizer.from_pretrained(_A		)
									snake_case_						:     Optional[Any]     =		TFOPTForCausalLM.from_pretrained(_A		)

									snake_case_						:     Optional[int]     =		'left'

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

									snake_case_						:     Tuple     =		tokenizer(_A						,				return_tensors='tf'						,				padding=_A		)
									snake_case_						:     Tuple     =		inputs['input_ids']

									snake_case_						:     int     =		model.generate(input_ids=_A						,				attention_mask=inputs['attention_mask']		)

									snake_case_						:     int     =		tokenizer(sentences[0]						,				return_tensors='tf'		).input_ids
									snake_case_						:     Optional[int]     =		model.generate(input_ids=_A		)

									snake_case_						:     Tuple     =		inputs_non_padded.shape[-1] - tf.math.reduce_sum(
									    tf.cast(inputs['attention_mask'][-1]						,				tf.intaa		)		)
									snake_case_						:     str     =		tokenizer(sentences[1]						,				return_tensors='tf'		).input_ids
									snake_case_						:     Optional[int]     =		model.generate(input_ids=_A						,				max_length=model.config.max_length - num_paddings		)

									snake_case_						:     int     =		tokenizer.batch_decode(_A						,				skip_special_tokens=_A		)
									snake_case_						:     List[str]     =		tokenizer.decode(output_non_padded[0]						,				skip_special_tokens=_A		)
									snake_case_						:     Optional[int]     =		tokenizer.decode(output_padded[0]						,				skip_special_tokens=_A		)

									snake_case_						:     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(_A						,				_A		)
									self.assertListEqual(_A						,				[non_padded_sentence, padded_sentence]		)






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







									"""simple docstring"""


									snake_case_						:     int     =		'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_						:     List[str]     =		[]
									snake_case_						:     Any     =		GPTaTokenizer.from_pretrained(_A		)
									snake_case_						:     Optional[int]     =		TFOPTForCausalLM.from_pretrained(_A		)

									for prompt in self.prompts:
												snake_case_						:     int     =		tokenizer(_A						,				return_tensors='tf'		).input_ids

												snake_case_						:     Tuple     =		model.generate(_A						,				max_length=10		)

												snake_case_						:     Optional[int]     =		tokenizer.batch_decode(_A						,				skip_special_tokens=_A		)
												predicted_outputs += generated_string

									self.assertListEqual(_A						,				_A		)








import argparse

from transformers import (
    TapasConfig,
    TapasForMaskedLM,
    TapasForQuestionAnswering,
    TapasForSequenceClassification,
    TapasModel,
    TapasTokenizer,
    load_tf_weights_in_tapas,
)
from transformers.utils import logging


logging.set_verbosity_info()


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				,	__a				,	__a				,	__a				):
	# Initialise PyTorch model.
	# If you want to convert a checkpoint that uses absolute position embeddings, make sure to set reset_position_index_per_cell of
	# TapasConfig to False.

	# initialize configuration from json file
	snake_case_						:     Optional[int]     =		TapasConfig.from_json_file(__a				)
	# set absolute/relative position embeddings parameter
	snake_case_						:     List[Any]     =		reset_position_index_per_cell

	# set remaining parameters of TapasConfig as well as the model based on the task
	if task == "SQA":
		snake_case_						:     int     =		TapasForQuestionAnswering(config=__a				)
	elif task == "WTQ":
		# run_task_main.py hparams
		snake_case_						:     Optional[int]     =		4
		snake_case_						:     List[str]     =		True
		# hparam_utils.py hparams
		snake_case_						:     Optional[Any]     =		0.664694
		snake_case_						:     Dict     =		0.207951
		snake_case_						:     Tuple     =		0.121194
		snake_case_						:     Dict     =		True
		snake_case_						:     int     =		True
		snake_case_						:     int     =		False
		snake_case_						:     str     =		0.0352513

		snake_case_						:     int     =		TapasForQuestionAnswering(config=__a				)
	elif task == "WIKISQL_SUPERVISED":
		# run_task_main.py hparams
		snake_case_						:     int     =		4
		snake_case_						:     Optional[int]     =		False
		# hparam_utils.py hparams
		snake_case_						:     str     =		36.4519
		snake_case_						:     Optional[Any]     =		0.903421
		snake_case_						:     List[Any]     =		222.088
		snake_case_						:     Optional[int]     =		True
		snake_case_						:     Optional[Any]     =		True
		snake_case_						:     str     =		True
		snake_case_						:     int     =		0.763141

		snake_case_						:     str     =		TapasForQuestionAnswering(config=__a				)
	elif task == "TABFACT":
		snake_case_						:     List[Any]     =		TapasForSequenceClassification(config=__a				)
	elif task == "MLM":
		snake_case_						:     Optional[int]     =		TapasForMaskedLM(config=__a				)
	elif task == "INTERMEDIATE_PRETRAINING":
		snake_case_						:     Tuple     =		TapasModel(config=__a				)
	else:
		raise ValueError(f"""Task {task} not supported."""				)

	print(f"""Building PyTorch model from configuration: {config}"""				)
	# Load weights from tf checkpoint
	load_tf_weights_in_tapas(__a				,	__a				,	__a				)

	# Save pytorch-model (weights and configuration)
	print(f"""Save PyTorch model to {pytorch_dump_path}"""				)
	model.save_pretrained(__a				)

	# Save tokenizer files
	print(f"""Save tokenizer files to {pytorch_dump_path}"""				)
	snake_case_						:     str     =		TapasTokenizer(vocab_file=tf_checkpoint_path[:-10] + 'vocab.txt'				,	model_max_length=5_12				)
	tokenizer.save_pretrained(__a				)

	print('Used relative position embeddings:'				,	model.config.reset_position_index_per_cell				)


if __name__ == "__main__":
						_SCREAMING_SNAKE_CASE 		=						argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    """--task""", default="""SQA""", type=str, help="""Model task for which to convert a checkpoint. Defaults to SQA."""
						)
						parser.add_argument(
						    """--reset_position_index_per_cell""",
						    default=False,
						    action="""store_true""",
						    help="""Whether to use relative position embeddings or not. Defaults to True.""",
						)
						parser.add_argument(
						    """--tf_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
						)
						parser.add_argument(
						    """--tapas_config_file""",
						    default=None,
						    type=str,
						    required=True,
						    help=(
						        """The config json file corresponding to the pre-trained TAPAS model. \n"""
						        """This specifies the model architecture."""
						    ),
						)
						parser.add_argument(
						    """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
						)
						_SCREAMING_SNAKE_CASE 		=						parser.parse_args()
						convert_tf_checkpoint_to_pytorch(
						    args.task,
						    args.reset_position_index_per_cell,
						    args.tf_checkpoint_path,
						    args.tapas_config_file,
						    args.pytorch_dump_path,
						)

