



import unittest

from transformers import SpeechaTextaConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers.models.speech_to_text_a.modeling_speech_to_text_a import (
          SpeechaTextaDecoder,
          SpeechaTextaForCausalLM,
      )

@require_torch
class    _SCREAMING_SNAKE_CASE						:





     def __init__(							self			:				int  ,	UpperCamelCase			:				Tuple  ,	UpperCamelCase			:				Union[str, Any]=9_9  ,	UpperCamelCase			:				Optional[Any]=1_3  ,	UpperCamelCase			:				Union[str, Any]=1_6  ,	UpperCamelCase			:				Any=7  ,	UpperCamelCase			:				int=True  ,	UpperCamelCase			:				Optional[Any]=True  ,	UpperCamelCase			:				Dict=True  ,	UpperCamelCase			:				Any=False  ,	UpperCamelCase			:				List[str]=True  ,	UpperCamelCase			:				Optional[Any]=2  ,	UpperCamelCase			:				str=3_2  ,	UpperCamelCase			:				int=4  ,	UpperCamelCase			:				Tuple=4  ,	UpperCamelCase			:				int=3_0  ,	UpperCamelCase			:				Tuple=0  ,	UpperCamelCase			:				Tuple=1  ,	UpperCamelCase			:				str=2  ,	UpperCamelCase			:				Any=None  ,	)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE							:     Tuple							=						parent
      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						batch_size
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						decoder_seq_length
      # For common tests
      __SCREAMING_SNAKE_CASE							:     int							=						self.decoder_seq_length
      __SCREAMING_SNAKE_CASE							:     Tuple							=						is_training
      __SCREAMING_SNAKE_CASE							:     List[str]							=						use_attention_mask
      __SCREAMING_SNAKE_CASE							:     int							=						use_labels

      __SCREAMING_SNAKE_CASE							:     str							=						vocab_size
      __SCREAMING_SNAKE_CASE							:     List[str]							=						d_model
      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						d_model
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						decoder_layers
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						decoder_layers
      __SCREAMING_SNAKE_CASE							:     str							=						decoder_ffn_dim
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						decoder_attention_heads
      __SCREAMING_SNAKE_CASE							:     str							=						decoder_attention_heads
      __SCREAMING_SNAKE_CASE							:     int							=						eos_token_id
      __SCREAMING_SNAKE_CASE							:     str							=						bos_token_id
      __SCREAMING_SNAKE_CASE							:     str							=						pad_token_id
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						decoder_start_token_id
      __SCREAMING_SNAKE_CASE							:     Any							=						use_cache
      __SCREAMING_SNAKE_CASE							:     Tuple							=						max_position_embeddings

      __SCREAMING_SNAKE_CASE							:     str							=						None
      __SCREAMING_SNAKE_CASE							:     List[str]							=						decoder_seq_length
      __SCREAMING_SNAKE_CASE							:     Dict							=						2
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						1





     def 						__snake_case     (							self			:				Any					)->Optional[int]:
      __SCREAMING_SNAKE_CASE							:     Dict							=						ids_tensor([self.batch_size, self.decoder_seq_length]  ,	self.vocab_size					)

      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						None
      if self.use_attention_mask:
       __SCREAMING_SNAKE_CASE							:     List[str]							=						ids_tensor([self.batch_size, self.decoder_seq_length]  ,	vocab_size=2					)

      __SCREAMING_SNAKE_CASE							:     str							=						None
      if self.use_labels:
       __SCREAMING_SNAKE_CASE							:     List[str]							=						ids_tensor([self.batch_size, self.decoder_seq_length]  ,	self.vocab_size					)

      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						SpeechaTextaConfig(
          vocab_size=self.vocab_size  ,	d_model=self.d_model  ,	decoder_layers=self.decoder_layers  ,	decoder_ffn_dim=self.decoder_ffn_dim  ,	decoder_attention_heads=self.decoder_attention_heads  ,	eos_token_id=self.eos_token_id  ,	bos_token_id=self.bos_token_id  ,	use_cache=self.use_cache  ,	pad_token_id=self.pad_token_id  ,	decoder_start_token_id=self.decoder_start_token_id  ,	max_position_embeddings=self.max_position_embeddings  ,	)

      return (
          config,
          input_ids,
          attention_mask,
          lm_labels,
      )





     def 						__snake_case     (							self			:				Optional[int]  ,	UpperCamelCase			:				Tuple  ,	UpperCamelCase			:				List[Any]  ,	UpperCamelCase			:				List[str]  ,	UpperCamelCase			:				Union[str, Any]  ,	)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						True
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						SpeechaTextaDecoder(config=UpperCamelCase					).to(UpperCamelCase					).eval()
      __SCREAMING_SNAKE_CASE							:     Any							=						input_ids[:2]

      input_ids[input_ids == 0] += 1
      # first forward pass
      __SCREAMING_SNAKE_CASE							:     int							=						model(UpperCamelCase  ,	use_cache=UpperCamelCase					)
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						model(UpperCamelCase					)
      __SCREAMING_SNAKE_CASE							:     int							=						model(UpperCamelCase  ,	use_cache=UpperCamelCase					)

      self.parent.assertTrue(len(UpperCamelCase					) == len(UpperCamelCase					)					)
      self.parent.assertTrue(len(UpperCamelCase					) == len(UpperCamelCase					) + 1					)

      __SCREAMING_SNAKE_CASE							:     str							=						outputs["past_key_values"]

      # create hypothetical next token and extent to next_input_ids
      __SCREAMING_SNAKE_CASE							:     List[str]							=						ids_tensor((2, 1)  ,	config.vocab_size - 1					) + 1

      # append to next input_ids and
      __SCREAMING_SNAKE_CASE							:     Any							=						torch.cat([input_ids, next_tokens]  ,	dim=-1					)

      print(UpperCamelCase					)

      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						model(UpperCamelCase					)["last_hidden_state"]
      __SCREAMING_SNAKE_CASE							:     Any							=						model(UpperCamelCase  ,	past_key_values=UpperCamelCase					)["last_hidden_state"]

      # select random slice
      __SCREAMING_SNAKE_CASE							:     int							=						ids_tensor((1,)  ,	output_from_past.shape[-1]					).item()
      __SCREAMING_SNAKE_CASE							:     Dict							=						output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
      __SCREAMING_SNAKE_CASE							:     Any							=						output_from_past[:, 0, random_slice_idx].detach()

      # test that outputs are equal for slice
      assert torch.allclose(UpperCamelCase  ,	UpperCamelCase  ,	atol=1E-3					)






     def 						__snake_case     (							self			:				Optional[int]					)->Optional[int]:
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						self.prepare_config_and_inputs()
      (
          (
          __SCREAMING_SNAKE_CASE
      ),		(
          __SCREAMING_SNAKE_CASE
      ),		(
          __SCREAMING_SNAKE_CASE
      ),		(
          __SCREAMING_SNAKE_CASE
      ),		
      )							:     Optional[int]							=						config_and_inputs

      __SCREAMING_SNAKE_CASE							:     List[Any]							=						{
          "input_ids": input_ids,
          "attention_mask": attention_mask,
      }
      return config, inputs_dict






