

'''simple docstring'''


from __future__ import annotations

import unittest

from transformers import FunnelConfig, is_tf_available
from transformers.testing_utils import require_tf

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 (
						    TFFunnelBaseModel,
						    TFFunnelForMaskedLM,
						    TFFunnelForMultipleChoice,
						    TFFunnelForPreTraining,
						    TFFunnelForQuestionAnswering,
						    TFFunnelForSequenceClassification,
						    TFFunnelForTokenClassification,
						    TFFunnelModel,
						)



class 					lowercase						:


							def __init__(			self    ,			_snake_case    ,			_snake_case=13    ,			_snake_case=7    ,			_snake_case=True    ,			_snake_case=True    ,			_snake_case=True    ,			_snake_case=True    ,			_snake_case=99    ,			_snake_case=[1, 1, 2]    ,			_snake_case=1    ,			_snake_case=32    ,			_snake_case=4    ,			_snake_case=8    ,			_snake_case=37    ,			_snake_case="gelu_new"    ,			_snake_case=0.1    ,			_snake_case=0.1    ,			_snake_case=0.0    ,			_snake_case=512    ,			_snake_case=3    ,			_snake_case=0.02    ,			_snake_case=3    ,			_snake_case=4    ,			_snake_case=None    ,			_snake_case=False    ,			)							->				Dict:
								UpperCAmelCase_		:      int       	=					parent
								UpperCAmelCase_		:      Union[str, Any]       	=					batch_size
								UpperCAmelCase_		:      List[str]       	=					seq_length
								UpperCAmelCase_		:      Optional[int]       	=					is_training
								UpperCAmelCase_		:      Union[str, Any]       	=					use_input_mask
								UpperCAmelCase_		:      str       	=					use_token_type_ids
								UpperCAmelCase_		:      Dict       	=					use_labels
								UpperCAmelCase_		:      Union[str, Any]       	=					vocab_size
								UpperCAmelCase_		:      Tuple       	=					block_sizes
								UpperCAmelCase_		:      Union[str, Any]       	=					num_decoder_layers
								UpperCAmelCase_		:      int       	=					d_model
								UpperCAmelCase_		:      Dict       	=					n_head
								UpperCAmelCase_		:      Tuple       	=					d_head
								UpperCAmelCase_		:      Dict       	=					d_inner
								UpperCAmelCase_		:      Dict       	=					hidden_act
								UpperCAmelCase_		:      List[str]       	=					hidden_dropout
								UpperCAmelCase_		:      Tuple       	=					attention_dropout
								UpperCAmelCase_		:      Optional[Any]       	=					activation_dropout
								UpperCAmelCase_		:      Union[str, Any]       	=					max_position_embeddings
								UpperCAmelCase_		:      List[Any]       	=					type_vocab_size
								UpperCAmelCase_		:      Dict       	=					2
								UpperCAmelCase_		:      Tuple       	=					num_labels
								UpperCAmelCase_		:      List[Any]       	=					num_choices
								UpperCAmelCase_		:      List[Any]       	=					scope
								UpperCAmelCase_		:      int       	=					initializer_std

								# Used in the tests to check the size of the first attention layer
								UpperCAmelCase_		:      List[Any]       	=					n_head
								# Used in the tests to check the size of the first hidden state
								UpperCAmelCase_		:      Union[str, Any]       	=					self.d_model
								# Used in the tests to check the number of output hidden states/attentions
								UpperCAmelCase_		:      Dict       	=					sum(self.block_sizes) + (0 if base else self.num_decoder_layers)
								# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
								# the last hidden state of the first block (which is the first hidden state of the decoder).
								if not base:
									UpperCAmelCase_		:      List[str]       	=					self.num_hidden_layers + 2


							def 						_snake_case							(			self)							->				int:
								UpperCAmelCase_		:      Dict       	=					ids_tensor([self.batch_size, self.seq_length]    ,			self.vocab_size)

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

								UpperCAmelCase_		:      Union[str, Any]       	=					None
								if self.use_token_type_ids:
									UpperCAmelCase_		:      Optional[Any]       	=					ids_tensor([self.batch_size, self.seq_length]    ,			self.type_vocab_size)

								UpperCAmelCase_		:      int       	=					None
								UpperCAmelCase_		:      Tuple       	=					None
								UpperCAmelCase_		:      Tuple       	=					None
								if self.use_labels:
									UpperCAmelCase_		:      str       	=					ids_tensor([self.batch_size]    ,			self.type_sequence_label_size)
									UpperCAmelCase_		:      int       	=					ids_tensor([self.batch_size, self.seq_length]    ,			self.num_labels)
									UpperCAmelCase_		:      List[Any]       	=					ids_tensor([self.batch_size]    ,			self.num_choices)

								UpperCAmelCase_		:      Tuple       	=					FunnelConfig(
								    vocab_size=self.vocab_size    ,			block_sizes=self.block_sizes    ,			num_decoder_layers=self.num_decoder_layers    ,			d_model=self.d_model    ,			n_head=self.n_head    ,			d_head=self.d_head    ,			d_inner=self.d_inner    ,			hidden_act=self.hidden_act    ,			hidden_dropout=self.hidden_dropout    ,			attention_dropout=self.attention_dropout    ,			activation_dropout=self.activation_dropout    ,			max_position_embeddings=self.max_position_embeddings    ,			type_vocab_size=self.type_vocab_size    ,			initializer_std=self.initializer_std    ,			)

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


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			)							->				Tuple:
								UpperCAmelCase_		:      List[Any]       	=					TFFunnelModel(config=_lowerCamelCase)
								UpperCAmelCase_		:      Any       	=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
								UpperCAmelCase_		:      List[Any]       	=					model(_lowerCamelCase)

								UpperCAmelCase_		:      Any       	=					[input_ids, input_mask]
								UpperCAmelCase_		:      List[Any]       	=					model(_lowerCamelCase)

								UpperCAmelCase_		:      Union[str, Any]       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.last_hidden_state.shape    ,			(self.batch_size, self.seq_length, self.d_model))

								UpperCAmelCase_		:      int       	=					False
								UpperCAmelCase_		:      Tuple       	=					TFFunnelModel(config=_lowerCamelCase)
								UpperCAmelCase_		:      Optional[Any]       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.last_hidden_state.shape    ,			(self.batch_size, self.seq_length, self.d_model))

								UpperCAmelCase_		:      Tuple       	=					False
								UpperCAmelCase_		:      Tuple       	=					TFFunnelModel(config=_lowerCamelCase)
								UpperCAmelCase_		:      Optional[Any]       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.last_hidden_state.shape    ,			(self.batch_size, self.seq_length, self.d_model))


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			)							->				int:
								UpperCAmelCase_		:      Optional[Any]       	=					TFFunnelBaseModel(config=_lowerCamelCase)
								UpperCAmelCase_		:      Optional[int]       	=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
								UpperCAmelCase_		:      Tuple       	=					model(_lowerCamelCase)

								UpperCAmelCase_		:      List[Any]       	=					[input_ids, input_mask]
								UpperCAmelCase_		:      int       	=					model(_lowerCamelCase)

								UpperCAmelCase_		:      Optional[int]       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.last_hidden_state.shape    ,			(self.batch_size, 2, self.d_model))

								UpperCAmelCase_		:      List[Any]       	=					False
								UpperCAmelCase_		:      str       	=					TFFunnelBaseModel(config=_lowerCamelCase)
								UpperCAmelCase_		:      Tuple       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.last_hidden_state.shape    ,			(self.batch_size, 3, self.d_model))

								UpperCAmelCase_		:      int       	=					False
								UpperCAmelCase_		:      Union[str, Any]       	=					TFFunnelBaseModel(config=_lowerCamelCase)
								UpperCAmelCase_		:      Tuple       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.last_hidden_state.shape    ,			(self.batch_size, 2, self.d_model))


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			)							->				Optional[Any]:
								UpperCAmelCase_		:      int       	=					TFFunnelForPreTraining(config=_lowerCamelCase)
								UpperCAmelCase_		:      Union[str, Any]       	=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
								UpperCAmelCase_		:      List[str]       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.logits.shape    ,			(self.batch_size, self.seq_length))


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			)							->				str:
								UpperCAmelCase_		:      Dict       	=					TFFunnelForMaskedLM(config=_lowerCamelCase)
								UpperCAmelCase_		:      Dict       	=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
								UpperCAmelCase_		:      str       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.logits.shape    ,			(self.batch_size, self.seq_length, self.vocab_size))


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			)							->				Any:
								UpperCAmelCase_		:      Any       	=					self.num_labels
								UpperCAmelCase_		:      Union[str, Any]       	=					TFFunnelForSequenceClassification(config=_lowerCamelCase)
								UpperCAmelCase_		:      Tuple       	=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
								UpperCAmelCase_		:      Union[str, Any]       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.logits.shape    ,			(self.batch_size, self.num_labels))


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			)							->				int:
								UpperCAmelCase_		:      List[Any]       	=					self.num_choices
								UpperCAmelCase_		:      Tuple       	=					TFFunnelForMultipleChoice(config=_lowerCamelCase)
								UpperCAmelCase_		:      List[str]       	=					tf.tile(tf.expand_dims(_lowerCamelCase    ,			1)    ,			(1, self.num_choices, 1))
								UpperCAmelCase_		:      Tuple       	=					tf.tile(tf.expand_dims(_lowerCamelCase    ,			1)    ,			(1, self.num_choices, 1))
								UpperCAmelCase_		:      List[str]       	=					tf.tile(tf.expand_dims(_lowerCamelCase    ,			1)    ,			(1, self.num_choices, 1))
								UpperCAmelCase_		:      List[str]       	=					{
								    'input_ids': multiple_choice_inputs_ids,
								    'attention_mask': multiple_choice_input_mask,
								    'token_type_ids': multiple_choice_token_type_ids,
								}
								UpperCAmelCase_		:      Optional[int]       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.logits.shape    ,			(self.batch_size, self.num_choices))


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			)							->				Dict:
								UpperCAmelCase_		:      Union[str, Any]       	=					self.num_labels
								UpperCAmelCase_		:      str       	=					TFFunnelForTokenClassification(config=_lowerCamelCase)
								UpperCAmelCase_		:      Optional[int]       	=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
								UpperCAmelCase_		:      str       	=					model(_lowerCamelCase)
								self.parent.assertEqual(result.logits.shape    ,			(self.batch_size, self.seq_length, self.num_labels))


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			)							->				Any:
								UpperCAmelCase_		:      List[Any]       	=					TFFunnelForQuestionAnswering(config=_lowerCamelCase)
								UpperCAmelCase_		:      Union[str, Any]       	=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
								UpperCAmelCase_		:      Any       	=					model(_lowerCamelCase)
								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 						_snake_case							(			self)							->				Optional[Any]:
								UpperCAmelCase_		:      Union[str, Any]       	=					self.prepare_config_and_inputs()
								(
								    (
								    UpperCAmelCase_	
								)	,    (
								    UpperCAmelCase_	
								)	,    (
								    UpperCAmelCase_	
								)	,    (
								    UpperCAmelCase_	
								)	,    (
								    UpperCAmelCase_	
								)	,    (
								    UpperCAmelCase_	
								)	,    (
								    UpperCAmelCase_	
								)	,    	
								)		:      Tuple       	=					config_and_inputs
								UpperCAmelCase_		:      List[Any]       	=					{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
								return config, inputs_dict



@require_tf
class 					lowercase						(  __UpperCAmelCase,					__UpperCAmelCase,					unittest.TestCase ):
							_lowerCamelCase		:      Union[str, Any]=   (
							    (
							        TFFunnelModel,
							        TFFunnelForMaskedLM,
							        TFFunnelForPreTraining,
							        TFFunnelForQuestionAnswering,
							        TFFunnelForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							_lowerCamelCase		:      Optional[int]=   (
							    {
							        "feature-extraction": (TFFunnelBaseModel, TFFunnelModel),
							        "fill-mask": TFFunnelForMaskedLM,
							        "question-answering": TFFunnelForQuestionAnswering,
							        "text-classification": TFFunnelForSequenceClassification,
							        "token-classification": TFFunnelForTokenClassification,
							        "zero-shot": TFFunnelForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							_lowerCamelCase		:      Dict=   False
							_lowerCamelCase		:      Union[str, Any]=   False


							def 						_snake_case							(			self)							->				Tuple:
								UpperCAmelCase_		:      Union[str, Any]       	=					TFFunnelModelTester(self)
								UpperCAmelCase_		:      Any       	=					ConfigTester(self    ,			config_class=_lowerCamelCase)


							def 						_snake_case							(			self)							->				Union[str, Any]:
								self.config_tester.run_common_tests()


							def 						_snake_case							(			self)							->				Union[str, Any]:
								UpperCAmelCase_		:      Dict       	=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_lowerCamelCase)


							def 						_snake_case							(			self)							->				Optional[int]:
								UpperCAmelCase_		:      Optional[Any]       	=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*_lowerCamelCase)


							def 						_snake_case							(			self)							->				Tuple:
								UpperCAmelCase_		:      List[Any]       	=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*_lowerCamelCase)


							def 						_snake_case							(			self)							->				int:
								UpperCAmelCase_		:      int       	=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase)





							def 						_snake_case							(			self)							->				Optional[int]:
								UpperCAmelCase_		:      Optional[int]       	=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*_lowerCamelCase)



