


from __future__ import annotations

import os
import tempfile
import unittest

from transformers import ConvBertConfig, is_tf_available
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 (
				    TFConvBertForMaskedLM,
				    TFConvBertForMultipleChoice,
				    TFConvBertForQuestionAnswering,
				    TFConvBertForSequenceClassification,
				    TFConvBertForTokenClassification,
				    TFConvBertModel,
				)






class    _UpperCamelCase     :

				"""simple docstring"""







				def __init__(    self					, a__					, a__=13					, a__=7					, a__=True					, a__=True					, a__=True					, a__=True					, a__=99					, a__=32					, a__=2					, a__=4					, a__=37					, a__="gelu"					, a__=0.1					, a__=0.1					, a__=512					, a__=16					, a__=2					, a__=0.02					, a__=3					, a__=4					, a__=None					, )						->					List[Any]:
					A             =					parent
					A             =					13
					A             =					7
					A             =					True
					A             =					True
					A             =					True
					A             =					True
					A             =					99
					A             =					384
					A             =					2
					A             =					4
					A             =					37
					A             =					"""gelu"""
					A             =					0.1
					A             =					0.1
					A             =					512
					A             =					16
					A             =					2
					A             =					0.02
					A             =					3
					A             =					4
					A             =					128
					A             =					2
					A             =					9
					A             =					1
					A             =					None

				def   _UpperCAmelCase       (    self )						->					Tuple:
					A             =					ids_tensor([self.batch_size, self.seq_length]					, self.vocab_size )

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

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

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

					A             =					ConvBertConfig(
					    vocab_size=self.vocab_size					, hidden_size=self.hidden_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_dict=a__					, )

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

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__ )						->					Optional[int]:
					A             =					TFConvBertModel(config=a__ )
					A             =					{"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}

					A             =					[input_ids, input_mask]
					A             =					model(a__ )

					A             =					model(a__ )

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

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__ )						->					Optional[Any]:
					A             =					TFConvBertForMaskedLM(config=a__ )
					A             =					{
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}
					A             =					model(a__ )
					self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.vocab_size) )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__ )						->					List[str]:
					A             =					self.num_labels
					A             =					TFConvBertForSequenceClassification(config=a__ )
					A             =					{
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}

					A             =					model(a__ )
					self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.num_labels) )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__ )						->					List[str]:
					A             =					self.num_choices
					A             =					TFConvBertForMultipleChoice(config=a__ )
					A             =					tf.tile(tf.expand_dims(a__					, 1 )					, (1, self.num_choices, 1) )
					A             =					tf.tile(tf.expand_dims(a__					, 1 )					, (1, self.num_choices, 1) )
					A             =					tf.tile(tf.expand_dims(a__					, 1 )					, (1, self.num_choices, 1) )
					A             =					{
					    """input_ids""": multiple_choice_inputs_ids,
					    """attention_mask""": multiple_choice_input_mask,
					    """token_type_ids""": multiple_choice_token_type_ids,
					}
					A             =					model(a__ )
					self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.num_choices) )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__ )						->					Tuple:
					A             =					self.num_labels
					A             =					TFConvBertForTokenClassification(config=a__ )
					A             =					{
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}
					A             =					model(a__ )
					self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.num_labels) )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__ )						->					Dict:
					A             =					TFConvBertForQuestionAnswering(config=a__ )
					A             =					{
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}

					A             =					model(a__ )
					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 )						->					List[Any]:
					A             =					self.prepare_config_and_inputs()
					(
					    (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      
					)             =					config_and_inputs
					A             =					{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
					return config, inputs_dict






@require_tf
class    _UpperCamelCase     (       __snake_case     ,		__snake_case     ,		unittest.TestCase					):

				"""simple docstring"""







				lowerCAmelCase					=			(
				    (
				        TFConvBertModel,
				        TFConvBertForMaskedLM,
				        TFConvBertForQuestionAnswering,
				        TFConvBertForSequenceClassification,
				        TFConvBertForTokenClassification,
				        TFConvBertForMultipleChoice,
				    )
				    if is_tf_available()
				    else ()
				)
				lowerCAmelCase					=			(
				    {
				        'feature-extraction': TFConvBertModel,
				        'fill-mask': TFConvBertForMaskedLM,
				        'question-answering': TFConvBertForQuestionAnswering,
				        'text-classification': TFConvBertForSequenceClassification,
				        'token-classification': TFConvBertForTokenClassification,
				        'zero-shot': TFConvBertForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				lowerCAmelCase					=			False
				lowerCAmelCase					=			False
				lowerCAmelCase					=			False

				def   _UpperCAmelCase       (    self )						->					Optional[Any]:
					A             =					TFConvBertModelTester(self )
					A             =					ConfigTester(self					, config_class=a__					, hidden_size=37 )

				def   _UpperCAmelCase       (    self )						->					List[str]:
					self.config_tester.run_common_tests()

				def   _UpperCAmelCase       (    self )						->					str:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*a__ )

				def   _UpperCAmelCase       (    self )						->					List[str]:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*a__ )

				def   _UpperCAmelCase       (    self )						->					Union[str, Any]:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_multiple_choice(*a__ )

				def   _UpperCAmelCase       (    self )						->					Optional[Any]:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_question_answering(*a__ )

				def   _UpperCAmelCase       (    self )						->					Tuple:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_sequence_classification(*a__ )

				def   _UpperCAmelCase       (    self )						->					Optional[Any]:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_token_classification(*a__ )

				@slow
				def   _UpperCAmelCase       (    self )						->					Dict:
					A       ,      A             =					self.model_tester.prepare_config_and_inputs_for_common()
					A             =					True
					A             =					True

					if hasattr(a__					, """use_cache""" ):
						A             =					True

					A             =					getattr(self.model_tester					, """encoder_seq_length"""					, self.model_tester.seq_length )
					A             =					getattr(self.model_tester					, """key_length"""					, a__ )

					for model_class in self.all_model_classes:
						A             =					self._prepare_for_class(a__					, a__ )
						A             =					model_class(a__ )
						A             =					len(model(a__ ) )

						with tempfile.TemporaryDirectory() as tmpdirname:
							model.save_pretrained(a__					, saved_model=a__ )
							A             =					os.path.join(a__					, """saved_model"""					, """1""" )
							A             =					tf.keras.models.load_model(a__ )
							A             =					model(a__ )

							if self.is_encoder_decoder:
								A             =					outputs["""encoder_hidden_states"""]
								A             =					outputs["""encoder_attentions"""]
							else:
								A             =					outputs["""hidden_states"""]
								A             =					outputs["""attentions"""]

							self.assertEqual(len(a__ )					, a__ )

							A             =					getattr(
							    self.model_tester					, """expected_num_hidden_layers"""					, self.model_tester.num_hidden_layers + 1 )

							self.assertEqual(len(a__ )					, a__ )
							self.assertListEqual(
							    list(output_hidden_states[0].shape[-2:] )					, [self.model_tester.seq_length, self.model_tester.hidden_size]					, )

							self.assertEqual(len(a__ )					, self.model_tester.num_hidden_layers )
							self.assertListEqual(
							    list(output_attentions[0].shape[-3:] )					, [self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]					, )

				@slow
				def   _UpperCAmelCase       (    self )						->					List[Any]:
					A             =					TFConvBertModel.from_pretrained("""YituTech/conv-bert-base""" )
					self.assertIsNotNone(a__ )

				def   _UpperCAmelCase       (    self )						->					List[Any]:
					A       ,      A             =					self.model_tester.prepare_config_and_inputs_for_common()
					A             =					True
					A             =					getattr(self.model_tester					, """decoder_seq_length"""					, self.model_tester.seq_length )
					A             =					getattr(self.model_tester					, """encoder_seq_length"""					, self.model_tester.seq_length )
					A             =					getattr(self.model_tester					, """key_length"""					, a__ )
					A             =					getattr(self.model_tester					, """key_length"""					, a__ )

					def check_decoder_attentions_output(a__ ):
						A             =					len(a__ )
						self.assertEqual(out_len % 2					, 0 )
						A             =					outputs.decoder_attentions
						self.assertEqual(len(a__ )					, self.model_tester.num_hidden_layers )
						self.assertListEqual(
						    list(decoder_attentions[0].shape[-3:] )					, [self.model_tester.num_attention_heads / 2, decoder_seq_length, decoder_key_length]					, )

					def check_encoder_attentions_output(a__ ):
						A             =					[
						    t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
						]
						self.assertEqual(len(a__ )					, self.model_tester.num_hidden_layers )
						self.assertListEqual(
						    list(attentions[0].shape[-3:] )					, [self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]					, )

					for model_class in self.all_model_classes:
						A             =					True
						A             =					False
						A             =					model_class(a__ )
						A             =					model(self._prepare_for_class(a__					, a__ ) )
						A             =					len(a__ )
						self.assertEqual(config.output_hidden_states					, a__ )
						check_encoder_attentions_output(a__ )

						if self.is_encoder_decoder:
							A             =					model_class(a__ )
							A             =					model(self._prepare_for_class(a__					, a__ ) )
							self.assertEqual(config.output_hidden_states					, a__ )
							check_decoder_attentions_output(a__ )

						# Check that output attentions can also be changed via the config
						del inputs_dict["output_attentions"]
						A             =					True
						A             =					model_class(a__ )
						A             =					model(self._prepare_for_class(a__					, a__ ) )
						self.assertEqual(config.output_hidden_states					, a__ )
						check_encoder_attentions_output(a__ )

						# Check attention is always last and order is fine
						A             =					True
						A             =					True
						A             =					model_class(a__ )
						A             =					model(self._prepare_for_class(a__					, a__ ) )

						self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)					, len(a__ ) )
						self.assertEqual(model.config.output_hidden_states					, a__ )
						check_encoder_attentions_output(a__ )






