


'''simple docstring'''







import argparse
import os
from functools import reduce

import fairseq
import torch
from datasets import load_dataset

from transformers import WavaVecaProcessor, logging
from transformers.models.dataavec.configuration_dataavec_audio import DataaVecAudioConfig

# Copied from https://github.com/pytorch/fairseq/blob/main/examples/data2vec/models/data2vec_audio.py
from transformers.models.dataavec.dataavec_audio import DataaVecAudioModel as Dummy  # noqa: F401
from transformers.models.dataavec.modeling_dataavec_audio import DataaVecAudioForCTC, DataaVecAudioModel


logging.set_verbosity_info()
_SCREAMING_SNAKE_CASE       :   List[Any]        = logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE       :   Union[str, Any]        = {
    "post_extract_proj": "feature_projection.projection",
    "models.0.layer_norm": "feature_projection.layer_norm",
    "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",
    "w2v_model.layer_norm": "feature_projection.layer_norm",
    "w2v_encoder.proj": "lm_head",
    "mask_emb": "masked_spec_embed",
}
_SCREAMING_SNAKE_CASE       :   Optional[Any]        = [
    "lm_head",
]



def 						_UpperCamelCase							( UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__		):






				"""simple docstring"""



				for attribute in key.split("."		):
								__magic_name__				:					Tuple						       =       getattr(UpperCamelCase__     ,   UpperCamelCase__		)

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

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



def 						_UpperCamelCase							( UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__		):






				"""simple docstring"""



				__magic_name__				:					Optional[int]						       =       []
				__magic_name__				:					str						       =       fairseq_model.state_dict()

				if not is_headless:
								__magic_name__				:					List[str]						       =       hf_model.dataavec_audio.feature_extractor
								__magic_name__				:					Union[str, Any]						       =       hf_model.dataavec_audio.encoder.pos_conv_embed

				else:
								__magic_name__				:					Tuple						       =       hf_model.feature_extractor
								__magic_name__				:					Union[str, Any]						       =       hf_model.encoder.pos_conv_embed

				for name, value in fairseq_dict.items():
								__magic_name__				:					int						       =       False
								if "conv_layers" in name:
												load_conv_layer(
												    UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   )
												__magic_name__				:					Optional[Any]						       =       True
								elif "pos_conv" in name:
												load_pos_conv_layer(
												    UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   )
												__magic_name__				:					Union[str, Any]						       =       True
								else:
												for key, mapped_key in MAPPING.items():
																if not is_headless:
																				__magic_name__				:					Optional[int]						       =       "data2vec_audio." + 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]:
																				__magic_name__				:					Dict						       =       True
																				if "*" in mapped_key:
																								__magic_name__				:					Optional[Any]						       =       name.split(UpperCamelCase__		)[0].split("."		)[-2]
																								__magic_name__				:					Dict						       =       mapped_key.replace("*"     ,   UpperCamelCase__		)
																				if "weight_g" in name:
																								__magic_name__				:					str						       =       "weight_g"
																				elif "weight_v" in name:
																								__magic_name__				:					List[Any]						       =       "weight_v"
																				elif "bias" in name:
																								__magic_name__				:					Optional[Any]						       =       "bias"
																				elif "weight" in name:
																								# TODO: don't match quantizer.weight_proj
																								__magic_name__				:					int						       =       "weight"
																				else:
																								__magic_name__				:					Any						       =       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 						_UpperCamelCase							( UpperCamelCase__     ,   UpperCamelCase__		):






				"""simple docstring"""



				__magic_name__				:					Tuple						       =       path.split("."		)
				return reduce(UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__		)



def 						_UpperCamelCase							( UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__		):






				"""simple docstring"""



				__magic_name__				:					Dict						       =       access_by_string(UpperCamelCase__     ,   UpperCamelCase__		)
				__magic_name__				:					List[Any]						       =       hf_weight.data

				if fsq_value.shape != hf_value.shape:
								raise ValueError(F"""{full_name} has size {fsq_value.shape}, but {hf_value.shape} was found."""		)
				__magic_name__				:					int						       =       fsq_value
				logger.info(F"""{full_name} was correctly initialized from {hf_weight_path}."""		)



