


"""simple docstring"""






import unittest

from transformers import MarkupLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    MarkupLMForQuestionAnswering,
					    MarkupLMForSequenceClassification,
					    MarkupLMForTokenClassification,
					    MarkupLMModel,
					)

# TODO check dependencies
from transformers import MarkupLMFeatureExtractor, MarkupLMProcessor, MarkupLMTokenizer


class 			SCREAMING_SNAKE_CASE_  :




							'''simple docstring'''





							def __init__(						self	,					lowerCamelCase__	,					lowerCamelCase__=13	,					lowerCamelCase__=7	,					lowerCamelCase__=True	,					lowerCamelCase__=True	,					lowerCamelCase__=True	,					lowerCamelCase__=True	,					lowerCamelCase__=99	,					lowerCamelCase__=32	,					lowerCamelCase__=5	,					lowerCamelCase__=4	,					lowerCamelCase__=37	,					lowerCamelCase__="gelu"	,					lowerCamelCase__=0.1	,					lowerCamelCase__=0.1	,					lowerCamelCase__=512	,					lowerCamelCase__=16	,					lowerCamelCase__=2	,					lowerCamelCase__=0.02	,					lowerCamelCase__=3	,					lowerCamelCase__=None	,					lowerCamelCase__=20	,					lowerCamelCase__=30	,					lowerCamelCase__=2	,					lowerCamelCase__=2	,					lowerCamelCase__=10	,					)   ->	Optional[Any]:







								'''simple docstring'''




								snake_case__  :						Union[str, Any]		=      parent
								snake_case__  :						Union[str, Any]		=      batch_size
								snake_case__  :						Optional[Any]		=      seq_length
								snake_case__  :						str		=      is_training
								snake_case__  :						Union[str, Any]		=      use_input_mask
								snake_case__  :						int		=      use_token_type_ids
								snake_case__  :						Tuple		=      use_labels
								snake_case__  :						Optional[Any]		=      vocab_size
								snake_case__  :						int		=      hidden_size
								snake_case__  :						Union[str, Any]		=      num_hidden_layers
								snake_case__  :						List[str]		=      num_attention_heads
								snake_case__  :						Optional[int]		=      intermediate_size
								snake_case__  :						Optional[int]		=      hidden_act
								snake_case__  :						List[Any]		=      hidden_dropout_prob
								snake_case__  :						List[Any]		=      attention_probs_dropout_prob
								snake_case__  :						int		=      max_position_embeddings
								snake_case__  :						List[Any]		=      type_vocab_size
								snake_case__  :						Optional[Any]		=      type_sequence_label_size
								snake_case__  :						List[str]		=      initializer_range
								snake_case__  :						List[Any]		=      num_labels
								snake_case__  :						Any		=      scope
								snake_case__  :						Tuple		=      max_xpath_tag_unit_embeddings
								snake_case__  :						List[str]		=      max_xpath_subs_unit_embeddings
								snake_case__  :						Any		=      tag_pad_id
								snake_case__  :						int		=      subs_pad_id
								snake_case__  :						List[str]		=      max_depth




							def        UpperCAmelCase (						self)   ->	str:







								'''simple docstring'''




								snake_case__  :						Dict		=      ids_tensor([self.batch_size, self.seq_length]	,					self.vocab_size)

								snake_case__  :						List[Any]		=      ids_tensor(
								    [self.batch_size, self.seq_length, self.max_depth]	,					self.max_xpath_tag_unit_embeddings)

								snake_case__  :						int		=      ids_tensor(
								    [self.batch_size, self.seq_length, self.max_depth]	,					self.max_xpath_subs_unit_embeddings)

								snake_case__  :						Dict		=      None
								if self.use_input_mask:
									snake_case__  :						Optional[int]		=      ids_tensor([self.batch_size, self.seq_length]	,					vocab_size=2)

								snake_case__  :						Union[str, Any]		=      None
								if self.use_token_type_ids:
									snake_case__  :						List[str]		=      ids_tensor([self.batch_size, self.seq_length]	,					self.type_vocab_size)

								snake_case__  :						Any		=      None
								snake_case__  :						str		=      None
								if self.use_labels:
									snake_case__  :						Tuple		=      ids_tensor([self.batch_size]	,					self.type_sequence_label_size)
									snake_case__  :						Tuple		=      ids_tensor([self.batch_size, self.seq_length]	,					self.num_labels)

								snake_case__  :						List[str]		=      self.get_config()

								return (
								    config,
								    input_ids,
								    xpath_tags_seq,
								    xpath_subs_seq,
								    token_type_ids,
								    input_mask,
								    sequence_labels,
								    token_labels,
								)




							def        UpperCAmelCase (						self)   ->	int:







								'''simple docstring'''




								return MarkupLMConfig(
								    vocab_size=self.vocab_size	,					hidden_size=self.hidden_size	,					num_hidden_layers=self.num_hidden_layers	,					num_attention_heads=self.num_attention_heads	,					intermediate_size=self.intermediate_size	,					hidden_act=self.hidden_act	,					hidden_dropout_prob=self.hidden_dropout_prob	,					attention_probs_dropout_prob=self.attention_probs_dropout_prob	,					max_position_embeddings=self.max_position_embeddings	,					type_vocab_size=self.type_vocab_size	,					initializer_range=self.initializer_range	,					max_xpath_tag_unit_embeddings=self.max_xpath_tag_unit_embeddings	,					max_xpath_subs_unit_embeddings=self.max_xpath_subs_unit_embeddings	,					tag_pad_id=self.tag_pad_id	,					subs_pad_id=self.subs_pad_id	,					max_depth=self.max_depth	,					)




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					)   ->	Union[str, Any]:







								'''simple docstring'''




								snake_case__  :						Optional[Any]		=      MarkupLMModel(config=lowerCamelCase__)
								model.to(lowerCamelCase__)
								model.eval()
								print("Configs:"	,					model.config.tag_pad_id	,					model.config.subs_pad_id)
								snake_case__  :						Any		=      model(lowerCamelCase__	,					attention_mask=lowerCamelCase__	,					token_type_ids=lowerCamelCase__)
								snake_case__  :						Optional[Any]		=      model(lowerCamelCase__	,					token_type_ids=lowerCamelCase__)
								snake_case__  :						Tuple		=      model(lowerCamelCase__)
								self.parent.assertEqual(result.last_hidden_state.shape	,					(self.batch_size, self.seq_length, self.hidden_size))
								self.parent.assertEqual(result.pooler_output.shape	,					(self.batch_size, self.hidden_size))




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					)   ->	int:







								'''simple docstring'''




								snake_case__  :						List[str]		=      self.num_labels
								snake_case__  :						int		=      MarkupLMForSequenceClassification(lowerCamelCase__)
								model.to(lowerCamelCase__)
								model.eval()
								snake_case__  :						Tuple		=      model(
								    lowerCamelCase__	,					xpath_tags_seq=lowerCamelCase__	,					xpath_subs_seq=lowerCamelCase__	,					attention_mask=lowerCamelCase__	,					token_type_ids=lowerCamelCase__	,					labels=lowerCamelCase__	,					)
								self.parent.assertEqual(result.logits.shape	,					(self.batch_size, self.num_labels))




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					)   ->	List[Any]:







								'''simple docstring'''




								snake_case__  :						Dict		=      self.num_labels
								snake_case__  :						List[str]		=      MarkupLMForTokenClassification(config=lowerCamelCase__)
								model.to(lowerCamelCase__)
								model.eval()
								snake_case__  :						Tuple		=      model(
								    lowerCamelCase__	,					xpath_tags_seq=lowerCamelCase__	,					xpath_subs_seq=lowerCamelCase__	,					attention_mask=lowerCamelCase__	,					token_type_ids=lowerCamelCase__	,					labels=lowerCamelCase__	,					)
								self.parent.assertEqual(result.logits.shape	,					(self.batch_size, self.seq_length, self.num_labels))




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					)   ->	List[Any]:







								'''simple docstring'''




								snake_case__  :						List[Any]		=      MarkupLMForQuestionAnswering(config=lowerCamelCase__)
								model.to(lowerCamelCase__)
								model.eval()
								snake_case__  :						List[str]		=      model(
								    lowerCamelCase__	,					xpath_tags_seq=lowerCamelCase__	,					xpath_subs_seq=lowerCamelCase__	,					attention_mask=lowerCamelCase__	,					token_type_ids=lowerCamelCase__	,					start_positions=lowerCamelCase__	,					end_positions=lowerCamelCase__	,					)
								self.parent.assertEqual(result.start_logits.shape	,					(self.batch_size, self.seq_length))
								self.parent.assertEqual(result.end_logits.shape	,					(self.batch_size, self.seq_length))







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







								'''simple docstring'''




								snake_case__  :						Any		=      self.prepare_config_and_inputs()
								(
								    (
								    snake_case__
								),					(
								    snake_case__
								),					(
								    snake_case__
								),					(
								    snake_case__
								),					(
								    snake_case__
								),					(
								    snake_case__
								),					(
								    snake_case__
								),					(
								    snake_case__
								),					
								)  :						str		=      config_and_inputs
								snake_case__  :						List[str]		=      {
								    "input_ids": input_ids,
								    "xpath_tags_seq": xpath_tags_seq,
								    "xpath_subs_seq": xpath_subs_seq,
								    "token_type_ids": token_type_ids,
								    "attention_mask": input_mask,
								}
								return config, inputs_dict