@require_tf
class 					lowercase						(  __UpperCAmelCase,					unittest.TestCase ):
							_lowerCamelCase		:      Any=   (
							    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
							)
							_lowerCamelCase		:      Optional[int]=   False
							_lowerCamelCase		:      List[str]=   False


							def 						_snake_case							(			self)							->				Union[str, Any]:
								UpperCAmelCase_		:      int       	=					TFFunnelModelTester(self    ,			base=_lowerCamelCase)
								UpperCAmelCase_		:      str       	=					ConfigTester(self    ,			config_class=_lowerCamelCase)


							def 						_snake_case							(			self)							->				List[Any]:
								self.config_tester.run_common_tests()


							def 						_snake_case							(			self)							->				Union[str, Any]:
								UpperCAmelCase_		:      List[Any]       	=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_base_model(*_lowerCamelCase)


							def 						_snake_case							(			self)							->				Optional[int]:
								UpperCAmelCase_		:      Optional[int]       	=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*_lowerCamelCase)





							def 						_snake_case							(			self)							->				List[Any]:
								UpperCAmelCase_		:      int       	=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*_lowerCamelCase)









'''simple docstring'''


import argparse

import torch

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


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

lowerCAmelCase__										= [
    ["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 						SCREAMING_SNAKE_CASE(  UpperCamelCase  )    ->    List[Any]:
	if k == "embeddings.weight":
		return "shared.weight"

	for parlai_name, hf_name in PATTERNS:
		UpperCAmelCase_		:      Optional[int]       	=					k.replace(UpperCamelCase   ,UpperCamelCase  )

	if k.startswith('encoder'  ):
		UpperCAmelCase_		:      Dict       	=					k.replace('.attn'   ,'.self_attn'  )
		UpperCAmelCase_		:      int       	=					k.replace('norm1'   ,'self_attn_layer_norm'  )
		UpperCAmelCase_		:      int       	=					k.replace('norm2'   ,'final_layer_norm'  )
	elif k.startswith('decoder'  ):
		UpperCAmelCase_		:      Union[str, Any]       	=					k.replace('norm1'   ,'self_attn_layer_norm'  )
		UpperCAmelCase_		:      int       	=					k.replace('norm2'   ,'encoder_attn_layer_norm'  )
		UpperCAmelCase_		:      List[str]       	=					k.replace('norm3'   ,'final_layer_norm'  )
	return k







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase  )    ->    Optional[Any]:
	UpperCAmelCase_		:      Union[str, Any]       	=					[
	    'model.encoder.layernorm_embedding.weight',
	    'model.encoder.layernorm_embedding.bias',
	    'model.decoder.layernorm_embedding.weight',
	    'model.decoder.layernorm_embedding.bias',
	]
	for k in keys:
		UpperCAmelCase_		:      Optional[int]       	=					sd.pop(UpperCamelCase  )
		UpperCAmelCase_		:      List[Any]       	=					k.replace('layernorm_embedding'   ,'layer_norm'  )
		assert new_k not in sd
		UpperCAmelCase_		:      Any       	=					v


lowerCAmelCase__										= ["START"]







@torch.no_grad()
def 						SCREAMING_SNAKE_CASE(  UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase  )    ->    Dict:
	UpperCAmelCase_		:      str       	=					torch.load(UpperCamelCase   ,map_location='cpu'  )
	UpperCAmelCase_		:      Tuple       	=					model['model']
	UpperCAmelCase_		:      Union[str, Any]       	=					BlenderbotConfig.from_json_file(UpperCamelCase  )
	UpperCAmelCase_		:      Union[str, Any]       	=					BlenderbotForConditionalGeneration(UpperCamelCase  )
	UpperCAmelCase_		:      Union[str, Any]       	=					m.model.state_dict().keys()
	UpperCAmelCase_		:      Optional[Any]       	=					[]
	UpperCAmelCase_		:      Dict       	=					{}
	for k, v in sd.items():
		if k in IGNORE_KEYS:
			continue

		UpperCAmelCase_		:      Union[str, Any]       	=					rename_state_dict_key(UpperCamelCase  )
		if new_k not in valid_keys:
			failures.append([k, new_k]  )
		else:
			UpperCAmelCase_		:      List[str]       	=					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__":
						lowerCAmelCase__										= 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"
						)
						lowerCAmelCase__										= parser.parse_args()
						convert_parlai_checkpoint(args.src_path, args.save_dir, args.hf_config_json)






