




'''simple docstring'''





import re
import tempfile
import unittest
from pathlib import Path

from datasets.utils.metadata import DatasetMetadata





def 	__a	(   _UpperCamelCase:    str    )						-> Optional[int]:

     """simple docstring"""




     _snake_case							=      min(re.search("^ +"					,						UpperCAmelCase__    ).end() if t.startswith(" "    ) else 0 for t in string.splitlines()    )
     return "\n".join([line[indent_level:] for line in string.splitlines() if indent_level < len(UpperCAmelCase__    )]    )


UpperCamelCase_						:			Any  	=						'''\
---
language:
- zh
- en
task_ids:
- sentiment-classification
---
# Begin of markdown

Some cool dataset card
'''

UpperCamelCase_						:			Optional[Any]  	=						'''\
---
---
# Begin of markdown

Some cool dataset card
'''


UpperCamelCase_						:			str  	=						'''\
# Begin of markdown

Some cool dataset card
'''





class    _a				(						unittest.TestCase							):






     def 				_lowercase	( self    )							->			List[str]:

          with tempfile.TemporaryDirectory() as tmp_dir:
               _snake_case							=      Path(_UpperCamelCase    ) / "README.md"
               with open(_UpperCamelCase					,"w+"    ) as readme_file:
                    readme_file.write(_UpperCamelCase    )
               _snake_case							=      DatasetMetadata.from_readme(_UpperCamelCase    )
               self.assertDictEqual(_UpperCamelCase					,{"language": ["zh", "en"], "task_ids": ["sentiment-classification"]}    )

               with open(_UpperCamelCase					,"w+"    ) as readme_file:
                    readme_file.write(_UpperCamelCase    )
               _snake_case							=      DatasetMetadata.from_readme(_UpperCamelCase    )
               self.assertDictEqual(_UpperCamelCase					,{}    )

               with open(_UpperCamelCase					,"w+"    ) as readme_file:
                    readme_file.write(_UpperCamelCase    )
               _snake_case							=      DatasetMetadata.from_readme(_UpperCamelCase    )
               self.assertEqual(_UpperCamelCase					,{}    )






     def 				_lowercase	( self    )							->			Union[str, Any]:
          _snake_case							=      _dedent(
              "\\n            annotations_creators:\n            - found\n            language_creators:\n            - found\n            language:\n            - en\n            license:\n            - unknown\n            multilinguality:\n            - monolingual\n            pretty_name: Test Dataset\n            size_categories:\n            - 10K<n<100K\n            source_datasets:\n            - extended|other-yahoo-webscope-l6\n            task_categories:\n            - question-answering\n            task_ids:\n            - open-domain-qa\n            "    )
          assert DatasetMetadata.from_yaml_string(_UpperCamelCase    )

          _snake_case							=      _dedent(
              "\\n            annotations_creators:\n            - found\n            language:\n            - en\n            license:\n            - unknown\n            multilinguality:\n            - monolingual\n            pretty_name: Test Dataset\n            size_categories:\n            - 10K<n<100K\n            source_datasets:\n            - extended|other-yahoo-webscope-l6\n            task_categories:\n            - question-answering\n            task_ids:\n            - open-domain-qa\n            task_ids:\n            - open-domain-qa\n            "    )
          with self.assertRaises(_UpperCamelCase    ):
               DatasetMetadata.from_yaml_string(_UpperCamelCase    )

          _snake_case							=      _dedent(
              "\\n            annotations_creators:\n            - found\n            language_creators:\n            - found\n            language:\n            - en\n            license:\n            - unknown\n            multilinguality:\n            - monolingual\n            pretty_name: Test Dataset\n            size_categories:\n            - 10K<n<100K\n            source_datasets:\n            - extended|other-yahoo-webscope-l6\n            task_categories:\n            - text-classification\n            task_ids:\n            - multi-class-classification\n            paperswithcode_id:\n            - squad\n            configs:\n            - en\n            train-eval-index:\n            - config: en\n              task: text-classification\n              task_id: multi_class_classification\n              splits:\n                train_split: train\n                eval_split: test\n              col_mapping:\n                text: text\n                label: target\n              metrics:\n                - type: accuracy\n                  name: Accuracy\n            extra_gated_prompt: |\n              By clicking on “Access repository” below, you also agree to ImageNet Terms of Access:\n              [RESEARCHER_FULLNAME] (the \"Researcher\") has requested permission to use the ImageNet database (the \"Database\") at Princeton University and Stanford University. In exchange for such permission, Researcher hereby agrees to the following terms and conditions:\n              1. Researcher shall use the Database only for non-commercial research and educational purposes.\n            extra_gated_fields:\n              Company: text\n              Country: text\n              I agree to use this model for non-commerical use ONLY: checkbox\n            "    )
          assert DatasetMetadata.from_yaml_string(_UpperCamelCase    )



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()
_lowerCamelCase    :      List[Any]	 =						logging.get_logger(__name__)