@require_torch
class 			SCREAMING_SNAKE_CASE_  (   _lowercase	,      _lowercase	,      unittest.TestCase):




							'''simple docstring'''





							__magic_name__					:     List[Any]				  =  (
							    (
							        MarkupLMModel,
							        MarkupLMForSequenceClassification,
							        MarkupLMForTokenClassification,
							        MarkupLMForQuestionAnswering,
							    )
							    if is_torch_available()
							    else None
							)
							__magic_name__					:     Any				  =  (
							    {
							        '''feature-extraction''': MarkupLMModel,
							        '''question-answering''': MarkupLMForQuestionAnswering,
							        '''text-classification''': MarkupLMForSequenceClassification,
							        '''token-classification''': MarkupLMForTokenClassification,
							        '''zero-shot''': MarkupLMForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__)   ->	Any:







								'''simple docstring'''




								return True




							def        UpperCAmelCase (						self)   ->	Optional[int]:







								'''simple docstring'''




								snake_case__  :						str		=      MarkupLMModelTester(self)
								snake_case__  :						List[str]		=      ConfigTester(self	,					config_class=lowerCamelCase__	,					hidden_size=37)




							def        UpperCAmelCase (						self)   ->	Union[str, Any]:







								'''simple docstring'''




								self.config_tester.run_common_tests()




							def        UpperCAmelCase (						self)   ->	Tuple:







								'''simple docstring'''




								snake_case__  :						Optional[int]		=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCamelCase__)




							def        UpperCAmelCase (						self)   ->	List[Any]:







								'''simple docstring'''




								snake_case__  :						Optional[int]		=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase__)




							def        UpperCAmelCase (						self)   ->	List[Any]:







								'''simple docstring'''




								snake_case__  :						Any		=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*lowerCamelCase__)







							def        UpperCAmelCase (						self)   ->	int:







								'''simple docstring'''




								snake_case__  :						List[str]		=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*lowerCamelCase__)