def 						_UpperCamelCase							( UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__		):






				"""simple docstring"""



				__magic_name__				:					Optional[Any]						       =       full_name.split("conv_layers."		)[-1]
				__magic_name__				:					List[str]						       =       name.split("."		)
				__magic_name__				:					Optional[Any]						       =       int(items[0]		)
				__magic_name__				:					Optional[int]						       =       int(items[1]		)

				__magic_name__				:					int						       =       name.split("."		)[-1]
				if type_id == 0:
								__magic_name__				:					Any						       =       "conv"
				elif type_id == 2:
								__magic_name__				:					int						       =       "layer_norm"
				else:
								unused_weights.append(UpperCamelCase__		)
								return

				set_weights(UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   F"""conv_layers.{layer_id}.{layer_type}.{weight_type}"""		)



def 						_UpperCamelCase							( UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__		):






				"""simple docstring"""



				__magic_name__				:					Tuple						       =       full_name.split("pos_conv."		)[-1]
				__magic_name__				:					str						       =       name.split("."		)
				__magic_name__				:					List[str]						       =       int(items[0]		)
				__magic_name__				:					Any						       =       int(items[1]		)

				__magic_name__				:					Dict						       =       name.split("."		)[-1]
				if type_id != 0:
								unused_weights.append(UpperCamelCase__		)
								return
				else:
								__magic_name__				:					List[Any]						       =       "conv"

				set_weights(UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   F"""layers.{layer_id}.{layer_type}.{weight_type}"""		)



@torch.no_grad()
def 						_UpperCamelCase							( UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__=None     ,   UpperCamelCase__=None     ,   UpperCamelCase__=True		):






				"""simple docstring"""



				if config_path is not None:
								__magic_name__				:					List[str]						       =       DataaVecAudioConfig.from_pretrained(UpperCamelCase__		)
				else:
								__magic_name__				:					List[str]						       =       DataaVecAudioConfig()

				if not is_finetuned:
								# Modify final_proj layer name
								__magic_name__				:					Tuple						       =       DataaVecAudioModel(UpperCamelCase__		)
								__magic_name__				:					Dict						       =       os.path.dirname(UpperCamelCase__		)

								__magic_name__				:					Tuple						       =       torch.load(UpperCamelCase__		)
								__magic_name__				:					List[Any]						       =       state_dict["model"].pop("final_proj.0.weight"		)
								__magic_name__				:					Dict						       =       state_dict["model"].pop("final_proj.0.bias"		)
								__magic_name__				:					List[Any]						       =       os.path.join(UpperCamelCase__     ,   "converted.pt"		)
								torch.save(UpperCamelCase__     ,   UpperCamelCase__		)
				else:
								__magic_name__				:					Optional[Any]						       =       DataaVecAudioForCTC(UpperCamelCase__		)
								__magic_name__				:					Optional[Any]						       =       checkpoint_path

				def load_dataavec(UpperCamelCase__		):
								__magic_name__       ,     __magic_name__       ,     __magic_name__				:					int						       =       fairseq.checkpoint_utils.load_model_ensemble_and_task([path]		)
								return model[0].eval()

				__magic_name__				:					List[str]						       =       load_dataavec(UpperCamelCase__		)

				recursively_load_weights(UpperCamelCase__     ,   UpperCamelCase__     ,   not is_finetuned		)

				__magic_name__				:					str						       =       WavaVecaProcessor.from_pretrained("facebook/wav2vec2-large-lv60"		)

				__magic_name__				:					Optional[Any]						       =       load_dataset("patrickvonplaten/librispeech_asr_dummy"     ,   "clean"     ,   split="validation"		)
				__magic_name__				:					str						       =       [x["array"] for x in ds[:4]["audio"]]

				__magic_name__				:					List[Any]						       =       processor(UpperCamelCase__     ,   return_tensors="pt"     ,   padding=UpperCamelCase__		)

				__magic_name__				:					Tuple						       =       inputs.input_values
				__magic_name__				:					Any						       =       inputs.attention_mask
				#    input_values = inputs.input_values[:, :-1]
				#    attention_mask = inputs.attention_mask[:, :-1]

				hf_wavavec.eval()
				model.eval()
				if is_finetuned:
								__magic_name__				:					Tuple						       =       model(source=UpperCamelCase__     ,   padding_mask=(1 - attention_mask)     ,   mask=UpperCamelCase__     ,   features_only=UpperCamelCase__		)[
								    "encoder_out"
								].transpose(0     ,   1		)
								__magic_name__				:					Any						       =       hf_wavavec(UpperCamelCase__     ,   attention_mask=UpperCamelCase__		)["logits"]

								__magic_name__				:					Optional[int]						       =       torch.argmax(UpperCamelCase__     ,   dim=-1		)
								__magic_name__				:					List[str]						       =       processor.batch_decode(UpperCamelCase__		)

								print(F"""Expected Output: {ds[:4]["text"]}, Pred: {output_string}"""		)
				else:
								__magic_name__				:					Optional[int]						       =       model(source=UpperCamelCase__     ,   padding_mask=(1 - attention_mask)     ,   mask=UpperCamelCase__     ,   features_only=UpperCamelCase__		)[
								    "layer_results"
								][-1][0].transpose(0     ,   1		)
								__magic_name__				:					List[str]						       =       hf_wavavec(UpperCamelCase__     ,   attention_mask=UpperCamelCase__		)["last_hidden_state"]

				print(our_output.shape     ,   their_output.shape		)
				__magic_name__				:					Optional[Any]						       =       torch.max(torch.abs(our_output - their_output		)		).item()
				print(F"""max_absolute_diff = {max_absolute_diff}"""		)  # ~ 1e-7
				__magic_name__				:					Any						       =       torch.allclose(UpperCamelCase__     ,   UpperCamelCase__     ,   atol=1E-3		)
				print("Do both models output the same tensors?"     ,   "🔥" if success else "💩"		)
				if not success:
								raise Exception("Something went wRoNg"		)

				hf_wavavec.save_pretrained(UpperCamelCase__		)

				if is_finetuned:
								processor.save_pretrained(UpperCamelCase__		)
				else:
								processor.feature_extractor.save_pretrained(UpperCamelCase__		)


