
import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging


logging.set_verbosity_info()
_UpperCAmelCase      :				Tuple          = logging.get_logger(__name__)

_UpperCAmelCase      :				str          = {
    "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",
    "encoder.layer_norm_for_extract": "layer_norm_for_extract",
    "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",
    "label_embs_concat": "label_embeddings_concat",
    "mask_emb": "masked_spec_embed",
    "spk_proj": "speaker_proj",
}
_UpperCAmelCase      :				Any          = [
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
    "label_embeddings_concat",
    "speaker_proj",
    "layer_norm_for_extract",
]


def    A (      lowercase					, lowercase					, lowercase					, lowercase					, lowercase				)     ->						Optional[int]:
		'''simple docstring'''

		for attribute in key.split('.'				):
				UpperCamelCase						=  getattr(lowercase					, lowercase				)

		if weight_type is not None:
				UpperCamelCase						=  getattr(lowercase					, lowercase				).shape
		else:
				UpperCamelCase						=  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":
				UpperCamelCase						=  value
		elif weight_type == "weight_g":
				UpperCamelCase						=  value
		elif weight_type == "weight_v":
				UpperCamelCase						=  value
		elif weight_type == "bias":
				UpperCamelCase						=  value
		else:
				UpperCamelCase						=  value

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


def    A (      lowercase					, lowercase				)     ->						Union[str, Any]:
		'''simple docstring'''

		UpperCamelCase						=  []
		UpperCamelCase						=  fairseq_model.state_dict()

		UpperCamelCase						=  hf_model.unispeech_sat.feature_extractor

		for name, value in fairseq_dict.items():
				UpperCamelCase						=  False
				if "conv_layers" in name:
						load_conv_layer(
						    lowercase					, lowercase					, lowercase					, lowercase					, hf_model.config.feat_extract_norm == 'group'					, )
						UpperCamelCase						=  True
				else:
						for key, mapped_key in MAPPING.items():
								UpperCamelCase						=  'unispeech_sat.' + 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]:
										if "layer_norm_for_extract" in name and (".".join(name.split('.'				)[:-1]				) != key):
												# special case since naming is very similar
												continue
										UpperCamelCase						=  True
										if "*" in mapped_key:
												UpperCamelCase						=  name.split(lowercase				)[0].split('.'				)[-2]
												UpperCamelCase						=  mapped_key.replace('*'					, lowercase				)
										if "weight_g" in name:
												UpperCamelCase						=  'weight_g'
										elif "weight_v" in name:
												UpperCamelCase						=  'weight_v'
										elif "bias" in name:
												UpperCamelCase						=  'bias'
										elif "weight" in name:
												# TODO: don't match quantizer.weight_proj
												UpperCamelCase						=  'weight'
										else:
												UpperCamelCase						=  None
										set_recursively(lowercase					, lowercase					, lowercase					, lowercase					, lowercase				)
								continue
				if not is_used:
						unused_weights.append(lowercase				)

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


def    A (      lowercase					, lowercase					, lowercase					, lowercase					, lowercase				)     ->						Optional[int]:
		'''simple docstring'''

		UpperCamelCase						=  full_name.split('conv_layers.'				)[-1]
		UpperCamelCase						=  name.split('.'				)
		UpperCamelCase						=  int(items[0]				)
		UpperCamelCase						=  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.'''				)
						UpperCamelCase						=  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.'''				)
						UpperCamelCase						=  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[layer_id].layer_norm.bias.data.shape} was found.'''				)
						UpperCamelCase						=  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[layer_id].layer_norm.weight.data.shape} was found.'''				)
						UpperCamelCase						=  value
						logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''				)
		else:
				unused_weights.append(lowercase				)


@torch.no_grad()
def    A (      lowercase					, lowercase					, lowercase=None					, lowercase=None					, lowercase=True				)     ->						List[Any]:
		'''simple docstring'''

		if config_path is not None:
				UpperCamelCase						=  UniSpeechSatConfig.from_pretrained(lowercase				)
		else:
				UpperCamelCase						=  UniSpeechSatConfig()

		UpperCamelCase						=  ''

		if is_finetuned:
				UpperCamelCase						=  UniSpeechSatForCTC(lowercase				)
		else:
				UpperCamelCase						=  UniSpeechSatForPreTraining(lowercase				)

		UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase						=  fairseq.checkpoint_utils.load_model_ensemble_and_task(
		    [checkpoint_path]					, arg_overrides={'data': '/'.join(dict_path.split('/'				)[:-1]				)}				)
		UpperCamelCase						=  model[0].eval()

		recursively_load_weights(lowercase					, lowercase				)

		hf_wavavec.save_pretrained(lowercase				)