def 			A__		(					)							->	Tuple:

	'''simple docstring'''


	snake_case__  :						Tuple		=      "\n    <!DOCTYPE html>\n    <html>\n    <head>\n    <title>Page Title</title>\n    </head>\n    <body>\n\n    <h1>This is a Heading</h1>\n    <p>This is a paragraph.</p>\n\n    </body>\n    </html>\n    "

	return html_string



@require_torch
class 			SCREAMING_SNAKE_CASE_  (   unittest.TestCase):




							'''simple docstring'''





							@cached_property
							def        UpperCAmelCase (						self)   ->	Dict:







								'''simple docstring'''




								snake_case__  :						Optional[int]		=      MarkupLMFeatureExtractor()
								snake_case__  :						List[str]		=      MarkupLMTokenizer.from_pretrained("microsoft/markuplm-base")

								return MarkupLMProcessor(lowerCamelCase__	,					lowerCamelCase__)







							@slow
							def        UpperCAmelCase (						self)   ->	Any:







								'''simple docstring'''




								snake_case__  :						str		=      MarkupLMModel.from_pretrained("microsoft/markuplm-base").to(lowerCamelCase__)

								snake_case__  :						List[Any]		=      self.default_processor

								snake_case__  :						Optional[Any]		=      processor(prepare_html_string()	,					return_tensors="pt")
								snake_case__  :						Any		=      inputs.to(lowerCamelCase__)

								# forward pass
								with torch.no_grad():
									snake_case__  :						Dict		=      model(**lowerCamelCase__)

								# verify the last hidden states
								snake_case__  :						Optional[int]		=      torch.Size([1, 14, 768])
								self.assertEqual(outputs.last_hidden_state.shape	,					lowerCamelCase__)

								snake_case__  :						Tuple		=      torch.tensor(
								    [[0.06_75, -0.00_52, 0.50_01], [-0.22_81, 0.08_02, 0.21_92], [-0.05_83, -0.33_11, 0.11_85]]).to(lowerCamelCase__)

								self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]	,					lowerCamelCase__	,					atol=1E-4))









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

