






from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
		import tensorflow as tf

		from transformers import TF_MODEL_FOR_PRETRAINING_MAPPING
		from transformers.models.albert.modeling_tf_albert import (
		    TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFAlbertForMaskedLM,
		    TFAlbertForMultipleChoice,
		    TFAlbertForPreTraining,
		    TFAlbertForQuestionAnswering,
		    TFAlbertForSequenceClassification,
		    TFAlbertForTokenClassification,
		    TFAlbertModel,
		)





class __lowerCamelCase					:





	"""simple docstring"""






	def __init__(				self			,       UpperCAmelCase			,       UpperCAmelCase=13			,       UpperCAmelCase=7			,       UpperCAmelCase=True			,       UpperCAmelCase=True			,       UpperCAmelCase=True			,       UpperCAmelCase=True			,       UpperCAmelCase=99			,       UpperCAmelCase=16			,       UpperCAmelCase=32			,       UpperCAmelCase=2			,       UpperCAmelCase=4			,       UpperCAmelCase=37			,       UpperCAmelCase="gelu"			,       UpperCAmelCase=0.1			,       UpperCAmelCase=0.1			,       UpperCAmelCase=512			,       UpperCAmelCase=16			,       UpperCAmelCase=2			,       UpperCAmelCase=0.02			,       UpperCAmelCase=3			,       UpperCAmelCase=4			,       UpperCAmelCase=None			,       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     parent
							_UpperCAmelCase							   =     13
							_UpperCAmelCase							   =     7
							_UpperCAmelCase							   =     True
							_UpperCAmelCase							   =     True
							_UpperCAmelCase							   =     True
							_UpperCAmelCase							   =     True
							_UpperCAmelCase							   =     99
							_UpperCAmelCase							   =     16
							_UpperCAmelCase							   =     32
							_UpperCAmelCase							   =     2
							_UpperCAmelCase							   =     4
							_UpperCAmelCase							   =     37
							_UpperCAmelCase							   =     'gelu'
							_UpperCAmelCase							   =     0.1
							_UpperCAmelCase							   =     0.1
							_UpperCAmelCase							   =     512
							_UpperCAmelCase							   =     16
							_UpperCAmelCase							   =     2
							_UpperCAmelCase							   =     0.02
							_UpperCAmelCase							   =     3
							_UpperCAmelCase							   =     4
							_UpperCAmelCase							   =     None

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     ids_tensor([self.batch_size, self.seq_length]			,       self.vocab_size       )

							_UpperCAmelCase							   =     None
							if self.use_input_mask:
													_UpperCAmelCase							   =     random_attention_mask([self.batch_size, self.seq_length]       )

							_UpperCAmelCase							   =     None
							if self.use_token_type_ids:
													_UpperCAmelCase							   =     ids_tensor([self.batch_size, self.seq_length]			,       self.type_vocab_size       )

							_UpperCAmelCase							   =     None
							_UpperCAmelCase							   =     None
							_UpperCAmelCase							   =     None
							if self.use_labels:
													_UpperCAmelCase							   =     ids_tensor([self.batch_size]			,       self.type_sequence_label_size       )
													_UpperCAmelCase							   =     ids_tensor([self.batch_size, self.seq_length]			,       self.num_labels       )
													_UpperCAmelCase							   =     ids_tensor([self.batch_size]			,       self.num_choices       )

							_UpperCAmelCase							   =     AlbertConfig(
							    vocab_size=self.vocab_size			,       hidden_size=self.hidden_size			,       embedding_size=self.embedding_size			,       num_hidden_layers=self.num_hidden_layers			,       num_attention_heads=self.num_attention_heads			,       intermediate_size=self.intermediate_size			,       hidden_act=self.hidden_act			,       hidden_dropout_prob=self.hidden_dropout_prob			,       attention_probs_dropout_prob=self.attention_probs_dropout_prob			,       max_position_embeddings=self.max_position_embeddings			,       type_vocab_size=self.type_vocab_size			,       initializer_range=self.initializer_range			,       )

							return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     TFAlbertModel(config=UpperCAmelCase       )
							# inputs = {'input_ids': input_ids,
							#           'attention_mask': input_mask,
							#           'token_type_ids': token_type_ids}
							# sequence_output, pooled_output = model(**inputs)
							_UpperCAmelCase							   =     {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
							_UpperCAmelCase							   =     model(UpperCAmelCase       )

							_UpperCAmelCase							   =     [input_ids, input_mask]
							_UpperCAmelCase							   =     model(UpperCAmelCase       )

							_UpperCAmelCase							   =     model(UpperCAmelCase       )

							self.parent.assertEqual(result.last_hidden_state.shape			,       (self.batch_size, self.seq_length, self.hidden_size)       )
							self.parent.assertEqual(result.pooler_output.shape			,       (self.batch_size, self.hidden_size)       )

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.num_labels
							_UpperCAmelCase							   =     TFAlbertForPreTraining(config=UpperCAmelCase       )
							_UpperCAmelCase							   =     {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
							_UpperCAmelCase							   =     model(UpperCAmelCase       )
							self.parent.assertEqual(result.prediction_logits.shape			,       (self.batch_size, self.seq_length, self.vocab_size)       )
							self.parent.assertEqual(result.sop_logits.shape			,       (self.batch_size, self.num_labels)       )

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     TFAlbertForMaskedLM(config=UpperCAmelCase       )
							_UpperCAmelCase							   =     {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
							_UpperCAmelCase							   =     model(UpperCAmelCase       )
							self.parent.assertEqual(result.logits.shape			,       (self.batch_size, self.seq_length, self.vocab_size)       )

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.num_labels
							_UpperCAmelCase							   =     TFAlbertForSequenceClassification(config=UpperCAmelCase       )
							_UpperCAmelCase							   =     {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
							_UpperCAmelCase							   =     model(UpperCAmelCase       )
							self.parent.assertEqual(result.logits.shape			,       (self.batch_size, self.num_labels)       )

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     TFAlbertForQuestionAnswering(config=UpperCAmelCase       )
							_UpperCAmelCase							   =     {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
							_UpperCAmelCase							   =     model(UpperCAmelCase       )
							self.parent.assertEqual(result.start_logits.shape			,       (self.batch_size, self.seq_length)       )
							self.parent.assertEqual(result.end_logits.shape			,       (self.batch_size, self.seq_length)       )

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.num_choices
							_UpperCAmelCase							   =     TFAlbertForMultipleChoice(config=UpperCAmelCase       )
							_UpperCAmelCase							   =     tf.tile(tf.expand_dims(UpperCAmelCase			,       1       )			,       (1, self.num_choices, 1)       )
							_UpperCAmelCase							   =     tf.tile(tf.expand_dims(UpperCAmelCase			,       1       )			,       (1, self.num_choices, 1)       )
							_UpperCAmelCase							   =     tf.tile(tf.expand_dims(UpperCAmelCase			,       1       )			,       (1, self.num_choices, 1)       )
							_UpperCAmelCase							   =     {
							    'input_ids': multiple_choice_inputs_ids,
							    'attention_mask': multiple_choice_input_mask,
							    'token_type_ids': multiple_choice_token_type_ids,
							}
							_UpperCAmelCase							   =     model(UpperCAmelCase       )
							self.parent.assertListEqual(list(result['logits'].shape       )			,       [self.batch_size, self.num_choices]       )

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.num_labels
							_UpperCAmelCase							   =     TFAlbertForTokenClassification(config=UpperCAmelCase       )
							_UpperCAmelCase							   =     {
							    'input_ids': input_ids,
							    'attention_mask': input_mask,
							    'token_type_ids': token_type_ids,
							}
							_UpperCAmelCase							   =     model(UpperCAmelCase       )
							self.parent.assertListEqual(list(result['logits'].shape       )			,       [self.batch_size, self.seq_length, self.num_labels]       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.prepare_config_and_inputs()
							(
							    (
							    _UpperCAmelCase
							)   ,					(
							    _UpperCAmelCase
							)   ,					(
							    _UpperCAmelCase
							)   ,					(
							    _UpperCAmelCase
							)   ,					(
							    _UpperCAmelCase
							)   ,					(
							    _UpperCAmelCase
							)   ,					(
							    _UpperCAmelCase
							)   ,					
							)							   =     config_and_inputs
							_UpperCAmelCase							   =     {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
							return config, inputs_dict





@require_tf
class __lowerCamelCase					( snake_case__	,			snake_case__	,			unittest.TestCase):





	"""simple docstring"""






	UpperCamelCase__								= (
	    (
	        TFAlbertModel,
	        TFAlbertForPreTraining,
	        TFAlbertForMaskedLM,
	        TFAlbertForSequenceClassification,
	        TFAlbertForQuestionAnswering,
	        TFAlbertForTokenClassification,
	        TFAlbertForMultipleChoice,
	    )
	    if is_tf_available()
	    else ()
	)
	UpperCamelCase__								= (
	    {
	        "feature-extraction": TFAlbertModel,
	        "fill-mask": TFAlbertForMaskedLM,
	        "question-answering": TFAlbertForQuestionAnswering,
	        "text-classification": TFAlbertForSequenceClassification,
	        "token-classification": TFAlbertForTokenClassification,
	        "zero-shot": TFAlbertForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	UpperCamelCase__								= False
	UpperCamelCase__								= False

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase=False       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     super()._prepare_for_class(UpperCAmelCase			,       UpperCAmelCase			,       return_labels=UpperCAmelCase       )

							if return_labels:
													if model_class in get_values(UpperCAmelCase       ):
																			_UpperCAmelCase							   =     tf.zeros(self.model_tester.batch_size			,       dtype=tf.intaa       )

							return inputs_dict

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     TFAlbertModelTester(self       )
							_UpperCAmelCase							   =     ConfigTester(self			,       config_class=UpperCAmelCase			,       hidden_size=37       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							self.config_tester.run_common_tests()

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_albert_model(*UpperCAmelCase       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_albert_for_pretraining(*UpperCAmelCase       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_albert_for_masked_lm(*UpperCAmelCase       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_albert_for_multiple_choice(*UpperCAmelCase       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_albert_for_sequence_classification(*UpperCAmelCase       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_albert_for_question_answering(*UpperCAmelCase       )

	@slow
	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													_UpperCAmelCase							   =     TFAlbertModel.from_pretrained(UpperCAmelCase       )
													self.assertIsNotNone(UpperCAmelCase       )





@require_tf
class __lowerCamelCase					( unittest.TestCase):





	"""simple docstring"""






	@slow
	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     TFAlbertForPreTraining.from_pretrained('albert-base-v2'       )
							_UpperCAmelCase							   =     tf.constant([[0, 1, 2, 3, 4, 5]]       )
							_UpperCAmelCase							   =     model(UpperCAmelCase       )[0]

							_UpperCAmelCase							   =     [1, 6, 3_0000]
							self.assertEqual(output.shape			,       UpperCAmelCase       )

							_UpperCAmelCase							   =     tf.constant(
							    [
							        [
							            [4.59_56_68, 0.74_46_27_54, -1.81_81_47],
							            [4.5_95_43_47, 0.7_45_41_84, -1.8_18_82_58],
							            [4.5_95_49_05, 0.7_44_82_35, -1.8_18_23_16],
							        ]
							    ]       )
							tf.debugging.assert_near(output[:, :3, :3]			,       UpperCAmelCase			,       atol=1e-4       )








import argparse
import pathlib

import fairseq
import torch
from fairseq.models.roberta import RobertaModel as FairseqRobertaModel
from fairseq.modules import TransformerSentenceEncoderLayer
from packaging import version

from transformers import RobertaConfig, RobertaForMaskedLM, RobertaForSequenceClassification
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


if version.parse(fairseq.__version__) < version.parse('''0.9.0'''):
		raise Exception('''requires fairseq >= 0.9.0''')


logging.set_verbosity_info()
_a       =    logging.get_logger(__name__)

_a       =    '''Hello world! cécé herlolip'''


def __A       ( __lowerCAmelCase					,						__lowerCAmelCase					,						__lowerCAmelCase      )->    List[Any]:





						"""simple docstring"""



						_UpperCAmelCase							   =     FairseqRobertaModel.from_pretrained(__lowerCAmelCase      )
						roberta.eval()  # disable dropout
						_UpperCAmelCase							   =     roberta.model.encoder.sentence_encoder
						_UpperCAmelCase							   =     RobertaConfig(
						    vocab_size=roberta_sent_encoder.embed_tokens.num_embeddings					,						hidden_size=roberta.args.encoder_embed_dim					,						num_hidden_layers=roberta.args.encoder_layers					,						num_attention_heads=roberta.args.encoder_attention_heads					,						intermediate_size=roberta.args.encoder_ffn_embed_dim					,						max_position_embeddings=514					,						type_vocab_size=1					,						layer_norm_eps=1E-5					,						)
						if classification_head:
												_UpperCAmelCase							   =     roberta.model.classification_heads['mnli'].out_proj.weight.shape[0]
						print('Our BERT config:'					,						__lowerCAmelCase      )

						_UpperCAmelCase							   =     RobertaForSequenceClassification(__lowerCAmelCase      ) if classification_head else RobertaForMaskedLM(__lowerCAmelCase      )
						model.eval()

						# Now let's copy all the weights.
						# Embeddings
						_UpperCAmelCase							   =     roberta_sent_encoder.embed_tokens.weight
						_UpperCAmelCase							   =     roberta_sent_encoder.embed_positions.weight
						_UpperCAmelCase							   =     torch.zeros_like(
						    model.roberta.embeddings.token_type_embeddings.weight      )  # just zero them out b/c RoBERTa doesn't use them.
						_UpperCAmelCase							   =     roberta_sent_encoder.emb_layer_norm.weight
						_UpperCAmelCase							   =     roberta_sent_encoder.emb_layer_norm.bias

						for i in range(config.num_hidden_layers      ):
												# Encoder: start of layer
												_UpperCAmelCase							   =     model.roberta.encoder.layer[i]
												_UpperCAmelCase							   =     roberta_sent_encoder.layers[i]

												# self attention
												_UpperCAmelCase							   =     layer.attention.self
												assert (
												    roberta_layer.self_attn.k_proj.weight.data.shape
												    == roberta_layer.self_attn.q_proj.weight.data.shape
												    == roberta_layer.self_attn.v_proj.weight.data.shape
												    == torch.Size((config.hidden_size, config.hidden_size)      )
												)

												_UpperCAmelCase							   =     roberta_layer.self_attn.q_proj.weight
												_UpperCAmelCase							   =     roberta_layer.self_attn.q_proj.bias
												_UpperCAmelCase							   =     roberta_layer.self_attn.k_proj.weight
												_UpperCAmelCase							   =     roberta_layer.self_attn.k_proj.bias
												_UpperCAmelCase							   =     roberta_layer.self_attn.v_proj.weight
												_UpperCAmelCase							   =     roberta_layer.self_attn.v_proj.bias

												# self-attention output
												_UpperCAmelCase							   =     layer.attention.output
												assert self_output.dense.weight.shape == roberta_layer.self_attn.out_proj.weight.shape
												_UpperCAmelCase							   =     roberta_layer.self_attn.out_proj.weight
												_UpperCAmelCase							   =     roberta_layer.self_attn.out_proj.bias
												_UpperCAmelCase							   =     roberta_layer.self_attn_layer_norm.weight
												_UpperCAmelCase							   =     roberta_layer.self_attn_layer_norm.bias

												# intermediate
												_UpperCAmelCase							   =     layer.intermediate
												assert intermediate.dense.weight.shape == roberta_layer.fca.weight.shape
												_UpperCAmelCase							   =     roberta_layer.fca.weight
												_UpperCAmelCase							   =     roberta_layer.fca.bias

												# output
												_UpperCAmelCase							   =     layer.output
												assert bert_output.dense.weight.shape == roberta_layer.fca.weight.shape
												_UpperCAmelCase							   =     roberta_layer.fca.weight
												_UpperCAmelCase							   =     roberta_layer.fca.bias
												_UpperCAmelCase							   =     roberta_layer.final_layer_norm.weight
												_UpperCAmelCase							   =     roberta_layer.final_layer_norm.bias
												# end of layer

						if classification_head:
												_UpperCAmelCase							   =     roberta.model.classification_heads['mnli'].dense.weight
												_UpperCAmelCase							   =     roberta.model.classification_heads['mnli'].dense.bias
												_UpperCAmelCase							   =     roberta.model.classification_heads['mnli'].out_proj.weight
												_UpperCAmelCase							   =     roberta.model.classification_heads['mnli'].out_proj.bias
						else:
												# LM Head
												_UpperCAmelCase							   =     roberta.model.encoder.lm_head.dense.weight
												_UpperCAmelCase							   =     roberta.model.encoder.lm_head.dense.bias
												_UpperCAmelCase							   =     roberta.model.encoder.lm_head.layer_norm.weight
												_UpperCAmelCase							   =     roberta.model.encoder.lm_head.layer_norm.bias
												_UpperCAmelCase							   =     roberta.model.encoder.lm_head.weight
												_UpperCAmelCase							   =     roberta.model.encoder.lm_head.bias

						# Let's check that we get the same results.
						_UpperCAmelCase							   =     roberta.encode(__lowerCAmelCase      ).unsqueeze(0      )  # batch of size 1

						_UpperCAmelCase							   =     model(__lowerCAmelCase      )[0]
						if classification_head:
												_UpperCAmelCase							   =     roberta.model.classification_heads['mnli'](roberta.extract_features(__lowerCAmelCase      )      )
						else:
												_UpperCAmelCase							   =     roberta.model(__lowerCAmelCase      )[0]
						print(our_output.shape					,						their_output.shape      )
						_UpperCAmelCase							   =     torch.max(torch.abs(our_output - their_output      )      ).item()
						print(F"""max_absolute_diff = {max_absolute_diff}"""      )  # ~ 1e-7
						_UpperCAmelCase							   =     torch.allclose(__lowerCAmelCase					,						__lowerCAmelCase					,						atol=1E-3      )
						print('Do both models output the same tensors?'					,						'🔥' if success else '💩'      )
						if not success:
												raise Exception('Something went wRoNg'      )

						pathlib.Path(__lowerCAmelCase      ).mkdir(parents=__lowerCAmelCase					,						exist_ok=__lowerCAmelCase      )
						print(F"""Saving model to {pytorch_dump_folder_path}"""      )
						model.save_pretrained(__lowerCAmelCase      )


if __name__ == "__main__":
		_a       =    argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--roberta_checkpoint_path''', default=None, type=str, required=True, help='''Path the official PyTorch dump.'''
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
		)
		parser.add_argument(
		    '''--classification_head''', action='''store_true''', help='''Whether to convert a final classification head.'''
		)
		_a       =    parser.parse_args()
		convert_roberta_checkpoint_to_pytorch(
		    args.roberta_checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
		)