if __name__ == "__main__":
					_UpperCAmelCase      :				List[str]          = 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"
					)
					_UpperCAmelCase      :				Dict          = parser.parse_args()
					convert_unispeech_sat_checkpoint(
					    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
					)


import os
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from huggingface_hub.file_download import http_get
from requests.exceptions import HTTPError

from transformers import (
    AlbertTokenizer,
    AutoTokenizer,
    BertTokenizer,
    BertTokenizerFast,
    GPTaTokenizerFast,
    is_tokenizers_available,
)
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_tokenizers
from transformers.tokenization_utils import Trie


sys.path.append(str(Path(__file__).parent.parent / "utils"))

from test_module.custom_tokenization import CustomTokenizer  # noqa E402


if is_tokenizers_available():
					from test_module.custom_tokenization_fast import CustomTokenizerFast


class        lowercase      (	unittest.TestCase     ):



		def __UpperCamelCase				(				self )						->      Tuple:
				"""simple docstring"""
				# A mock response for an HTTP head request to emulate server down
				UpperCamelCase						=  mock.Mock()
				UpperCamelCase						=  500
				UpperCamelCase						=  {}
				UpperCamelCase						=  HTTPError
				UpperCamelCase						=  {}

				# Download this model to make sure it's in the cache.
				UpperCamelCase						=  BertTokenizer.from_pretrained('hf-internal-testing/tiny-random-bert' )

				# Under the mock environment we get a 500 error when trying to reach the tokenizer.
				with mock.patch('requests.Session.request'       ,   return_value=A_ ) as mock_head:
						UpperCamelCase						=  BertTokenizer.from_pretrained('hf-internal-testing/tiny-random-bert' )
						# This check we did call the fake head request
						mock_head.assert_called()



		@require_tokenizers
		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				# A mock response for an HTTP head request to emulate server down
				UpperCamelCase						=  mock.Mock()
				UpperCamelCase						=  500
				UpperCamelCase						=  {}
				UpperCamelCase						=  HTTPError
				UpperCamelCase						=  {}

				# Download this model to make sure it's in the cache.
				UpperCamelCase						=  GPTaTokenizerFast.from_pretrained('gpt2' )

				# Under the mock environment we get a 500 error when trying to reach the tokenizer.
				with mock.patch('requests.Session.request'       ,   return_value=A_ ) as mock_head:
						UpperCamelCase						=  GPTaTokenizerFast.from_pretrained('gpt2' )
						# This check we did call the fake head request
						mock_head.assert_called()



		def __UpperCamelCase				(				self )						->      Optional[Any]:
				"""simple docstring"""
				# This test is for deprecated behavior and can be removed in v5
				try:
						UpperCamelCase						=  tempfile.mktemp()
						with open(A_       ,   'wb' ) as f:
								http_get('https://huggingface.co/albert-base-v1/resolve/main/spiece.model'       ,   A_ )

						UpperCamelCase						=  AlbertTokenizer.from_pretrained(A_ )
				finally:
						os.remove(A_ )

				# Supporting this legacy load introduced a weird bug where the tokenizer would load local files if they are in
				# the current folder and have the right name.
				if os.path.isfile('tokenizer.json' ):
						# We skip the test if the user has a `tokenizer.json` in this folder to avoid deleting it.
						return
				try:
						with open('tokenizer.json'       ,   'wb' ) as f:
								http_get('https://huggingface.co/hf-internal-testing/tiny-random-bert/blob/main/tokenizer.json'       ,   A_ )
						UpperCamelCase						=  AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2' )
						# The tiny random BERT has a vocab size of 1024, tiny gpt2 as a vocab size of 1000
						self.assertEqual(tokenizer.vocab_size       ,   1_000 )
						# Tokenizer should depend on the remote checkpoint, not the local tokenizer.json file.

				finally:
						os.remove('tokenizer.json' )



		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				# This test is for deprecated behavior and can be removed in v5
				UpperCamelCase						=  AlbertTokenizer.from_pretrained('https://huggingface.co/albert-base-v1/resolve/main/spiece.model' )