lowercase  		=			{
    """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""",
}
lowercase  		=			[
    """lm_head""",
]



def 			A__		(					_UpperCAmelCase    :      str				,    _UpperCAmelCase    :      List[Any]				,    _UpperCAmelCase    :      List[str]				,    _UpperCAmelCase    :      Optional[int]				,    _UpperCAmelCase    :      Tuple	)							->	Union[str, Any]:

	'''simple docstring'''


	for attribute in key.split("."	):
		snake_case__  :						Dict		=      getattr(_UpperCAmelCase				,    _UpperCAmelCase	)

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

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



def 			A__		(					_UpperCAmelCase    :      List[Any]				,    _UpperCAmelCase    :      Dict				,    _UpperCAmelCase    :      Tuple	)							->	List[Any]:

	'''simple docstring'''


	snake_case__  :						Optional[Any]		=      []
	snake_case__  :						Any		=      fairseq_model.state_dict()

	if not is_headless:
		snake_case__  :						int		=      hf_model.dataavec_audio.feature_extractor
		snake_case__  :						Optional[Any]		=      hf_model.dataavec_audio.encoder.pos_conv_embed

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

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

	'''simple docstring'''


	snake_case__  :						Optional[Any]		=      path.split("."	)
	return reduce(_UpperCAmelCase				,    _UpperCAmelCase				,    _UpperCAmelCase	)