@require_tf
class    _UpperCamelCase     (       unittest.TestCase					):

				"""simple docstring"""







				@slow
				def   _UpperCAmelCase       (    self )						->					Any:
					A             =					TFConvBertModel.from_pretrained("""YituTech/conv-bert-base""" )
					A             =					tf.constant([[0, 1, 2, 3, 4, 5]] )
					A             =					model(a__ )[0]

					A             =					[1, 6, 768]
					self.assertEqual(output.shape					, a__ )

					A             =					tf.constant(
					    [
					        [
					            [-0.03_47_54_93, -0.4_68_60_34, -0.30_63_88_32],
					            [0.22_63_72_48, -0.26_98_86_46, -0.7_42_34_24],
					            [0.10_32_48_68, -0.45_01_35_08, -0.58_28_07_84],
					        ]
					    ] )
					tf.debugging.assert_near(output[:, :3, :3]					, a__					, atol=1e-4 )







import argparse

import torch

from transformers import BlenderbotConfig, BlenderbotForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
_lowercase							:   List[str]						=						logging.get_logger(__name__)

_lowercase							:   int						=						[
    ["attention", "attn"],
    ["encoder_attention", "encoder_attn"],
    ["q_lin", "q_proj"],
    ["k_lin", "k_proj"],
    ["v_lin", "v_proj"],
    ["out_lin", "out_proj"],
    ["norm_embeddings", "layernorm_embedding"],
    ["position_embeddings", "embed_positions"],
    ["embeddings", "embed_tokens"],
    ["ffn.lin", "fc"],
]