@is_staging_test
class        lowercase      (	unittest.TestCase     ):
		__lowercase : int			  =       ["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]", "bla", "blou"]



		@classmethod
		def __UpperCamelCase				(				cls )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  TOKEN
				HfFolder.save_token(A_ )



		@classmethod
		def __UpperCamelCase				(				cls )						->      Optional[int]:
				"""simple docstring"""
				try:
						delete_repo(token=cls._token       ,   repo_id='test-tokenizer' )
				except HTTPError:
						pass

				try:
						delete_repo(token=cls._token       ,   repo_id='valid_org/test-tokenizer-org' )
				except HTTPError:
						pass

				try:
						delete_repo(token=cls._token       ,   repo_id='test-dynamic-tokenizer' )
				except HTTPError:
						pass



		def __UpperCamelCase				(				self )						->      Any:
				"""simple docstring"""
				with tempfile.TemporaryDirectory() as tmp_dir:
						UpperCamelCase						=  os.path.join(A_       ,   'vocab.txt' )
						with open(A_       ,   'w'       ,   encoding='utf-8' ) as vocab_writer:
								vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens] ) )
						UpperCamelCase						=  BertTokenizer(A_ )

				tokenizer.push_to_hub('test-tokenizer'       ,   use_auth_token=self._token )
				UpperCamelCase						=  BertTokenizer.from_pretrained(F'''{USER}/test-tokenizer''' )
				self.assertDictEqual(new_tokenizer.vocab       ,   tokenizer.vocab )

				# Reset repo
				delete_repo(token=self._token       ,   repo_id='test-tokenizer' )

				# Push to hub via save_pretrained
				with tempfile.TemporaryDirectory() as tmp_dir:
						tokenizer.save_pretrained(A_       ,   repo_id='test-tokenizer'       ,   push_to_hub=A_       ,   use_auth_token=self._token )

				UpperCamelCase						=  BertTokenizer.from_pretrained(F'''{USER}/test-tokenizer''' )
				self.assertDictEqual(new_tokenizer.vocab       ,   tokenizer.vocab )



		def __UpperCamelCase				(				self )						->      int:
				"""simple docstring"""
				with tempfile.TemporaryDirectory() as tmp_dir:
						UpperCamelCase						=  os.path.join(A_       ,   'vocab.txt' )
						with open(A_       ,   'w'       ,   encoding='utf-8' ) as vocab_writer:
								vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens] ) )
						UpperCamelCase						=  BertTokenizer(A_ )

				tokenizer.push_to_hub('valid_org/test-tokenizer-org'       ,   use_auth_token=self._token )
				UpperCamelCase						=  BertTokenizer.from_pretrained('valid_org/test-tokenizer-org' )
				self.assertDictEqual(new_tokenizer.vocab       ,   tokenizer.vocab )

				# Reset repo
				delete_repo(token=self._token       ,   repo_id='valid_org/test-tokenizer-org' )

				# Push to hub via save_pretrained
				with tempfile.TemporaryDirectory() as tmp_dir:
						tokenizer.save_pretrained(
						    A_       ,   repo_id='valid_org/test-tokenizer-org'       ,   push_to_hub=A_       ,   use_auth_token=self._token )

				UpperCamelCase						=  BertTokenizer.from_pretrained('valid_org/test-tokenizer-org' )
				self.assertDictEqual(new_tokenizer.vocab       ,   tokenizer.vocab )



		@require_tokenizers
		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				CustomTokenizer.register_for_auto_class()
				with tempfile.TemporaryDirectory() as tmp_dir:
						UpperCamelCase						=  os.path.join(A_       ,   'vocab.txt' )
						with open(A_       ,   'w'       ,   encoding='utf-8' ) as vocab_writer:
								vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens] ) )
						UpperCamelCase						=  CustomTokenizer(A_ )

				# No fast custom tokenizer
				tokenizer.push_to_hub('test-dynamic-tokenizer'       ,   use_auth_token=self._token )

				UpperCamelCase						=  AutoTokenizer.from_pretrained(F'''{USER}/test-dynamic-tokenizer'''       ,   trust_remote_code=A_ )
				# Can't make an isinstance check because the new_model.config is from the CustomTokenizer class of a dynamic module
				self.assertEqual(tokenizer.__class__.__name__       ,   'CustomTokenizer' )

				# Fast and slow custom tokenizer
				CustomTokenizerFast.register_for_auto_class()
				with tempfile.TemporaryDirectory() as tmp_dir:
						UpperCamelCase						=  os.path.join(A_       ,   'vocab.txt' )
						with open(A_       ,   'w'       ,   encoding='utf-8' ) as vocab_writer:
								vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens] ) )

						UpperCamelCase						=  BertTokenizerFast.from_pretrained(A_ )
						bert_tokenizer.save_pretrained(A_ )
						UpperCamelCase						=  CustomTokenizerFast.from_pretrained(A_ )

				tokenizer.push_to_hub('test-dynamic-tokenizer'       ,   use_auth_token=self._token )

				UpperCamelCase						=  AutoTokenizer.from_pretrained(F'''{USER}/test-dynamic-tokenizer'''       ,   trust_remote_code=A_ )
				# Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
				self.assertEqual(tokenizer.__class__.__name__       ,   'CustomTokenizerFast' )
				UpperCamelCase						=  AutoTokenizer.from_pretrained(
				    F'''{USER}/test-dynamic-tokenizer'''       ,   use_fast=A_       ,   trust_remote_code=A_ )
				# Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
				self.assertEqual(tokenizer.__class__.__name__       ,   'CustomTokenizer' )






