

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()
_UpperCamelCase			       =			logging.get_logger(__name__)

_UpperCamelCase			       =			{
    "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",
}
_UpperCamelCase			       =			[
    "lm_head",
]
def 						UpperCamelCase_(			snake_case__:  Tuple  ,    snake_case__:  Dict  ,    snake_case__:  Optional[int]  ,    snake_case__:  str  ,    snake_case__:  Optional[int]   )       -> Dict:
 for attribute in key.split('.'   ):
  UpperCAmelCase__							=				getattr(UpperCamelCase_  ,    UpperCamelCase_   )

 if weight_type is not None:
  UpperCAmelCase__							=				getattr(UpperCamelCase_  ,    UpperCamelCase_   ).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 						UpperCamelCase_(			snake_case__:  Tuple  ,    snake_case__:  Any  ,    snake_case__:  List[str]   )       -> Union[str, Any]:
 UpperCAmelCase__							=				[]
 UpperCAmelCase__							=				fairseq_model.state_dict()

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

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

 for name, value in fairseq_dict.items():
  UpperCAmelCase__							=				False
  if "conv_layers" in name:
   load_conv_layer(
       UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_  ,    )
   UpperCAmelCase__							=				True
  elif "pos_conv" in name:
   load_pos_conv_layer(
       UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_  ,    )
   UpperCAmelCase__							=				True
  else:
   for key, mapped_key in MAPPING.items():
    if not is_headless:
     UpperCAmelCase__							=				"""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]:
     UpperCAmelCase__							=				True
     if "*" in mapped_key:
      UpperCAmelCase__							=				name.split(UpperCamelCase_   )[0].split('.'   )[-2]
      UpperCAmelCase__							=				mapped_key.replace('*'  ,    UpperCamelCase_   )
     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(UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_   )
    continue
  if not is_used:
   unused_weights.append(UpperCamelCase_   )

 logger.warning(f"Unused weights: {unused_weights}"   )
def 						UpperCamelCase_(			snake_case__:  Any  ,    snake_case__:  Dict   )       -> Union[str, Any]:
 UpperCAmelCase__							=				path.split('.'   )
 return reduce(UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_   )
def 						UpperCamelCase_(			snake_case__:  Union[str, Any]  ,    snake_case__:  Optional[Any]  ,    snake_case__:  List[Any]  ,    snake_case__:  Dict   )       -> Dict:
 UpperCAmelCase__							=				access_by_string(UpperCamelCase_  ,    UpperCamelCase_   )
 UpperCAmelCase__							=				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."   )
 UpperCAmelCase__							=				fsq_value
 logger.info(f"{full_name} was correctly initialized from {hf_weight_path}."   )
def 						UpperCamelCase_(			snake_case__:  Any  ,    snake_case__:  Dict  ,    snake_case__:  List[Any]  ,    snake_case__:  Optional[Any]   )       -> Optional[Any]:
 UpperCAmelCase__							=				full_name.split('conv_layers.'   )[-1]
 UpperCAmelCase__							=				name.split('.'   )
 UpperCAmelCase__							=				int(items[0]   )
 UpperCAmelCase__							=				int(items[1]   )

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

 set_weights(UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_  ,    f"conv_layers.{layer_id}.{layer_type}.{weight_type}"   )
def 						UpperCamelCase_(			snake_case__:  Dict  ,    snake_case__:  int  ,    snake_case__:  List[Any]  ,    snake_case__:  str   )       -> Tuple:
 UpperCAmelCase__							=				full_name.split('pos_conv.'   )[-1]
 UpperCAmelCase__							=				name.split('.'   )
 UpperCAmelCase__							=				int(items[0]   )
 UpperCAmelCase__							=				int(items[1]   )

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

 set_weights(UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_  ,    f"layers.{layer_id}.{layer_type}.{weight_type}"   )
@torch.no_grad()
def 						UpperCamelCase_(			snake_case__:  Union[str, Any]  ,    snake_case__:  Optional[int]  ,    snake_case__:  Union[str, Any]=None  ,    snake_case__:  Union[str, Any]=None  ,    snake_case__:  Union[str, Any]=True   )       -> Any:
 if config_path is not None:
  UpperCAmelCase__							=				DataaVecAudioConfig.from_pretrained(UpperCamelCase_   )
 else:
  UpperCAmelCase__							=				DataaVecAudioConfig()

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

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

 def load_dataavec(snake_case__:  Dict   ):
  UpperCAmelCase__							=				fairseq.checkpoint_utils.load_model_ensemble_and_task([path]   )
  return model[0].eval()

 UpperCAmelCase__							=				load_dataavec(UpperCamelCase_   )

 recursively_load_weights(UpperCamelCase_  ,    UpperCamelCase_  ,    not is_finetuned   )

 UpperCAmelCase__							=				WavaVecaProcessor.from_pretrained('facebook/wav2vec2-large-lv60'   )

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

 UpperCAmelCase__							=				processor(UpperCamelCase_  ,    return_tensors='pt'  ,    padding=UpperCamelCase_   )

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

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

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

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

 print(our_output.shape  ,    their_output.shape   )
 UpperCAmelCase__							=				torch.max(torch.abs(our_output - their_output   )   ).item()
 print(f"max_absolute_diff = {max_absolute_diff}"   )  # ~ 1e-7
 UpperCAmelCase__							=				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__":
 _UpperCamelCase			       =			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			       =			parser.parse_args()
 convert_wavaveca_checkpoint(
     args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
 )









import json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_longformer import LongformerTokenizer


_lowerCamelCase					:     Optional[Any]					  =  logging.get_logger(__name__)

_lowerCamelCase					:     Dict					  =  {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}

_lowerCamelCase					:     List[str]					  =  {
    "vocab_file": {
        "allenai/longformer-base-4096": "https://huggingface.co/allenai/longformer-base-4096/resolve/main/vocab.json",
        "allenai/longformer-large-4096": (
            "https://huggingface.co/allenai/longformer-large-4096/resolve/main/vocab.json"
        ),
        "allenai/longformer-large-4096-finetuned-triviaqa": (
            "https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/vocab.json"
        ),
        "allenai/longformer-base-4096-extra.pos.embd.only": (
            "https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/vocab.json"
        ),
        "allenai/longformer-large-4096-extra.pos.embd.only": (
            "https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/vocab.json"
        ),
    },
    "merges_file": {
        "allenai/longformer-base-4096": "https://huggingface.co/allenai/longformer-base-4096/resolve/main/merges.txt",
        "allenai/longformer-large-4096": (
            "https://huggingface.co/allenai/longformer-large-4096/resolve/main/merges.txt"
        ),
        "allenai/longformer-large-4096-finetuned-triviaqa": (
            "https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/merges.txt"
        ),
        "allenai/longformer-base-4096-extra.pos.embd.only": (
            "https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/merges.txt"
        ),
        "allenai/longformer-large-4096-extra.pos.embd.only": (
            "https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/merges.txt"
        ),
    },
    "tokenizer_file": {
        "allenai/longformer-base-4096": (
            "https://huggingface.co/allenai/longformer-base-4096/resolve/main/tokenizer.json"
        ),
        "allenai/longformer-large-4096": (
            "https://huggingface.co/allenai/longformer-large-4096/resolve/main/tokenizer.json"
        ),
        "allenai/longformer-large-4096-finetuned-triviaqa": (
            "https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/tokenizer.json"
        ),
        "allenai/longformer-base-4096-extra.pos.embd.only": (
            "https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/tokenizer.json"
        ),
        "allenai/longformer-large-4096-extra.pos.embd.only": (
            "https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/tokenizer.json"
        ),
    },
}

_lowerCamelCase					:     Dict					  =  {
    "allenai/longformer-base-4096": 4_0_9_6,
    "allenai/longformer-large-4096": 4_0_9_6,
    "allenai/longformer-large-4096-finetuned-triviaqa": 4_0_9_6,
    "allenai/longformer-base-4096-extra.pos.embd.only": 4_0_9_6,
    "allenai/longformer-large-4096-extra.pos.embd.only": 4_0_9_6,
}
class 				__snake_case     (_a						):
				lowerCAmelCase__			=   VOCAB_FILES_NAMES
				lowerCAmelCase__			=   PRETRAINED_VOCAB_FILES_MAP
				lowerCAmelCase__			=   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				lowerCAmelCase__			=   ["input_ids", "attention_mask"]
				lowerCAmelCase__			=   LongformerTokenizer



				def __init__(					self					:						Optional[Any]						,		_UpperCAmelCase					:						Tuple=None						,		_UpperCAmelCase					:						List[str]=None						,		_UpperCAmelCase					:						List[str]=None						,		_UpperCAmelCase					:						Union[str, Any]="replace"						,		_UpperCAmelCase					:						str="<s>"						,		_UpperCAmelCase					:						str="</s>"						,		_UpperCAmelCase					:						int="</s>"						,		_UpperCAmelCase					:						str="<s>"						,		_UpperCAmelCase					:						Optional[Any]="<unk>"						,		_UpperCAmelCase					:						Optional[Any]="<pad>"						,		_UpperCAmelCase					:						Any="<mask>"						,		_UpperCAmelCase					:						Optional[int]=False						,		_UpperCAmelCase					:						List[str]=True						,		**_UpperCAmelCase					:						Optional[int]						,		)     ->			str:





								'''simple docstring'''




								super().__init__(
								    _UpperCAmelCase						,		_UpperCAmelCase						,		tokenizer_file=_UpperCAmelCase						,		errors=_UpperCAmelCase						,		bos_token=_UpperCAmelCase						,		eos_token=_UpperCAmelCase						,		sep_token=_UpperCAmelCase						,		cls_token=_UpperCAmelCase						,		unk_token=_UpperCAmelCase						,		pad_token=_UpperCAmelCase						,		mask_token=_UpperCAmelCase						,		add_prefix_space=_UpperCAmelCase						,		trim_offsets=_UpperCAmelCase						,		**_UpperCAmelCase						,		)

								_lowerCAmelCase			:			Tuple        =					json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()					)
								if pre_tok_state.get("""add_prefix_space"""						,		_UpperCAmelCase					) != add_prefix_space:
												_lowerCAmelCase			:			str        =					getattr(_UpperCAmelCase						,		pre_tok_state.pop("""type"""					)					)
												_lowerCAmelCase			:			Dict        =					add_prefix_space
												_lowerCAmelCase			:			Union[str, Any]        =					pre_tok_class(**_UpperCAmelCase					)

								_lowerCAmelCase			:			List[Any]        =					add_prefix_space

								_lowerCAmelCase			:			List[str]        =					"""post_processor"""
								_lowerCAmelCase			:			str        =					getattr(self.backend_tokenizer						,		_UpperCAmelCase						,		_UpperCAmelCase					)
								if tokenizer_component_instance:
												_lowerCAmelCase			:			Any        =					json.loads(tokenizer_component_instance.__getstate__()					)

												# The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
												if "sep" in state:
																_lowerCAmelCase			:			str        =					tuple(state["""sep"""]					)
												if "cls" in state:
																_lowerCAmelCase			:			int        =					tuple(state["""cls"""]					)

												_lowerCAmelCase			:			Any        =					False

												if state.get("""add_prefix_space"""						,		_UpperCAmelCase					) != add_prefix_space:
																_lowerCAmelCase			:			int        =					add_prefix_space
																_lowerCAmelCase			:			List[str]        =					True

												if state.get("""trim_offsets"""						,		_UpperCAmelCase					) != trim_offsets:
																_lowerCAmelCase			:			List[str]        =					trim_offsets
																_lowerCAmelCase			:			List[str]        =					True

												if changes_to_apply:
																_lowerCAmelCase			:			List[str]        =					getattr(_UpperCAmelCase						,		state.pop("""type"""					)					)
																_lowerCAmelCase			:			Any        =					component_class(**_UpperCAmelCase					)
																setattr(self.backend_tokenizer						,		_UpperCAmelCase						,		_UpperCAmelCase					)



				@property
				def 					SCREAMING_SNAKE_CASE				(					self					:						List[Any]					)     ->			str:





								'''simple docstring'''




								if self._mask_token is None:
												if self.verbose:
																logger.error("""Using mask_token, but it is not set yet."""					)
												return None
								return str(self._mask_token					)



				@mask_token.setter
				def 					SCREAMING_SNAKE_CASE				(					self					:						Any						,		_UpperCAmelCase					:						Any					)     ->			Tuple:





								'''simple docstring'''




								_lowerCAmelCase			:			Tuple        =					AddedToken(_UpperCAmelCase						,		lstrip=_UpperCAmelCase						,		rstrip=_UpperCAmelCase					) if isinstance(_UpperCAmelCase						,		_UpperCAmelCase					) else value
								_lowerCAmelCase			:			List[Any]        =					value



				def 					SCREAMING_SNAKE_CASE				(					self					:						int						,		*_UpperCAmelCase					:						str						,		**_UpperCAmelCase					:						Union[str, Any]					)     ->			BatchEncoding:





								'''simple docstring'''




								_lowerCAmelCase			:			Dict        =					kwargs.get("""is_split_into_words"""						,		_UpperCAmelCase					)
								assert self.add_prefix_space or not is_split_into_words, (
								    f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
								    "to use it with pretokenized inputs."
								)

								return super()._batch_encode_plus(*_UpperCAmelCase						,		**_UpperCAmelCase					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						str						,		*_UpperCAmelCase					:						Any						,		**_UpperCAmelCase					:						Optional[Any]					)     ->			BatchEncoding:





								'''simple docstring'''




								_lowerCAmelCase			:			Any        =					kwargs.get("""is_split_into_words"""						,		_UpperCAmelCase					)

								assert self.add_prefix_space or not is_split_into_words, (
								    f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
								    "to use it with pretokenized inputs."
								)

								return super()._encode_plus(*_UpperCAmelCase						,		**_UpperCAmelCase					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						Dict						,		_UpperCAmelCase					:						str						,		_UpperCAmelCase					:						Optional[str] = None					)     ->			Tuple[str]:





								'''simple docstring'''




								_lowerCAmelCase			:			int        =					self._tokenizer.model.save(_UpperCAmelCase						,		name=_UpperCAmelCase					)
								return tuple(_UpperCAmelCase					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						Any						,		_UpperCAmelCase					:						Any						,		_UpperCAmelCase					:						str=None					)     ->			List[Any]:





								'''simple docstring'''




								_lowerCAmelCase			:			Any        =					[self.bos_token_id] + token_ids_a + [self.eos_token_id]
								if token_ids_a is None:
												return output

								return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]



				def 					SCREAMING_SNAKE_CASE				(					self					:						List[Any]						,		_UpperCAmelCase					:						List[int]						,		_UpperCAmelCase					:						Optional[List[int]] = None					)     ->			List[int]:





								'''simple docstring'''




								_lowerCAmelCase			:			Dict        =					[self.sep_token_id]
								_lowerCAmelCase			:			Optional[Any]        =					[self.cls_token_id]

								if token_ids_a is None:
												return len(cls + token_ids_a + sep					) * [0]
								return len(cls + token_ids_a + sep + sep + token_ids_a + sep					) * [0]