def 	_lowerCAmelCase					( UpperCamelCase__:  str    )			-> int:





	"""simple docstring"""







	if k == "embeddings.weight":
		return "shared.weight"

	for parlai_name, hf_name in PATTERNS:
		A             =					k.replace(UpperCamelCase__				,       UpperCamelCase__    )

	if k.startswith("""encoder"""    ):
		A             =					k.replace(""".attn"""				,       """.self_attn"""    )
		A             =					k.replace("""norm1"""				,       """self_attn_layer_norm"""    )
		A             =					k.replace("""norm2"""				,       """final_layer_norm"""    )
	elif k.startswith("""decoder"""    ):
		A             =					k.replace("""norm1"""				,       """self_attn_layer_norm"""    )
		A             =					k.replace("""norm2"""				,       """encoder_attn_layer_norm"""    )
		A             =					k.replace("""norm3"""				,       """final_layer_norm"""    )
	return k





def 	_lowerCAmelCase					( UpperCamelCase__:  Tuple    )			-> str:





	"""simple docstring"""







	A             =					[
	    """model.encoder.layernorm_embedding.weight""",
	    """model.encoder.layernorm_embedding.bias""",
	    """model.decoder.layernorm_embedding.weight""",
	    """model.decoder.layernorm_embedding.bias""",
	]
	for k in keys:
		A             =					sd.pop(UpperCamelCase__    )
		A             =					k.replace("""layernorm_embedding"""				,       """layer_norm"""    )
		assert new_k not in sd
		A             =					v


_lowercase							:   List[Any]						=						["START"]





@torch.no_grad()
def 	_lowerCAmelCase					( UpperCamelCase__:  List[Any]				,       UpperCamelCase__:  Any				,       UpperCamelCase__:  Optional[int]    )			-> int:





	"""simple docstring"""







	A             =					torch.load(UpperCamelCase__				,       map_location="""cpu"""    )
	A             =					model["""model"""]
	A             =					BlenderbotConfig.from_json_file(UpperCamelCase__    )
	A             =					BlenderbotForConditionalGeneration(UpperCamelCase__    )
	A             =					m.model.state_dict().keys()
	A             =					[]
	A             =					{}
	for k, v in sd.items():
		if k in IGNORE_KEYS:
			continue

		A             =					rename_state_dict_key(UpperCamelCase__    )
		if new_k not in valid_keys:
			failures.append([k, new_k]    )
		else:
			A             =					v
	if cfg.normalize_before:  # Blenderbot-3B checkpoints. Rename layernorm_embedding -> layer_norm
		rename_layernorm_keys(UpperCamelCase__    )
	m.model.load_state_dict(UpperCamelCase__				,       strict=UpperCamelCase__    )
	m.half()
	m.save_pretrained(UpperCamelCase__    )


if __name__ == "__main__":
				_lowercase							:   str						=						argparse.ArgumentParser()
				# Required parameters
				parser.add_argument("--src_path", type=str, help="like blenderbot-model.bin")
				parser.add_argument("--save_dir", default="hf_blenderbot", type=str, help="Where to save converted model.")
				parser.add_argument(
				    "--hf_config_json", default="blenderbot-3b-config.json", type=str, help="Path to config to use"
				)
				_lowercase							:   List[str]						=						parser.parse_args()
				convert_parlai_checkpoint(args.src_path, args.save_dir, args.hf_config_json)