class        lowercase      (	unittest.TestCase     ):



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('Hello 友達' )
				self.assertEqual(trie.data       ,   {'H': {'e': {'l': {'l': {'o': {' ': {'友': {'達': {'': 1}}}}}}}}} )
				trie.add('Hello' )
				trie.data
				self.assertEqual(trie.data       ,   {'H': {'e': {'l': {'l': {'o': {'': 1, ' ': {'友': {'達': {'': 1}}}}}}}}} )



		def __UpperCamelCase				(				self )						->      str:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				self.assertEqual(trie.split('[CLS] This is a extra_id_100' )       ,   ['[CLS] This is a extra_id_100'] )
				trie.add('[CLS]' )
				trie.add('extra_id_1' )
				trie.add('extra_id_100' )
				self.assertEqual(trie.split('[CLS] This is a extra_id_100' )       ,   ['[CLS]', ' This is a ', 'extra_id_100'] )



		def __UpperCamelCase				(				self )						->      List[Any]:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('A' )
				self.assertEqual(trie.split('ABC' )       ,   ['A', 'BC'] )
				self.assertEqual(trie.split('BCA' )       ,   ['BC', 'A'] )



		def __UpperCamelCase				(				self )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('TOKEN]' )
				trie.add('[SPECIAL_TOKEN]' )
				self.assertEqual(trie.split('This is something [SPECIAL_TOKEN]' )       ,   ['This is something ', '[SPECIAL_TOKEN]'] )



		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('A' )
				trie.add('P' )
				trie.add('[SPECIAL_TOKEN]' )
				self.assertEqual(trie.split('This is something [SPECIAL_TOKEN]' )       ,   ['This is something ', '[SPECIAL_TOKEN]'] )



		def __UpperCamelCase				(				self )						->      Optional[Any]:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('AB' )
				trie.add('B' )
				trie.add('C' )
				self.assertEqual(trie.split('ABC' )       ,   ['AB', 'C'] )



		def __UpperCamelCase				(				self )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('ABC' )
				trie.add('B' )
				trie.add('CD' )
				self.assertEqual(trie.split('ABCD' )       ,   ['ABC', 'D'] )



		def __UpperCamelCase				(				self )						->      int:
				"""simple docstring"""
				# Even if the offsets are wrong, we necessarily output correct string
				# parts.
				UpperCamelCase						=  Trie()
				UpperCamelCase						=  trie.cut_text('ABC'       ,   [0, 0, 2, 1, 2, 3] )
				self.assertEqual(A_       ,   ['AB', 'C'] )