_lowerCamelCase    :      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''',
}
_lowerCamelCase    :      int	 =						[
    '''lm_head''',
]





def 		__lowerCamelCase							(UpperCAmelCase__   :							Optional[int]					,       UpperCAmelCase__   :							List[str]					,       UpperCAmelCase__   :							Optional[int]					,       UpperCAmelCase__   :							Tuple					,       UpperCAmelCase__   :							Tuple     ):
					for attribute in key.split("."     ):
										SCREAMING_SNAKE_CASE				=				getattr(UpperCAmelCase__					,       UpperCAmelCase__     )

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

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





def 		__lowerCamelCase							(UpperCAmelCase__   :							Union[str, Any]					,       UpperCAmelCase__   :							Optional[Any]					,       UpperCAmelCase__   :							int     ):
					SCREAMING_SNAKE_CASE				=				[]
					SCREAMING_SNAKE_CASE				=				fairseq_model.state_dict()

					if not is_headless:
										SCREAMING_SNAKE_CASE				=				hf_model.dataavec_audio.feature_extractor
										SCREAMING_SNAKE_CASE				=				hf_model.dataavec_audio.encoder.pos_conv_embed

					else:
										SCREAMING_SNAKE_CASE				=				hf_model.feature_extractor
										SCREAMING_SNAKE_CASE				=				hf_model.encoder.pos_conv_embed

					for name, value in fairseq_dict.items():
										SCREAMING_SNAKE_CASE				=				False
										if "conv_layers" in name:
															load_conv_layer(
															    UpperCAmelCase__					,       UpperCAmelCase__					,       UpperCAmelCase__					,       UpperCAmelCase__					,       )
															SCREAMING_SNAKE_CASE				=				True
										elif "pos_conv" in name:
															load_pos_conv_layer(
															    UpperCAmelCase__					,       UpperCAmelCase__					,       UpperCAmelCase__					,       UpperCAmelCase__					,       )
															SCREAMING_SNAKE_CASE				=				True
										else:
															for key, mapped_key in MAPPING.items():
																				if not is_headless:
																									SCREAMING_SNAKE_CASE				=				"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]:
																									SCREAMING_SNAKE_CASE				=				True
																									if "*" in mapped_key:
																														SCREAMING_SNAKE_CASE				=				name.split(UpperCAmelCase__     )[0].split("."     )[-2]
																														SCREAMING_SNAKE_CASE				=				mapped_key.replace("*"					,       UpperCAmelCase__     )
																									if "weight_g" in name:
																														SCREAMING_SNAKE_CASE				=				"weight_g"
																									elif "weight_v" in name:
																														SCREAMING_SNAKE_CASE				=				"weight_v"
																									elif "bias" in name:
																														SCREAMING_SNAKE_CASE				=				"bias"
																									elif "weight" in name:
																														# TODO: don't match quantizer.weight_proj
																														SCREAMING_SNAKE_CASE				=				"weight"
																									else:
																														SCREAMING_SNAKE_CASE				=				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__   :							Any					,       UpperCAmelCase__   :							List[str]     ):
					SCREAMING_SNAKE_CASE				=				path.split("."     )
					return reduce(UpperCAmelCase__					,       UpperCAmelCase__					,       UpperCAmelCase__     )





def 		__lowerCamelCase							(UpperCAmelCase__   :							Optional[int]					,       UpperCAmelCase__   :							List[Any]					,       UpperCAmelCase__   :							List[Any]					,       UpperCAmelCase__   :							List[str]     ):
					SCREAMING_SNAKE_CASE				=				access_by_string(UpperCAmelCase__					,       UpperCAmelCase__     )
					SCREAMING_SNAKE_CASE				=				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."     )
					SCREAMING_SNAKE_CASE				=				fsq_value
					logger.info(F"{full_name} was correctly initialized from {hf_weight_path}."     )





def 		__lowerCamelCase							(UpperCAmelCase__   :							Optional[Any]					,       UpperCAmelCase__   :							Optional[int]					,       UpperCAmelCase__   :							Dict					,       UpperCAmelCase__   :							int     ):
					SCREAMING_SNAKE_CASE				=				full_name.split("conv_layers."     )[-1]
					SCREAMING_SNAKE_CASE				=				name.split("."     )
					SCREAMING_SNAKE_CASE				=				int(items[0]     )
					SCREAMING_SNAKE_CASE				=				int(items[1]     )

					SCREAMING_SNAKE_CASE				=				name.split("."     )[-1]
					if type_id == 0:
										SCREAMING_SNAKE_CASE				=				"conv"
					elif type_id == 2:
										SCREAMING_SNAKE_CASE				=				"layer_norm"
					else:
										unused_weights.append(UpperCAmelCase__     )
										return

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





def 		__lowerCamelCase							(UpperCAmelCase__   :							str					,       UpperCAmelCase__   :							Tuple					,       UpperCAmelCase__   :							int					,       UpperCAmelCase__   :							Optional[int]     ):
					SCREAMING_SNAKE_CASE				=				full_name.split("pos_conv."     )[-1]
					SCREAMING_SNAKE_CASE				=				name.split("."     )
					SCREAMING_SNAKE_CASE				=				int(items[0]     )
					SCREAMING_SNAKE_CASE				=				int(items[1]     )

					SCREAMING_SNAKE_CASE				=				name.split("."     )[-1]
					if type_id != 0:
										unused_weights.append(UpperCAmelCase__     )
										return
					else:
										SCREAMING_SNAKE_CASE				=				"conv"

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





@torch.no_grad()
def 		__lowerCamelCase							(UpperCAmelCase__   :							Any					,       UpperCAmelCase__   :							Optional[Any]					,       UpperCAmelCase__   :							Optional[Any]=None					,       UpperCAmelCase__   :							List[Any]=None					,       UpperCAmelCase__   :							Optional[int]=True     ):
					if config_path is not None:
										SCREAMING_SNAKE_CASE				=				DataaVecAudioConfig.from_pretrained(UpperCAmelCase__     )
					else:
										SCREAMING_SNAKE_CASE				=				DataaVecAudioConfig()

					if not is_finetuned:
										# Modify final_proj layer name
										SCREAMING_SNAKE_CASE				=				DataaVecAudioModel(UpperCAmelCase__     )
										SCREAMING_SNAKE_CASE				=				os.path.dirname(UpperCAmelCase__     )

										SCREAMING_SNAKE_CASE				=				torch.load(UpperCAmelCase__     )
										SCREAMING_SNAKE_CASE				=				state_dict["model"].pop("final_proj.0.weight"     )
										SCREAMING_SNAKE_CASE				=				state_dict["model"].pop("final_proj.0.bias"     )
										SCREAMING_SNAKE_CASE				=				os.path.join(UpperCAmelCase__					,       "converted.pt"     )
										torch.save(UpperCAmelCase__					,       UpperCAmelCase__     )
					else:
										SCREAMING_SNAKE_CASE				=				DataaVecAudioForCTC(UpperCAmelCase__     )
										SCREAMING_SNAKE_CASE				=				checkpoint_path

					def load_dataavec(UpperCAmelCase__   :							Optional[Any]     ):
										SCREAMING_SNAKE_CASE	,      SCREAMING_SNAKE_CASE	,      SCREAMING_SNAKE_CASE				=				fairseq.checkpoint_utils.load_model_ensemble_and_task([path]     )
										return model[0].eval()

					SCREAMING_SNAKE_CASE				=				load_dataavec(UpperCAmelCase__     )

					recursively_load_weights(UpperCAmelCase__					,       UpperCAmelCase__					,       not is_finetuned     )

					SCREAMING_SNAKE_CASE				=				WavaVecaProcessor.from_pretrained("facebook/wav2vec2-large-lv60"     )

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

					SCREAMING_SNAKE_CASE				=				processor(UpperCAmelCase__					,       return_tensors="pt"					,       padding=UpperCAmelCase__     )

					SCREAMING_SNAKE_CASE				=				inputs.input_values
					SCREAMING_SNAKE_CASE				=				inputs.attention_mask
					#    input_values = inputs.input_values[:, :-1]
					#    attention_mask = inputs.attention_mask[:, :-1]

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

										SCREAMING_SNAKE_CASE				=				torch.argmax(UpperCAmelCase__					,       dim=-1     )
										SCREAMING_SNAKE_CASE				=				processor.batch_decode(UpperCAmelCase__     )

										print(F"Expected Output: {ds[:4]['text']}, Pred: {output_string}"     )
					else:
										SCREAMING_SNAKE_CASE				=				model(source=UpperCAmelCase__					,       padding_mask=(1 - attention_mask)					,       mask=UpperCAmelCase__					,       features_only=UpperCAmelCase__     )[
										    "layer_results"
										][-1][0].transpose(0					,       1     )
										SCREAMING_SNAKE_CASE				=				hf_wavavec(UpperCAmelCase__					,       attention_mask=UpperCAmelCase__     )["last_hidden_state"]

					print(our_output.shape					,       their_output.shape     )
					SCREAMING_SNAKE_CASE				=				torch.max(torch.abs(our_output - their_output     )     ).item()
					print(F"max_absolute_diff = {max_absolute_diff}"     )  # ~ 1e-7
					SCREAMING_SNAKE_CASE				=				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__":
							_lowerCamelCase    :      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'''
							)
							_lowerCamelCase    :      Union[str, 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
							)