if __name__ == "__main__":
		_SCREAMING_SNAKE_CASE       :   Any        = 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"
		)
		_SCREAMING_SNAKE_CASE       :   Optional[Any]        = parser.parse_args()
		convert_wavaveca_checkpoint(
		    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
		)



'''simple docstring'''







import json
import os
import unittest

from transformers.models.gpt_neox_japanese.tokenization_gpt_neox_japanese import (
    VOCAB_FILES_NAMES,
    GPTNeoXJapaneseTokenizer,
)
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin






@require_tokenizers
class 		_snake_case    (				snake_case_  ,  unittest.TestCase		):






						'''simple docstring'''







						__snake_case   =      GPTNeoXJapaneseTokenizer
						__snake_case   =      False
						__snake_case   =      {"do_clean_text": False, "add_prefix_space": False}







						def lowerCAmelCase__ (  self:	Any     )							->	List[str]:
										super().setUp()

										__magic_name__				:					Union[str, Any]						       =       [
										    "こん",
										    "こんに",
										    "にちは",
										    "ばんは",
										    "世界,㔺界",
										    "、",
										    "。",
										    "<BR>",
										    "<SP>",
										    "<TAB>",
										    "<URL>",
										    "<EMAIL>",
										    "<TEL>",
										    "<DATE>",
										    "<PRICE>",
										    "<BLOCK>",
										    "<KIGOU>",
										    "<U2000U2BFF>",
										    "<|emoji1|>",
										    "<unk>",
										    "<|startoftext|>",
										    "<|endoftext|>",
										]
										__magic_name__				:					List[str]						       =       {"emoji": {"\ud83d\ude00": "<|emoji1|>"}, "emoji_inv": {"<|emoji1|>": "\ud83d\ude00"}}  # 😀
										__magic_name__				:					str						       =       {"unk_token": "<unk>"}

										__magic_name__				:					Dict						       =       os.path.join(self.tmpdirname      ,							VOCAB_FILES_NAMES["vocab_file"]     )
										__magic_name__				:					Dict						       =       os.path.join(self.tmpdirname      ,							VOCAB_FILES_NAMES["emoji_file"]     )
										with open(self.vocab_file      ,							"w"      ,							encoding="utf-8"     ) as vocab_writer:
														vocab_writer.write("".join([x + "\n" for x in vocab_tokens]     )     )
										with open(self.emoji_file      ,							"w"     ) as emoji_writer:
														emoji_writer.write(json.dumps(__UpperCamelCase     )     )







						def lowerCAmelCase__ (  self:	Tuple      ,							**__UpperCamelCase:	str     )							->	Union[str, Any]:
										kwargs.update(self.special_tokens_map     )
										return GPTNeoXJapaneseTokenizer.from_pretrained(self.tmpdirname      ,							**__UpperCamelCase     )







						def lowerCAmelCase__ (  self:	List[str]      ,							__UpperCamelCase:	str     )							->	List[str]:
										__magic_name__				:					str						       =       "こんにちは、世界。 \nこんばんは、㔺界。😀"
										__magic_name__				:					str						       =       "こんにちは、世界。 \nこんばんは、世界。😀"
										return input_text, output_text







						def lowerCAmelCase__ (  self:	Union[str, Any]      ,							__UpperCamelCase:	List[Any]     )							->	Dict:
										__magic_name__       ,     __magic_name__				:					Union[str, Any]						       =       self.get_input_output_texts(__UpperCamelCase     )
										__magic_name__				:					List[Any]						       =       tokenizer.encode(__UpperCamelCase      ,							add_special_tokens=__UpperCamelCase     )
										__magic_name__				:					Optional[Any]						       =       tokenizer.decode(__UpperCamelCase      ,							clean_up_tokenization_spaces=__UpperCamelCase     )
										return text, ids







						def lowerCAmelCase__ (  self:	Optional[int]     )							->	Union[str, Any]:
										pass  # TODO add if relevant







						def lowerCAmelCase__ (  self:	str     )							->	List[str]:
										pass  # TODO add if relevant







						def lowerCAmelCase__ (  self:	str     )							->	str:
										pass  # TODO add if relevant







						def lowerCAmelCase__ (  self:	Union[str, Any]     )							->	Dict:
										__magic_name__				:					List[str]						       =       self.get_tokenizer()

										# Testing tokenization
										__magic_name__				:					Any						       =       "こんにちは、世界。　こんばんは、㔺界。"
										__magic_name__				:					Optional[int]						       =       ["こん", "にちは", "、", "世界", "。", "<SP>", "こん", "ばんは", "、", "㔺界", "。"]
										__magic_name__				:					List[Any]						       =       tokenizer.tokenize(__UpperCamelCase     )
										self.assertListEqual(__UpperCamelCase      ,							__UpperCamelCase     )

										# Testing conversion to ids without special tokens
										__magic_name__				:					int						       =       [0, 2, 5, 4, 6, 8, 0, 3, 5, 4, 6]
										__magic_name__				:					str						       =       tokenizer.convert_tokens_to_ids(__UpperCamelCase     )
										self.assertListEqual(__UpperCamelCase      ,							__UpperCamelCase     )

										# Testing conversion to ids with special tokens
										__magic_name__				:					str						       =       tokens + [tokenizer.unk_token]
										__magic_name__				:					Tuple						       =       [0, 2, 5, 4, 6, 8, 0, 3, 5, 4, 6, 19]
										__magic_name__				:					Dict						       =       tokenizer.convert_tokens_to_ids(__UpperCamelCase     )
										self.assertListEqual(__UpperCamelCase      ,							__UpperCamelCase     )







						@slow
						def lowerCAmelCase__ (  self:	List[Any]     )							->	Union[str, Any]:
										__magic_name__				:					List[str]						       =       self.tokenizer_class.from_pretrained("abeja/gpt-neox-japanese-2.7b"     )

										__magic_name__				:					List[Any]						       =       tokenizer.encode("ありがとう。"      ,							add_special_tokens=__UpperCamelCase     )
										__magic_name__				:					Optional[Any]						       =       tokenizer.encode("どういたしまして。"      ,							add_special_tokens=__UpperCamelCase     )

										__magic_name__				:					List[Any]						       =       tokenizer.build_inputs_with_special_tokens(__UpperCamelCase     )
										__magic_name__				:					str						       =       tokenizer.build_inputs_with_special_tokens(__UpperCamelCase      ,							__UpperCamelCase     )

										assert encoded_sentence == ids_a
										assert encoded_pair == ids_a + ids_a







						def lowerCAmelCase__ (  self:	Tuple     )							->	Dict:
										# Intentionally convert some words to accommodate character fluctuations unique to Japanese
										pass







						def lowerCAmelCase__ (  self:	Dict     )							->	Optional[int]:
										# tokenizer has no padding token
										pass