@require_torch
class    _SCREAMING_SNAKE_CASE						(UpperCamelCase      ,     UpperCamelCase      ,     UpperCamelCase      ,     unittest.TestCase			):
     lowerCAmelCase       =      (SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
     lowerCAmelCase       =      (SpeechaTextaForCausalLM,) if is_torch_available() else ()
     lowerCAmelCase       =      {"""text-generation""": SpeechaTextaForCausalLM} if is_torch_available() else {}
     lowerCAmelCase       =      True
     lowerCAmelCase       =      False





     def 						__snake_case     (							self			:				List[str]  ,	)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						SpeechaTextaStandaloneDecoderModelTester(self  ,	is_training=UpperCamelCase					)
      __SCREAMING_SNAKE_CASE							:     Any							=						ConfigTester(self  ,	config_class=UpperCamelCase					)





     def 						__snake_case     (							self			:				Dict					)->Optional[Any]:
      pass





     def 						__snake_case     (							self			:				List[Any]					)->List[Any]:
      pass





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





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





     def 						__snake_case     (							self			:				List[str]					)->List[Any]:
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_decoder_model_past(*UpperCamelCase					)






     def 						__snake_case     (							self			:				Optional[int]					)->str:
      return











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    (					__lowerCamelCase :  Union[str, Any]		,						__lowerCamelCase :  Optional[int]		,						__lowerCamelCase :  List[Any]		,						__lowerCamelCase :  str		,						__lowerCamelCase :  Tuple  ):



 """simple docstring"""





 for attribute in key.split("."  ):
  __SCREAMING_SNAKE_CASE							:     Any							=						getattr(__lowerCamelCase		,						__lowerCamelCase  )

 if weight_type is not None:
  __SCREAMING_SNAKE_CASE							:     List[Any]							=						getattr(__lowerCamelCase		,						__lowerCamelCase  ).shape
 else:
  __SCREAMING_SNAKE_CASE							:     Dict							=						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":
  __SCREAMING_SNAKE_CASE							:     str							=						value
 elif weight_type == "weight_g":
  __SCREAMING_SNAKE_CASE							:     str							=						value
 elif weight_type == "weight_v":
  __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						value
 elif weight_type == "bias":
  __SCREAMING_SNAKE_CASE							:     int							=						value
 else:
  __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						value

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






def 						_lowerCAmelCase    (					__lowerCamelCase :  int		,						__lowerCamelCase :  Tuple  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     Dict							=						[]
 __SCREAMING_SNAKE_CASE							:     str							=						fairseq_model.state_dict()

 __SCREAMING_SNAKE_CASE							:     List[Any]							=						hf_model.feature_extractor

 for name, value in fairseq_dict.items():
  __SCREAMING_SNAKE_CASE							:     str							=						False
  if "conv_layers" in name:
   load_conv_layer(
       __lowerCamelCase		,						__lowerCamelCase		,						__lowerCamelCase		,						__lowerCamelCase		,						hf_model.config.feat_extract_norm == "group"		,						)
   __SCREAMING_SNAKE_CASE							:     Optional[int]							=						True
  else:
   for key, mapped_key in MAPPING.items():
    __SCREAMING_SNAKE_CASE							:     List[str]							=						mapped_key

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

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






def 						_lowerCAmelCase    (					__lowerCamelCase :  List[Any]		,						__lowerCamelCase :  Optional[Any]		,						__lowerCamelCase :  str		,						__lowerCamelCase :  str		,						__lowerCamelCase :  Union[str, Any]  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     int							=						full_name.split("conv_layers."  )[-1]
 __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						name.split("."  )
 __SCREAMING_SNAKE_CASE							:     Dict							=						int(items[0]  )
 __SCREAMING_SNAKE_CASE							:     Dict							=						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."""
   )
   __SCREAMING_SNAKE_CASE							:     Dict							=						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."""
   )
   __SCREAMING_SNAKE_CASE							:     int							=						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."
   )
   __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						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."""
   )
   __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						value
   logger.info(F"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""  )
 else:
  unused_weights.append(__lowerCamelCase  )






def 						_lowerCAmelCase    (					__lowerCamelCase :  List[Any]  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						HubertConfig()
 __SCREAMING_SNAKE_CASE							:     List[str]							=						model.config

 __SCREAMING_SNAKE_CASE							:     str							=						fs_config.activation_dropout
 __SCREAMING_SNAKE_CASE							:     Any							=						False
 __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						fs_config.attention_dropout
 __SCREAMING_SNAKE_CASE							:     List[str]							=						False
 __SCREAMING_SNAKE_CASE							:     Optional[int]							=						eval(fs_config.extractor_conv_feature_layers  )
 __SCREAMING_SNAKE_CASE							:     Optional[int]							=						[x[0] for x in conv_layers]
 __SCREAMING_SNAKE_CASE							:     Any							=						[x[1] for x in conv_layers]
 __SCREAMING_SNAKE_CASE							:     Dict							=						[x[2] for x in conv_layers]
 __SCREAMING_SNAKE_CASE							:     Optional[int]							=						"gelu"
 __SCREAMING_SNAKE_CASE							:     List[Any]							=						"layer" if fs_config.extractor_mode == "layer_norm" else "group"
 __SCREAMING_SNAKE_CASE							:     Any							=						False
 __SCREAMING_SNAKE_CASE							:     Tuple							=						0.0
 __SCREAMING_SNAKE_CASE							:     Optional[int]							=						0.0
 __SCREAMING_SNAKE_CASE							:     Any							=						fs_config.activation_fn
 __SCREAMING_SNAKE_CASE							:     Optional[int]							=						fs_config.dropout
 __SCREAMING_SNAKE_CASE							:     List[Any]							=						fs_config.encoder_embed_dim
 __SCREAMING_SNAKE_CASE							:     int							=						0.02
 __SCREAMING_SNAKE_CASE							:     Tuple							=						fs_config.encoder_ffn_embed_dim
 __SCREAMING_SNAKE_CASE							:     List[Any]							=						1E-5
 __SCREAMING_SNAKE_CASE							:     Any							=						0.0
 __SCREAMING_SNAKE_CASE							:     Optional[int]							=						fs_config.encoder_attention_heads
 __SCREAMING_SNAKE_CASE							:     Dict							=						fs_config.conv_pos_groups
 __SCREAMING_SNAKE_CASE							:     Tuple							=						fs_config.conv_pos
 __SCREAMING_SNAKE_CASE							:     Tuple							=						len(__lowerCamelCase  )
 __SCREAMING_SNAKE_CASE							:     Optional[int]							=						fs_config.encoder_layers

 return config






@torch.no_grad()
def 						_lowerCAmelCase    (					__lowerCamelCase :  Any		,						__lowerCamelCase :  Optional[Any]=None  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     Dict							=						distilhubert().model.model

 if config_path is not None:
  __SCREAMING_SNAKE_CASE							:     List[Any]							=						HubertConfig.from_pretrained(__lowerCamelCase  )
 else:
  __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						convert_config(__lowerCamelCase  )
 __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						model.eval()

 __SCREAMING_SNAKE_CASE							:     Optional[int]							=						WavaVecaFeatureExtractor(
     feature_size=1		,						sampling_rate=16000		,						padding_value=0		,						do_normalize=__lowerCamelCase		,						return_attention_mask=__lowerCamelCase		,						)
 __SCREAMING_SNAKE_CASE							:     List[Any]							=						HubertModel(__lowerCamelCase  )

 recursively_load_weights(__lowerCamelCase		,						__lowerCamelCase  )

 feature_extractor.save_pretrained(__lowerCamelCase  )
 hf_model.save_pretrained(__lowerCamelCase  )


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)