def 			A__		(					_UpperCAmelCase    :      List[str]				,    _UpperCAmelCase    :      str				,    _UpperCAmelCase    :      Dict				,    _UpperCAmelCase    :      str	)							->	Tuple:

	'''simple docstring'''


	snake_case__  :						List[str]		=      access_by_string(_UpperCAmelCase				,    _UpperCAmelCase	)
	snake_case__  :						Optional[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."""	)
	snake_case__  :						Tuple		=      fsq_value
	logger.info(F"""{full_name} was correctly initialized from {hf_weight_path}."""	)



def 			A__		(					_UpperCAmelCase    :      Any				,    _UpperCAmelCase    :      Any				,    _UpperCAmelCase    :      List[str]				,    _UpperCAmelCase    :      Tuple	)							->	Dict:

	'''simple docstring'''


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

	snake_case__  :						Tuple		=      name.split("."	)[-1]
	if type_id == 0:
		snake_case__  :						Dict		=      "conv"
	elif type_id == 2:
		snake_case__  :						Dict		=      "layer_norm"
	else:
		unused_weights.append(_UpperCAmelCase	)
		return

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



def 			A__		(					_UpperCAmelCase    :      Any				,    _UpperCAmelCase    :      Dict				,    _UpperCAmelCase    :      Dict				,    _UpperCAmelCase    :      str	)							->	Tuple:

	'''simple docstring'''


	snake_case__  :						Optional[Any]		=      full_name.split("pos_conv."	)[-1]
	snake_case__  :						Optional[Any]		=      name.split("."	)
	snake_case__  :						Dict		=      int(items[0]	)
	snake_case__  :						Optional[Any]		=      int(items[1]	)

	snake_case__  :						int		=      name.split("."	)[-1]
	if type_id != 0:
		unused_weights.append(_UpperCAmelCase	)
		return
	else:
		snake_case__  :						Optional[int]		=      "conv"

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



@torch.no_grad()
def 			A__		(					_UpperCAmelCase    :      Dict				,    _UpperCAmelCase    :      Union[str, Any]				,    _UpperCAmelCase    :      List[str]=None				,    _UpperCAmelCase    :      Optional[Any]=None				,    _UpperCAmelCase    :      int=True	)							->	Dict:

	'''simple docstring'''


	if config_path is not None:
		snake_case__  :						List[str]		=      DataaVecAudioConfig.from_pretrained(_UpperCAmelCase	)
	else:
		snake_case__  :						int		=      DataaVecAudioConfig()

	if not is_finetuned:
		# Modify final_proj layer name
		snake_case__  :						int		=      DataaVecAudioModel(_UpperCAmelCase	)
		snake_case__  :						Union[str, Any]		=      os.path.dirname(_UpperCAmelCase	)

		snake_case__  :						Optional[Any]		=      torch.load(_UpperCAmelCase	)
		snake_case__  :						str		=      state_dict["model"].pop("final_proj.0.weight"	)
		snake_case__  :						List[str]		=      state_dict["model"].pop("final_proj.0.bias"	)
		snake_case__  :						Optional[Any]		=      os.path.join(_UpperCAmelCase				,    "converted.pt"	)
		torch.save(_UpperCAmelCase				,    _UpperCAmelCase	)
	else:
		snake_case__  :						Optional[Any]		=      DataaVecAudioForCTC(_UpperCAmelCase	)
		snake_case__  :						int		=      checkpoint_path

	def load_dataavec(_UpperCAmelCase    :      Union[str, Any]	):
		snake_case__,					snake_case__,					snake_case__  :						Dict		=      fairseq.checkpoint_utils.load_model_ensemble_and_task([path]	)
		return model[0].eval()

	snake_case__  :						Optional[int]		=      load_dataavec(_UpperCAmelCase	)

	recursively_load_weights(_UpperCAmelCase				,    _UpperCAmelCase				,    not is_finetuned	)

	snake_case__  :						Any		=      WavaVecaProcessor.from_pretrained("facebook/wav2vec2-large-lv60"	)

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

	snake_case__  :						Union[str, Any]		=      processor(_UpperCAmelCase				,    return_tensors="pt"				,    padding=_UpperCAmelCase	)

	snake_case__  :						Any		=      inputs.input_values
	snake_case__  :						List[str]		=      inputs.attention_mask
	#    input_values = inputs.input_values[:, :-1]
	#    attention_mask = inputs.attention_mask[:, :-1]

	hf_wavavec.eval()
	model.eval()
	if is_finetuned:
		snake_case__  :						Optional[Any]		=      model(source=_UpperCAmelCase				,    padding_mask=(1 - attention_mask)				,    mask=_UpperCAmelCase				,    features_only=_UpperCAmelCase	)[
		    "encoder_out"
		].transpose(0				,    1	)
		snake_case__  :						Union[str, Any]		=      hf_wavavec(_UpperCAmelCase				,    attention_mask=_UpperCAmelCase	)["logits"]

		snake_case__  :						Optional[int]		=      torch.argmax(_UpperCAmelCase				,    dim=-1	)
		snake_case__  :						Union[str, Any]		=      processor.batch_decode(_UpperCAmelCase	)

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

	print(our_output.shape				,    their_output.shape	)
	snake_case__  :						Union[str, Any]		=      torch.max(torch.abs(our_output - their_output	)	).item()
	print(F"""max_absolute_diff = {max_absolute_diff}"""	)  # ~ 1e-7
	snake_case__  :						str		=      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__":
					lowercase  		=			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"""
					)
					lowercase  		=			parser.parse_args()
					convert_wavaveca_checkpoint(
					    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
					)





