

from __future__ import annotations

import unittest

from transformers import CTRLConfig, 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.models.ctrl.modeling_tf_ctrl import (
							    TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
							    TFCTRLForSequenceClassification,
							    TFCTRLLMHeadModel,
							    TFCTRLModel,
							)







class 						lowercase     (					_lowercase				):


		"""simple docstring"""




		def __init__(					self						,							__snake_case						,							):
			_UpperCamelCase :		Optional[Any]								=    parent
			_UpperCamelCase :		Optional[Any]								=    13
			_UpperCamelCase :		int								=    7
			_UpperCamelCase :		Union[str, Any]								=    True
			_UpperCamelCase :		Any								=    True
			_UpperCamelCase :		Optional[int]								=    True
			_UpperCamelCase :		Union[str, Any]								=    True
			_UpperCamelCase :		Optional[Any]								=    True
			_UpperCamelCase :		Optional[int]								=    99
			_UpperCamelCase :		str								=    32
			_UpperCamelCase :		Dict								=    2
			_UpperCamelCase :		Any								=    4
			_UpperCamelCase :		List[Any]								=    37
			_UpperCamelCase :		Optional[int]								=    'gelu'
			_UpperCamelCase :		Optional[Any]								=    0.1
			_UpperCamelCase :		Optional[Any]								=    0.1
			_UpperCamelCase :		List[str]								=    5_12
			_UpperCamelCase :		str								=    16
			_UpperCamelCase :		Optional[Any]								=    2
			_UpperCamelCase :		Optional[int]								=    0.0_2
			_UpperCamelCase :		int								=    3
			_UpperCamelCase :		Optional[Any]								=    4
			_UpperCamelCase :		Tuple								=    None
			_UpperCamelCase :		Union[str, Any]								=    self.vocab_size - 1


		def 		A__       (					self):
			_UpperCamelCase :		Optional[int]								=    ids_tensor([self.batch_size, self.seq_length]						,							self.vocab_size)

			_UpperCamelCase :		Dict								=    None
			if self.use_input_mask:
				_UpperCamelCase :		Tuple								=    random_attention_mask([self.batch_size, self.seq_length])

			_UpperCamelCase :		Dict								=    None
			if self.use_token_type_ids:
				_UpperCamelCase :		Optional[int]								=    ids_tensor([self.batch_size, self.seq_length]						,							self.type_vocab_size)

			_UpperCamelCase :		Optional[Any]								=    None
			if self.use_mc_token_ids:
				_UpperCamelCase :		int								=    ids_tensor([self.batch_size, self.num_choices]						,							self.seq_length)

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

			_UpperCamelCase :		Tuple								=    CTRLConfig(
			    vocab_size=self.vocab_size						,							n_embd=self.hidden_size						,							n_layer=self.num_hidden_layers						,							n_head=self.num_attention_heads						,							n_positions=self.max_position_embeddings						,							pad_token_id=self.pad_token_id						,							)

			_UpperCamelCase :		str								=    ids_tensor([self.num_hidden_layers, self.num_attention_heads]						,							2)

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


		def 		A__       (					self						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							*__snake_case):
			_UpperCamelCase :		Tuple								=    TFCTRLModel(config=__snake_case)
			_UpperCamelCase :		List[str]								=    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
			_UpperCamelCase :		List[str]								=    model(__snake_case)

			_UpperCamelCase :		Tuple								=    [input_ids, None, input_mask]  # None is the input for 'past'
			_UpperCamelCase :		Optional[Any]								=    model(__snake_case)

			_UpperCamelCase :		Optional[Any]								=    model(__snake_case)

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


		def 		A__       (					self						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							*__snake_case):
			_UpperCamelCase :		int								=    TFCTRLLMHeadModel(config=__snake_case)
			_UpperCamelCase :		Optional[int]								=    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
			_UpperCamelCase :		Union[str, Any]								=    model(__snake_case)
			self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.seq_length, self.vocab_size))


		def 		A__       (					self						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							*__snake_case):
			_UpperCamelCase :		Optional[int]								=    self.num_labels
			_UpperCamelCase :		Any								=    ids_tensor([self.batch_size]						,							self.type_sequence_label_size)
			_UpperCamelCase :		Union[str, Any]								=    {
			    'input_ids': input_ids,
			    'token_type_ids': token_type_ids,
			    'labels': sequence_labels,
			}
			_UpperCamelCase :		Any								=    TFCTRLForSequenceClassification(__snake_case)
			_UpperCamelCase :		Any								=    model(__snake_case)
			self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.num_labels))






		def 		A__       (					self):
			_UpperCamelCase :		Dict								=    self.prepare_config_and_inputs()

			(
			    (
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		
			) :		Optional[int]								=    config_and_inputs

			_UpperCamelCase :		List[str]								=    {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
			return config, inputs_dict







@require_tf
class 						lowercase     (					_lowercase     ,				_lowercase     ,				unittest.TestCase				):


		"""simple docstring"""




		a__       							=	(TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
		a__       							=	(TFCTRLLMHeadModel,) if is_tf_available() else ()
		a__       							=	(
		    {
		        "feature-extraction": TFCTRLModel,
		        "text-classification": TFCTRLForSequenceClassification,
		        "text-generation": TFCTRLLMHeadModel,
		        "zero-shot": TFCTRLForSequenceClassification,
		    }
		    if is_tf_available()
		    else {}
		)
		a__       							=	False
		a__       							=	False


		def 		A__       (					self						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case):
			if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
				# Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
				# `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
				# config could not be created.
				return True

			return False


		def 		A__       (					self):
			_UpperCamelCase :		List[Any]								=    TFCTRLModelTester(self)
			_UpperCamelCase :		Any								=    ConfigTester(self						,							config_class=__snake_case						,							n_embd=37)


		def 		A__       (					self):
			self.config_tester.run_common_tests()


		def 		A__       (					self):
			_UpperCamelCase :		Optional[int]								=    self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_ctrl_model(*__snake_case)


		def 		A__       (					self):
			_UpperCamelCase :		Optional[int]								=    self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_ctrl_lm_head(*__snake_case)


		def 		A__       (					self):
			_UpperCamelCase :		Optional[int]								=    self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_ctrl_for_sequence_classification(*__snake_case)


		def 		A__       (					self):
			_UpperCamelCase						,		_UpperCamelCase :		Optional[Any]								=    self.model_tester.prepare_config_and_inputs_for_common()
			_UpperCamelCase :		str								=    [TFCTRLLMHeadModel]
			_UpperCamelCase :		Union[str, Any]								=    [TFCTRLForSequenceClassification]

			for model_class in self.all_model_classes:
				_UpperCamelCase :		int								=    model_class(__snake_case)
				model.build()  # may be needed for the get_bias() call below
				assert isinstance(model.get_input_embeddings()						,							tf.keras.layers.Layer)

				if model_class in list_lm_models:
					_UpperCamelCase :		Tuple								=    model.get_output_embeddings()
					assert isinstance(__snake_case						,							tf.keras.layers.Layer)
					_UpperCamelCase :		Any								=    model.get_bias()
					assert isinstance(__snake_case						,							__snake_case)
					for k, v in name.items():
						assert isinstance(__snake_case						,							tf.Variable)
				elif model_class in list_other_models_with_output_ebd:
					_UpperCamelCase :		int								=    model.get_output_embeddings()
					assert isinstance(__snake_case						,							tf.keras.layers.Layer)
					_UpperCamelCase :		List[Any]								=    model.get_bias()
					assert name is None
				else:
					_UpperCamelCase :		Union[str, Any]								=    model.get_output_embeddings()
					assert x is None
					_UpperCamelCase :		Any								=    model.get_bias()
					assert name is None






		@slow
		def 		A__       (					self):
			for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
				_UpperCamelCase :		Optional[int]								=    TFCTRLModel.from_pretrained(__snake_case)
				self.assertIsNotNone(__snake_case)







@require_tf
class 						lowercase     (					unittest.TestCase				):


		"""simple docstring"""




		@slow
		def 		A__       (					self):
			_UpperCamelCase :		Optional[Any]								=    TFCTRLLMHeadModel.from_pretrained('ctrl')
			_UpperCamelCase :		Tuple								=    tf.convert_to_tensor([[1_18_59, 0, 16_11, 8]]						,							dtype=tf.intaa)  # Legal the president is
			_UpperCamelCase :		Optional[Any]								=    [
			    1_18_59,
			    0,
			    16_11,
			    8,
			    5,
			    1_50,
			    2_64_49,
			    2,
			    19,
			    3_48,
			    4_69,
			    3,
			    25_95,
			    48,
			    2_07_40,
			    24_65_33,
			    24_65_33,
			    19,
			    30,
			    5,
			]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

			_UpperCamelCase :		Any								=    model.generate(__snake_case						,							do_sample=__snake_case)
			self.assertListEqual(output_ids[0].numpy().tolist()						,							__snake_case)







import argparse

import torch
from saprl.hub import distilhubert

from transformers import HubertConfig, HubertModel, WavaVecaFeatureExtractor, logging


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

lowerCAmelCase__             =		{
    """post_extract_proj""": """feature_projection.projection""",
    """encoder.pos_conv.0""": """encoder.pos_conv_embed.conv""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """mask_emb""": """masked_spec_embed""",
}






def        lowerCamelCase_  (				UpperCAmelCase_				:  Dict				,							UpperCAmelCase_				:  List[Any]				,							UpperCAmelCase_				:  Optional[int]				,							UpperCAmelCase_				:  str				,							UpperCAmelCase_				:  List[Any]						)    ->					Any:




	'''simple docstring'''




	for attribute in key.split('.'						):
		_UpperCamelCase :		Any								=    getattr(UpperCAmelCase_				,							UpperCAmelCase_						)

	if weight_type is not None:
		_UpperCamelCase :		List[str]								=    getattr(UpperCAmelCase_				,							UpperCAmelCase_						).shape
	else:
		_UpperCamelCase :		Optional[int]								=    hf_pointer.shape

	assert hf_shape == value.shape, (
	    F'''Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'''
	    F''' {value.shape} for {full_name}'''
	)

	if weight_type == "weight":
		_UpperCamelCase :		Tuple								=    value
	elif weight_type == "weight_g":
		_UpperCamelCase :		Optional[Any]								=    value
	elif weight_type == "weight_v":
		_UpperCamelCase :		Tuple								=    value
	elif weight_type == "bias":
		_UpperCamelCase :		Any								=    value
	else:
		_UpperCamelCase :		Tuple								=    value

	logger.info(F'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'''						)






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




	'''simple docstring'''




	_UpperCamelCase :		Union[str, Any]								=    []
	_UpperCamelCase :		List[Any]								=    fairseq_model.state_dict()

	_UpperCamelCase :		int								=    hf_model.feature_extractor

	for name, value in fairseq_dict.items():
		_UpperCamelCase :		Dict								=    False
		if "conv_layers" in name:
			load_conv_layer(
			    UpperCAmelCase_				,							UpperCAmelCase_				,							UpperCAmelCase_				,							UpperCAmelCase_				,							hf_model.config.feat_extract_norm == 'group'				,							)
			_UpperCamelCase :		Union[str, Any]								=    True
		else:
			for key, mapped_key in MAPPING.items():
				_UpperCamelCase :		Tuple								=    mapped_key

				if key in name:
					_UpperCamelCase :		List[Any]								=    True
					if "*" in mapped_key:
						_UpperCamelCase :		List[Any]								=    name.split(UpperCAmelCase_						)[0].split('.'						)[-2]
						_UpperCamelCase :		Optional[Any]								=    mapped_key.replace('*'				,							UpperCAmelCase_						)
					if "weight_g" in name:
						_UpperCamelCase :		List[Any]								=    'weight_g'
					elif "weight_v" in name:
						_UpperCamelCase :		Tuple								=    'weight_v'
					elif "weight" in name:
						_UpperCamelCase :		Optional[Any]								=    'weight'
					elif "bias" in name:
						_UpperCamelCase :		List[Any]								=    'bias'
					else:
						_UpperCamelCase :		str								=    None
					set_recursively(UpperCAmelCase_				,							UpperCAmelCase_				,							UpperCAmelCase_				,							UpperCAmelCase_				,							UpperCAmelCase_						)
				continue
		if not is_used:
			unused_weights.append(UpperCAmelCase_						)

	logger.warning(F'''Unused weights: {unused_weights}'''						)






def        lowerCamelCase_  (				UpperCAmelCase_				:  Union[str, Any]				,							UpperCAmelCase_				:  List[Any]				,							UpperCAmelCase_				:  Optional[Any]				,							UpperCAmelCase_				:  int				,							UpperCAmelCase_				:  int						)    ->					Dict:




	'''simple docstring'''




	_UpperCamelCase :		Tuple								=    full_name.split('conv_layers.'						)[-1]
	_UpperCamelCase :		str								=    name.split('.'						)
	_UpperCamelCase :		List[Any]								=    int(items[0]						)
	_UpperCamelCase :		Tuple								=    int(items[1]						)

	if type_id == 0:
		if "bias" in name:
			assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
			    F'''{full_name} has size {value.shape}, but'''
			    F''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''
			)
			_UpperCamelCase :		Any								=    value
			logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''						)
		elif "weight" in name:
			assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
			    F'''{full_name} has size {value.shape}, but'''
			    F''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''
			)
			_UpperCamelCase :		List[str]								=    value
			logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''						)
	elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
		if "bias" in name:
			assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
			    F'''{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'''
			    " found."
			)
			_UpperCamelCase :		str								=    value
			logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''						)
		elif "weight" in name:
			assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
			    F'''{full_name} has size {value.shape}, but'''
			    F''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''
			)
			_UpperCamelCase :		Dict								=    value
			logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''						)
	else:
		unused_weights.append(UpperCAmelCase_						)






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




	'''simple docstring'''




	_UpperCamelCase :		Union[str, Any]								=    HubertConfig()
	_UpperCamelCase :		List[str]								=    model.config

	_UpperCamelCase :		Any								=    fs_config.activation_dropout
	_UpperCamelCase :		Tuple								=    False
	_UpperCamelCase :		Tuple								=    fs_config.attention_dropout
	_UpperCamelCase :		Dict								=    False
	_UpperCamelCase :		Optional[Any]								=    eval(fs_config.extractor_conv_feature_layers						)
	_UpperCamelCase :		Dict								=    [x[0] for x in conv_layers]
	_UpperCamelCase :		Any								=    [x[1] for x in conv_layers]
	_UpperCamelCase :		Any								=    [x[2] for x in conv_layers]
	_UpperCamelCase :		Dict								=    'gelu'
	_UpperCamelCase :		List[Any]								=    'layer' if fs_config.extractor_mode == 'layer_norm' else 'group'
	_UpperCamelCase :		List[Any]								=    False
	_UpperCamelCase :		List[Any]								=    0.0
	_UpperCamelCase :		Optional[int]								=    0.0
	_UpperCamelCase :		Optional[int]								=    fs_config.activation_fn
	_UpperCamelCase :		Optional[int]								=    fs_config.dropout
	_UpperCamelCase :		Union[str, Any]								=    fs_config.encoder_embed_dim
	_UpperCamelCase :		Tuple								=    0.0_2
	_UpperCamelCase :		int								=    fs_config.encoder_ffn_embed_dim
	_UpperCamelCase :		int								=    1e-5
	_UpperCamelCase :		Optional[Any]								=    0.0
	_UpperCamelCase :		List[str]								=    fs_config.encoder_attention_heads
	_UpperCamelCase :		List[str]								=    fs_config.conv_pos_groups
	_UpperCamelCase :		Union[str, Any]								=    fs_config.conv_pos
	_UpperCamelCase :		int								=    len(UpperCAmelCase_						)
	_UpperCamelCase :		Tuple								=    fs_config.encoder_layers

	return config






@torch.no_grad()
def        lowerCamelCase_  (				UpperCAmelCase_				:  str				,							UpperCAmelCase_				:  Dict=None						)    ->					int:




	'''simple docstring'''




	_UpperCamelCase :		Tuple								=    distilhubert().model.model

	if config_path is not None:
		_UpperCamelCase :		Optional[Any]								=    HubertConfig.from_pretrained(UpperCAmelCase_						)
	else:
		_UpperCamelCase :		Optional[int]								=    convert_config(UpperCAmelCase_						)
	_UpperCamelCase :		List[Any]								=    model.eval()

	_UpperCamelCase :		List[Any]								=    WavaVecaFeatureExtractor(
	    feature_size=1				,							sampling_rate=1_6_0_0_0				,							padding_value=0				,							do_normalize=UpperCAmelCase_				,							return_attention_mask=UpperCAmelCase_				,							)
	_UpperCamelCase :		Optional[Any]								=    HubertModel(UpperCAmelCase_						)

	recursively_load_weights(UpperCAmelCase_				,							UpperCAmelCase_						)

	feature_extractor.save_pretrained(UpperCAmelCase_						)
	hf_model.save_pretrained(UpperCAmelCase_						)


if __name__ == "__main__":
							lowerCAmelCase__             =		argparse.ArgumentParser()
							parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
							parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
							lowerCAmelCase__             =		parser.parse_args()
							convert_hubert_checkpoint(args.pytorch_dump_folder_path, args.config_path)




