




import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConfig,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaForPreTraining,
    WavaVecaProcessor,
    logging,
)
from transformers.models.wavaveca.modeling_wavaveca import WavaVecaForSequenceClassification


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

_A     =			{
    """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""",
    """adapter_layer""": """encoder.layers.*.adapter_layer""",
    """w2v_model.layer_norm""": """feature_projection.layer_norm""",
    """quantizer.weight_proj""": """quantizer.weight_proj""",
    """quantizer.vars""": """quantizer.codevectors""",
    """project_q""": """project_q""",
    """final_proj""": """project_hid""",
    """w2v_encoder.proj""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
    """pooling_layer.linear""": """projector""",
    """pooling_layer.projection""": """classifier""",
}
_A     =			[
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
    """projector""",
    """classifier""",
]





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_      			=  {}
						with open(lowerCAmelCase__ ,  '''r'''  ) as file:
												for line_number, line in enumerate(lowerCAmelCase__  ):
																		lowerCAmelCase_      			=  line.strip()
																		if line:
																								lowerCAmelCase_      			=  line.split()
																								lowerCAmelCase_      			=  line_number
																								lowerCAmelCase_      			=  words[0]
																								lowerCAmelCase_      			=  value
						return result





def __UpperCamelCase			(       _A ,  _A ,  _A ,  _A ,  _A  ):




						for attribute in key.split('''.'''  ):
												lowerCAmelCase_      			=  getattr(lowerCAmelCase__ ,  lowerCAmelCase__  )

						lowerCAmelCase_      			=  None
						for param_key in PARAM_MAPPING.keys():
												if full_name.endswith(lowerCAmelCase__  ):
																		lowerCAmelCase_      			=  PARAM_MAPPING[full_name.split('''.'''  )[-1]]
																		lowerCAmelCase_      			=  'param'

						if weight_type is not None and weight_type != "param":
												lowerCAmelCase_      			=  getattr(lowerCAmelCase__ ,  lowerCAmelCase__  ).shape
						elif weight_type is not None and weight_type == "param":
												lowerCAmelCase_      			=  hf_pointer
												for attribute in hf_param_name.split('''.'''  ):
																		lowerCAmelCase_      			=  getattr(lowerCAmelCase__ ,  lowerCAmelCase__  )
												lowerCAmelCase_      			=  shape_pointer.shape

												# let's reduce dimension
												lowerCAmelCase_      			=  value[0]
						else:
												lowerCAmelCase_      			=  hf_pointer.shape

						if hf_shape != value.shape:
												raise ValueError(
												    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":
												lowerCAmelCase_      			=  value
						elif weight_type == "weight_g":
												lowerCAmelCase_      			=  value
						elif weight_type == "weight_v":
												lowerCAmelCase_      			=  value
						elif weight_type == "bias":
												lowerCAmelCase_      			=  value
						elif weight_type == "param":
												for attribute in hf_param_name.split('''.'''  ):
																		lowerCAmelCase_      			=  getattr(lowerCAmelCase__ ,  lowerCAmelCase__  )
												lowerCAmelCase_      			=  value
						else:
												lowerCAmelCase_      			=  value

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





def __UpperCamelCase			(       _A ,  _A ,  _A ,  _A ,  _A  ):
						lowerCAmelCase_      			=  None
						for param_key in PARAM_MAPPING.keys():
												if full_name.endswith(lowerCAmelCase__  ):
																		lowerCAmelCase_      			=  PARAM_MAPPING[full_name.split('''.'''  )[-1]]
																		lowerCAmelCase_      			=  'param'

						if weight_type is not None and weight_type != "param":
												lowerCAmelCase_      			=  '.'.join([key, weight_type]  )
						elif weight_type is not None and weight_type == "param":
												lowerCAmelCase_      			=  '.'.join([key, hf_param_name]  )
						else:
												lowerCAmelCase_      			=  key

						lowerCAmelCase_      			=  value if 'lm_head' in full_key else value[0]


_A     =			{
    """W_a""": """linear_1.weight""",
    """W_b""": """linear_2.weight""",
    """b_a""": """linear_1.bias""",
    """b_b""": """linear_2.bias""",
    """ln_W""": """norm.weight""",
    """ln_b""": """norm.bias""",
}





def __UpperCamelCase			(       _A ,  _A ,  _A=None ,  _A=None  ):
						lowerCAmelCase_      			=  False
						for key, mapped_key in MAPPING.items():
												lowerCAmelCase_      			=  'wav2vec2.' + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
												if key in name or key.split('''w2v_model.'''  )[-1] == name.split('''.'''  )[0]:
																		lowerCAmelCase_      			=  True
																		if "*" in mapped_key:
																								lowerCAmelCase_      			=  name.split(lowerCAmelCase__  )[0].split('''.'''  )[-2]
																								lowerCAmelCase_      			=  mapped_key.replace('''*''' ,  lowerCAmelCase__  )
																		if "weight_g" in name:
																								lowerCAmelCase_      			=  'weight_g'
																		elif "weight_v" in name:
																								lowerCAmelCase_      			=  'weight_v'
																		elif "bias" in name:
																								lowerCAmelCase_      			=  'bias'
																		elif "weight" in name:
																								# TODO: don't match quantizer.weight_proj
																								lowerCAmelCase_      			=  'weight'
																		else:
																								lowerCAmelCase_      			=  None
																		if hf_dict is not None:
																								rename_dict(lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__  )
																		else:
																								set_recursively(lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__  )
																		return is_used
						return is_used





def __UpperCamelCase			(       _A ,  _A ,  _A  ):
						lowerCAmelCase_      			=  []
						lowerCAmelCase_      			=  fairseq_model.state_dict()

						lowerCAmelCase_      			=  hf_model.wavaveca.feature_extractor

						for name, value in fairseq_dict.items():
												lowerCAmelCase_      			=  False
												if "conv_layers" in name:
																		load_conv_layer(
																		    lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  hf_model.config.feat_extract_norm == '''group''' ,  )
																		lowerCAmelCase_      			=  True
												else:
																		lowerCAmelCase_      			=  load_wavaveca_layer(lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__  )
												if not is_used:
																		unused_weights.append(lowerCAmelCase__  )

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





def __UpperCamelCase			(       _A ,  _A ,  _A ,  _A ,  _A  ):
						lowerCAmelCase_      			=  full_name.split('''conv_layers.'''  )[-1]
						lowerCAmelCase_      			=  name.split('''.'''  )
						lowerCAmelCase_      			=  int(items[0]  )
						lowerCAmelCase_      			=  int(items[1]  )

						if type_id == 0:
												if "bias" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
																								raise ValueError(
																								    f"{full_name} has size {value.shape}, but"
																								    f" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found."  )
																		lowerCAmelCase_      			=  value
																		logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}."  )
												elif "weight" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
																								raise ValueError(
																								    f"{full_name} has size {value.shape}, but"
																								    f" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found."  )
																		lowerCAmelCase_      			=  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:
																		if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
																								raise ValueError(
																								    f"{full_name} has size {value.shape}, but"
																								    f" {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found."  )
																		lowerCAmelCase_      			=  value
																		logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."  )
												elif "weight" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
																								raise ValueError(
																								    f"{full_name} has size {value.shape}, but"
																								    f" {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found."  )
																		lowerCAmelCase_      			=  value
																		logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."  )
						else:
												unused_weights.append(lowerCAmelCase__  )





@torch.no_grad()
def __UpperCamelCase			(       _A ,  _A ,  _A=None ,  _A=None ,  _A=True ,  _A=False  ):




						if config_path is not None:
												lowerCAmelCase_      			=  WavaVecaConfig.from_pretrained(lowerCAmelCase__  )
						else:
												lowerCAmelCase_      			=  WavaVecaConfig()

						if is_seq_class:
												lowerCAmelCase_      			=  read_txt_into_dict(lowerCAmelCase__  )
												lowerCAmelCase_      			=  idalabel
												lowerCAmelCase_      			=  WavaVecaForSequenceClassification(lowerCAmelCase__  )
												lowerCAmelCase_      			=  WavaVecaFeatureExtractor(
												    feature_size=1 ,  sampling_rate=16000 ,  padding_value=0 ,  do_normalize=lowerCAmelCase__ ,  return_attention_mask=lowerCAmelCase__ ,  )
												feature_extractor.save_pretrained(lowerCAmelCase__  )

						elif is_finetuned:
												if dict_path:
																		lowerCAmelCase_      			=  Dictionary.load(lowerCAmelCase__  )

																		# important change bos & pad token id since CTC symbol is <pad> and
																		# not <s> as in fairseq
																		lowerCAmelCase_      			=  target_dict.pad_index
																		lowerCAmelCase_      			=  target_dict.bos_index
																		lowerCAmelCase_      			=  target_dict.eos_index
																		lowerCAmelCase_      			=  len(target_dict.symbols  )
																		lowerCAmelCase_      			=  os.path.join(lowerCAmelCase__ ,  '''vocab.json'''  )
																		if not os.path.isdir(lowerCAmelCase__  ):
																								logger.error('''--pytorch_dump_folder_path ({}) should be a directory'''.format(lowerCAmelCase__  )  )
																								return
																		os.makedirs(lowerCAmelCase__ ,  exist_ok=lowerCAmelCase__  )
																		lowerCAmelCase_      			=  target_dict.indices

																		# fairseq has the <pad> and <s> switched
																		lowerCAmelCase_      			=  0
																		lowerCAmelCase_      			=  1
																		with open(lowerCAmelCase__ ,  '''w''' ,  encoding='''utf-8'''  ) as vocab_handle:
																								json.dump(lowerCAmelCase__ ,  lowerCAmelCase__  )
																		lowerCAmelCase_      			=  WavaVecaCTCTokenizer(
																		    lowerCAmelCase__ ,  unk_token=target_dict.unk_word ,  pad_token=target_dict.pad_word ,  bos_token=target_dict.bos_word ,  eos_token=target_dict.eos_word ,  word_delimiter_token='''|''' ,  do_lower_case=lowerCAmelCase__ ,  )
																		lowerCAmelCase_      			=  True if config.feat_extract_norm == 'layer' else False
																		lowerCAmelCase_      			=  WavaVecaFeatureExtractor(
																		    feature_size=1 ,  sampling_rate=16000 ,  padding_value=0 ,  do_normalize=lowerCAmelCase__ ,  return_attention_mask=lowerCAmelCase__ ,  )
																		lowerCAmelCase_      			=  WavaVecaProcessor(feature_extractor=lowerCAmelCase__ ,  tokenizer=lowerCAmelCase__  )
																		processor.save_pretrained(lowerCAmelCase__  )

												lowerCAmelCase_      			=  WavaVecaForCTC(lowerCAmelCase__  )
						else:
												lowerCAmelCase_      			=  WavaVecaForPreTraining(lowerCAmelCase__  )

						if is_finetuned or is_seq_class:
												lowerCAmelCase_      			=  fairseq.checkpoint_utils.load_model_ensemble_and_task(
												    [checkpoint_path] ,  arg_overrides={'''data''': '''/'''.join(dict_path.split('''/'''  )[:-1]  )}  )
						else:
												lowerCAmelCase_      			=  argparse.Namespace(task='''audio_pretraining'''  )
												lowerCAmelCase_      			=  fairseq.tasks.setup_task(lowerCAmelCase__  )

												lowerCAmelCase_      			=  fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path] ,  task=lowerCAmelCase__  )

						lowerCAmelCase_      			=  model[0].eval()

						recursively_load_weights(lowerCAmelCase__ ,  lowerCAmelCase__ ,  not is_finetuned  )

						hf_wavavec.save_pretrained(lowerCAmelCase__  )


if __name__ == "__main__":
						_A     =			argparse.ArgumentParser()
						parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
						parser.add_argument('''--checkpoint_path''', default=None, type=str, help='''Path to fairseq checkpoint''')
						parser.add_argument('''--dict_path''', default=None, type=str, help='''Path to dict of fine-tuned model''')
						parser.add_argument('''--config_path''', default=None, type=str, help='''Path to hf config.json of model to convert''')
						parser.add_argument(
						    '''--not_finetuned''', action='''store_true''', help='''Whether the model to convert is a fine-tuned model or not'''
						)
						parser.add_argument(
						    '''--is_seq_class''',
						    action='''store_true''',
						    help='''Whether the model to convert is a fine-tuned sequence classification model or not''',
						)
						_A     =			parser.parse_args()

						_A     =			not args.not_finetuned and not args.is_seq_class
						convert_wavaveca_checkpoint(
						    args.checkpoint_path,
						    args.pytorch_dump_folder_path,
						    args.config_path,
						    args.dict_path,
						    is_finetuned,
						    args.is_seq_class,
						)









"""simple docstring"""







import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


lowercase__		:				List[Any]          =	logging.getLogger(__name__)

def 				UpperCamelCase_    (    )						->   Dict:

  """simple docstring"""




  lowerCAmelCase_   :				Tuple								= argparse.ArgumentParser(
      description='Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset.'						)
  parser.add_argument(
      '--dataset_name'    ,   type=lowerCAmelCase__    ,   default='wikitext'    ,   help='Name of the training. Explore datasets at: hf.co/datasets.'    ,   )
  parser.add_argument(
      '--dataset_config'    ,   type=lowerCAmelCase__    ,   default='wikitext-103-raw-v1'    ,   help='Configuration name of the dataset.'						)
  parser.add_argument(
      '--tokenizer_name_or_path'    ,   type=lowerCAmelCase__    ,   default='sayakpaul/unigram-tokenizer-wikitext'    ,   help='Tokenizer identifier. Can be a local filepath or a Hub identifier.'    ,   )
  parser.add_argument(
      '--shard_size'    ,   type=lowerCAmelCase__    ,   default=1000    ,   help='Number of entries to go in a single shard.'    ,   )
  parser.add_argument('--split'    ,   type=lowerCAmelCase__    ,   default='train'    ,   choices=['train', 'test', 'validation']						)
  parser.add_argument(
      '--limit'    ,   default=lowerCAmelCase__    ,   type=lowerCAmelCase__    ,   help='Limit the number of shards (used for debugging).'    ,   )
  parser.add_argument(
      '--max_length'    ,   type=lowerCAmelCase__    ,   default=512    ,   help='Maximum sequence length. For training on TPUs, it helps to have a maximum'
      ' sequence length that is a multiple of 8.'    ,   )
  parser.add_argument(
      '--output_dir'    ,   default='tf-tpu'    ,   type=lowerCAmelCase__    ,   help='Output directory where the TFRecord shards will be saved. If the'
      ' path is appended with `gs://` (\'gs://tf-tpu\', for example) then the TFRecord'
      ' shards will be directly saved to a Google Cloud Storage bucket.'    ,   )

  lowerCAmelCase_   :				List[Any]								= parser.parse_args()
  return args

def 				UpperCamelCase_    (    lowerCAmelCase__      :				int						)						->   Union[str, Any]:

  """simple docstring"""




  def fn(lowerCAmelCase__      :				Optional[Any]						):
    return tokenizer(examples['text']						)

  return fn

def 				UpperCamelCase_    (    lowerCAmelCase__      :				Tuple						)						->   Dict:

  """simple docstring"""




  lowerCAmelCase_   :				int								= []
  for i in range(len(tokenized_data['input_ids']						)						):
    lowerCAmelCase_   :				Tuple								= {
        'input_ids': tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data['input_ids'][i]						)						),
        'attention_mask': tf.train.Feature(
            intaa_list=tf.train.IntaaList(value=tokenized_data['attention_mask'][i]						)						),
    }
    lowerCAmelCase_   :				Union[str, Any]								= tf.train.Features(feature=lowerCAmelCase__						)
    lowerCAmelCase_   :				Dict								= tf.train.Example(features=lowerCAmelCase__						)
    lowerCAmelCase_   :				Optional[int]								= example.SerializeToString()
    records.append(lowerCAmelCase__						)
  return records

def 				UpperCamelCase_    (    lowerCAmelCase__      :				Optional[Any]						)						->   Optional[int]:

  """simple docstring"""




  lowerCAmelCase_   :				str								= datasets.load_dataset(args.dataset_name    ,   args.dataset_config    ,   split=args.split						)

  if args.limit is not None:
    lowerCAmelCase_   :				Tuple								= min(len(lowerCAmelCase__						)    ,   args.limit						)
    lowerCAmelCase_   :				Any								= dataset.select(range(lowerCAmelCase__						)						)
    print(f"Limiting the dataset to {args.limit} entries."						)

  lowerCAmelCase_   :				List[str]								= AutoTokenizer.from_pretrained(args.tokenizer_name_or_path						)

  # Handle output directory creation.
  # For serializing into a Google Cloud Storage Bucket, one needs to first
  # create a bucket.
  if "gs" not in args.output_dir:
    if not os.path.exists(args.output_dir						):
      os.makedirs(args.output_dir						)
    lowerCAmelCase_   :				int								= os.path.join(args.output_dir    ,   args.split						)
    if not os.path.exists(lowerCAmelCase__						):
      os.makedirs(lowerCAmelCase__						)
  else:
    lowerCAmelCase_   :				Dict								= os.path.join(args.output_dir    ,   args.split						)

  # Tokenize the whole dataset at once.
  lowerCAmelCase_   :				Dict								= tokenize_function(lowerCAmelCase__						)
  lowerCAmelCase_   :				List[Any]								= dataset.map(lowerCAmelCase__    ,   batched=lowerCAmelCase__    ,   num_proc=4    ,   remove_columns=['text']						)

  # We need to concatenate all our texts together, and then split the result
  # into chunks of a fixed size, which we will call block_size. To do this, we
  # will use the map method again, with the option batched=True. When we use batched=True,
  # the function we pass to map() will be passed multiple inputs at once, allowing us
  # to group them into more or fewer examples than we had in the input.
  # This allows us to create our new fixed-length samples. The advantage of this
  # method is that we don't lose a whole lot of content from the dataset compared to the
  # case where we simply tokenize with a pre-defined max_length.

  def group_texts(lowerCAmelCase__      :				List[Any]						):
    # Concatenate all texts.
    lowerCAmelCase_   :				int								= {k: sum(examples[k]    ,   []						) for k in examples.keys()}
    lowerCAmelCase_   :				Any								= len(concatenated_examples[list(examples.keys()						)[0]]						)
    # We drop the small remainder, though you could add padding instead if the model supports it
    # In this, as in all things, we advise you to follow your heart 🫀
    lowerCAmelCase_   :				Union[str, Any]								= (total_length // args.max_length) * args.max_length
    # Split by chunks of max_len.
    lowerCAmelCase_   :				Optional[Any]								= {
        k: [t[i : i + args.max_length] for i in range(0    ,   lowerCAmelCase__    ,   args.max_length						)]
        for k, t in concatenated_examples.items()
    }
    return result

  lowerCAmelCase_   :				Optional[int]								= dataset_tokenized.map(lowerCAmelCase__    ,   batched=lowerCAmelCase__    ,   batch_size=1000    ,   num_proc=4						)

  lowerCAmelCase_   :				Optional[int]								= 0
  lowerCAmelCase_   :				Optional[Any]								= 0
  for shard in range(0    ,   len(lowerCAmelCase__						)    ,   args.shard_size						):
    lowerCAmelCase_   :				Dict								= grouped_dataset[shard : shard + args.shard_size]
    lowerCAmelCase_   :				Tuple								= len(dataset_snapshot['input_ids']						)
    lowerCAmelCase_   :				Optional[Any]								= os.path.join(lowerCAmelCase__    ,   f"dataset-{shard_count}-{records_containing}.tfrecord"						)
    lowerCAmelCase_   :				Tuple								= get_serialized_examples(lowerCAmelCase__						)

    with tf.io.TFRecordWriter(lowerCAmelCase__						) as out_file:
      for i in range(len(lowerCAmelCase__						)						):
        lowerCAmelCase_   :				Dict								= serialized_examples[i]
        out_file.write(lowerCAmelCase__						)
      print('Wrote file {} containing {} records'.format(lowerCAmelCase__    ,   lowerCAmelCase__						)						)

    shard_count += 1
    total_records += records_containing

  with open(f"split-{args.split}-records-count.txt"    ,   'w'						) as f:
    print(f"Total {args.split} records: {total_records}"    ,   file=lowerCAmelCase__						)


if __name__ == "__main__":
     lowercase__		:				int          =	parse_args()
     main(args)





