



"""simple docstring"""
import json
import os
import shutil
import tempfile
import unittest
from multiprocessing import get_context
from pathlib import Path

import datasets
import numpy as np
from datasets import load_dataset
from parameterized import parameterized

from transformers import AutoProcessor
from transformers.models.wavaveca import WavaVecaCTCTokenizer, WavaVecaFeatureExtractor
from transformers.models.wavaveca.tokenization_wavaveca import VOCAB_FILES_NAMES
from transformers.testing_utils import require_pyctcdecode, require_torch, require_torchaudio, slow
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_pyctcdecode_available, is_torch_available

from ..wavaveca.test_feature_extraction_wavaveca import floats_list


if is_pyctcdecode_available():
      from huggingface_hub import snapshot_download
      from pyctcdecode import BeamSearchDecoderCTC

      from transformers.models.wavaveca_with_lm import WavaVecaProcessorWithLM
      from transformers.models.wavaveca_with_lm.processing_wavaveca_with_lm import WavaVecaDecoderWithLMOutput

if is_torch_available():
      from transformers import WavaVecaForCTC



@require_pyctcdecode
class       SCREAMING_SNAKE_CASE__  (					unittest.TestCase ):

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    UpperCAmelCase	:					Any	=	"""| <pad> <unk> <s> </s> a b c d e f g h i j k""".split()
    UpperCAmelCase	:					List[str]	=	dict(zip(_SCREAMING_SNAKE_CASE      ,     range(len(_SCREAMING_SNAKE_CASE						)						)						)						)

    UpperCAmelCase	:					Optional[int]	=	{
        """unk_token""": """<unk>""",
        """bos_token""": """<s>""",
        """eos_token""": """</s>""",
    }
    UpperCAmelCase	:					List[Any]	=	{
        """feature_size""": 1,
        """padding_value""": 0.0,
        """sampling_rate""": 16000,
        """return_attention_mask""": False,
        """do_normalize""": True,
    }

    UpperCAmelCase	:					Optional[Any]	=	tempfile.mkdtemp()
    UpperCAmelCase	:					Optional[int]	=	os.path.join(self.tmpdirname      ,     VOCAB_FILES_NAMES["""vocab_file"""]						)
    UpperCAmelCase	:					int	=	os.path.join(self.tmpdirname      ,     _SCREAMING_SNAKE_CASE						)
    with open(self.vocab_file      ,     """w"""      ,     encoding="""utf-8"""						) as fp:
     fp.write(json.dumps(_SCREAMING_SNAKE_CASE						) + """\n"""						)

    with open(self.feature_extraction_file      ,     """w"""      ,     encoding="""utf-8"""						) as fp:
     fp.write(json.dumps(_SCREAMING_SNAKE_CASE						) + """\n"""						)

    # load decoder from hub
    UpperCAmelCase	:					Optional[Any]	=	"""hf-internal-testing/ngram-beam-search-decoder"""

   def 	SCREAMING_SNAKE_CASE     (   self      ,     **_SCREAMING_SNAKE_CASE						)  ->      Tuple:





    '''simple docstring'''




    UpperCAmelCase	:					List[str]	=	self.add_kwargs_tokens_map.copy()
    kwargs.update(_SCREAMING_SNAKE_CASE						)
    return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname      ,     **_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self      ,     **_SCREAMING_SNAKE_CASE						)  ->      Any:





    '''simple docstring'''




    return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname      ,     **_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self      ,     **_SCREAMING_SNAKE_CASE						)  ->      int:





    '''simple docstring'''




    return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name      ,     **_SCREAMING_SNAKE_CASE						)

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





    '''simple docstring'''




    shutil.rmtree(self.tmpdirname						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      int:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[int]	=	self.get_tokenizer()
    UpperCAmelCase	:					Tuple	=	self.get_feature_extractor()
    UpperCAmelCase	:					Dict	=	self.get_decoder()

    UpperCAmelCase	:					List[Any]	=	WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE      ,     feature_extractor=_SCREAMING_SNAKE_CASE      ,     decoder=_SCREAMING_SNAKE_CASE						)

    processor.save_pretrained(self.tmpdirname						)
    UpperCAmelCase	:					int	=	WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname						)

    # tokenizer
    self.assertEqual(processor.tokenizer.get_vocab()      ,     tokenizer.get_vocab()						)
    self.assertIsInstance(processor.tokenizer      ,     _SCREAMING_SNAKE_CASE						)

    # feature extractor
    self.assertEqual(processor.feature_extractor.to_json_string()      ,     feature_extractor.to_json_string()						)
    self.assertIsInstance(processor.feature_extractor      ,     _SCREAMING_SNAKE_CASE						)

    # decoder
    self.assertEqual(processor.decoder._alphabet.labels      ,     decoder._alphabet.labels						)
    self.assertEqual(
        processor.decoder.model_container[decoder._model_key]._unigram_set      ,     decoder.model_container[decoder._model_key]._unigram_set      ,     )
    self.assertIsInstance(processor.decoder      ,     _SCREAMING_SNAKE_CASE						)

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





    '''simple docstring'''




    UpperCAmelCase	:					List[Any]	=	WavaVecaProcessorWithLM(
        tokenizer=self.get_tokenizer()      ,     feature_extractor=self.get_feature_extractor()      ,     decoder=self.get_decoder()						)
    processor.save_pretrained(self.tmpdirname						)

    # make sure that error is thrown when decoder alphabet doesn't match
    UpperCAmelCase	:					Optional[int]	=	WavaVecaProcessorWithLM.from_pretrained(
        self.tmpdirname      ,     alpha=5.0      ,     beta=3.0      ,     score_boundary=-7.0      ,     unk_score_offset=3						)

    # decoder
    self.assertEqual(processor.language_model.alpha      ,     5.0						)
    self.assertEqual(processor.language_model.beta      ,     3.0						)
    self.assertEqual(processor.language_model.score_boundary      ,     -7.0						)
    self.assertEqual(processor.language_model.unk_score_offset      ,     3						)

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





    '''simple docstring'''




    UpperCAmelCase	:					List[Any]	=	self.get_tokenizer()
    # add token to trigger raise
    tokenizer.add_tokens(["""xx"""]						)
    with self.assertRaisesRegex(_SCREAMING_SNAKE_CASE      ,     """include"""						):
     WavaVecaProcessorWithLM(
         tokenizer=_SCREAMING_SNAKE_CASE      ,     feature_extractor=self.get_feature_extractor()      ,     decoder=self.get_decoder()						)

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





    '''simple docstring'''




    UpperCAmelCase	:					Dict	=	self.get_feature_extractor()
    UpperCAmelCase	:					int	=	self.get_tokenizer()
    UpperCAmelCase	:					List[str]	=	self.get_decoder()

    UpperCAmelCase	:					List[str]	=	WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE      ,     feature_extractor=_SCREAMING_SNAKE_CASE      ,     decoder=_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Dict	=	floats_list((3, 1000)						)

    UpperCAmelCase	:					Optional[int]	=	feature_extractor(_SCREAMING_SNAKE_CASE      ,     return_tensors="""np"""						)
    UpperCAmelCase	:					int	=	processor(_SCREAMING_SNAKE_CASE      ,     return_tensors="""np"""						)

    for key in input_feat_extract.keys():
     self.assertAlmostEqual(input_feat_extract[key].sum()      ,     input_processor[key].sum()      ,     delta=1E-2						)

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





    '''simple docstring'''




    UpperCAmelCase	:					Tuple	=	self.get_feature_extractor()
    UpperCAmelCase	:					Union[str, Any]	=	self.get_tokenizer()
    UpperCAmelCase	:					str	=	self.get_decoder()

    UpperCAmelCase	:					int	=	WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE      ,     feature_extractor=_SCREAMING_SNAKE_CASE      ,     decoder=_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Optional[Any]	=	"""This is a test string"""

    UpperCAmelCase	:					Optional[Any]	=	processor(text=_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Union[str, Any]	=	tokenizer(_SCREAMING_SNAKE_CASE						)

    for key in encoded_tok.keys():
     self.assertListEqual(encoded_tok[key]      ,     encoded_processor[key]						)

   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE=(2, 10, 16)      ,     _SCREAMING_SNAKE_CASE=77						)  ->      Dict:





    '''simple docstring'''




    np.random.seed(_SCREAMING_SNAKE_CASE						)
    return np.random.rand(*_SCREAMING_SNAKE_CASE						)

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





    '''simple docstring'''




    UpperCAmelCase	:					Optional[int]	=	self.get_feature_extractor()
    UpperCAmelCase	:					int	=	self.get_tokenizer()
    UpperCAmelCase	:					Union[str, Any]	=	self.get_decoder()

    UpperCAmelCase	:					Optional[Any]	=	WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE      ,     feature_extractor=_SCREAMING_SNAKE_CASE      ,     decoder=_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Tuple	=	self._get_dummy_logits(shape=(10, 16)      ,     seed=13						)

    UpperCAmelCase	:					Optional[int]	=	processor.decode(_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					str	=	decoder.decode_beams(_SCREAMING_SNAKE_CASE						)[0]

    self.assertEqual(decoded_decoder[0]      ,     decoded_processor.text						)
    self.assertEqual("""</s> <s> </s>"""      ,     decoded_processor.text						)
    self.assertEqual(decoded_decoder[-2]      ,     decoded_processor.logit_score						)
    self.assertEqual(decoded_decoder[-1]      ,     decoded_processor.lm_score						)

   @parameterized.expand([[None], ["""fork"""], ["""spawn"""]]						)
   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE						)  ->      Any:





    '''simple docstring'''




    UpperCAmelCase	:					int	=	self.get_feature_extractor()
    UpperCAmelCase	:					Any	=	self.get_tokenizer()
    UpperCAmelCase	:					Optional[int]	=	self.get_decoder()

    UpperCAmelCase	:					Dict	=	WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE      ,     feature_extractor=_SCREAMING_SNAKE_CASE      ,     decoder=_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					int	=	self._get_dummy_logits()

    # note: pool should be instantiated *after* Wav2Vec2ProcessorWithLM.
    #       otherwise, the LM won't be available to the pool's sub-processes.
    # manual logic used to allow parameterized test for both pool=None and pool=Pool(...)
    if pool_context is None:
     UpperCAmelCase	:					Tuple	=	processor.batch_decode(_SCREAMING_SNAKE_CASE						)
    else:
     with get_context(_SCREAMING_SNAKE_CASE						).Pool() as pool:
      UpperCAmelCase	:					List[str]	=	processor.batch_decode(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Optional[Any]	=	list(_SCREAMING_SNAKE_CASE						)

    with get_context("""fork"""						).Pool() as p:
     UpperCAmelCase	:					Union[str, Any]	=	decoder.decode_beams_batch(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	:					Tuple	=	[], [], []
    for beams in decoded_beams:
     texts_decoder.append(beams[0][0]						)
     logit_scores_decoder.append(beams[0][-2]						)
     lm_scores_decoder.append(beams[0][-1]						)

    self.assertListEqual(_SCREAMING_SNAKE_CASE      ,     decoded_processor.text						)
    self.assertListEqual(["""<s> <s> </s>""", """<s> <s> <s>"""]      ,     decoded_processor.text						)
    self.assertListEqual(_SCREAMING_SNAKE_CASE      ,     decoded_processor.logit_score						)
    self.assertListEqual(_SCREAMING_SNAKE_CASE      ,     decoded_processor.lm_score						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      str:





    '''simple docstring'''




    UpperCAmelCase	:					Any	=	self.get_feature_extractor()
    UpperCAmelCase	:					str	=	self.get_tokenizer()
    UpperCAmelCase	:					Any	=	self.get_decoder()

    UpperCAmelCase	:					int	=	WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE      ,     feature_extractor=_SCREAMING_SNAKE_CASE      ,     decoder=_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					List[Any]	=	self._get_dummy_logits()

    UpperCAmelCase	:					Tuple	=	15
    UpperCAmelCase	:					str	=	-20.0
    UpperCAmelCase	:					Any	=	-4.0

    UpperCAmelCase	:					Optional[int]	=	processor.batch_decode(
        _SCREAMING_SNAKE_CASE      ,     beam_width=_SCREAMING_SNAKE_CASE      ,     beam_prune_logp=_SCREAMING_SNAKE_CASE      ,     token_min_logp=_SCREAMING_SNAKE_CASE      ,     )
    UpperCAmelCase	:					List[str]	=	decoded_processor_out.text

    UpperCAmelCase	:					List[str]	=	list(_SCREAMING_SNAKE_CASE						)

    with get_context("""fork"""						).Pool() as pool:
     UpperCAmelCase	:					Union[str, Any]	=	decoder.decode_beams_batch(
         _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     beam_width=_SCREAMING_SNAKE_CASE      ,     beam_prune_logp=_SCREAMING_SNAKE_CASE      ,     token_min_logp=_SCREAMING_SNAKE_CASE      ,     )

    UpperCAmelCase	:					str	=	[d[0][0] for d in decoded_decoder_out]
    UpperCAmelCase	:					List[str]	=	[d[0][2] for d in decoded_decoder_out]
    UpperCAmelCase	:					List[Any]	=	[d[0][3] for d in decoded_decoder_out]

    self.assertListEqual(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)
    self.assertListEqual(["""</s> <s> <s>""", """<s> <s> <s>"""]      ,     _SCREAMING_SNAKE_CASE						)

    self.assertTrue(np.array_equal(_SCREAMING_SNAKE_CASE      ,     decoded_processor_out.logit_score						)						)
    self.assertTrue(np.allclose([-20.054, -18.447]      ,     _SCREAMING_SNAKE_CASE      ,     atol=1E-3						)						)

    self.assertTrue(np.array_equal(_SCREAMING_SNAKE_CASE      ,     decoded_processor_out.lm_score						)						)
    self.assertTrue(np.allclose([-15.554, -13.9474]      ,     _SCREAMING_SNAKE_CASE      ,     atol=1E-3						)						)

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





    '''simple docstring'''




    UpperCAmelCase	:					List[str]	=	self.get_feature_extractor()
    UpperCAmelCase	:					int	=	self.get_tokenizer()
    UpperCAmelCase	:					List[Any]	=	self.get_decoder()

    UpperCAmelCase	:					Tuple	=	WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE      ,     feature_extractor=_SCREAMING_SNAKE_CASE      ,     decoder=_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Union[str, Any]	=	self._get_dummy_logits()

    UpperCAmelCase	:					str	=	2.0
    UpperCAmelCase	:					Tuple	=	5.0
    UpperCAmelCase	:					List[Any]	=	-20.0
    UpperCAmelCase	:					Any	=	True

    UpperCAmelCase	:					Union[str, Any]	=	processor.batch_decode(
        _SCREAMING_SNAKE_CASE      ,     alpha=_SCREAMING_SNAKE_CASE      ,     beta=_SCREAMING_SNAKE_CASE      ,     unk_score_offset=_SCREAMING_SNAKE_CASE      ,     lm_score_boundary=_SCREAMING_SNAKE_CASE      ,     )
    UpperCAmelCase	:					str	=	decoded_processor_out.text

    UpperCAmelCase	:					Any	=	list(_SCREAMING_SNAKE_CASE						)
    decoder.reset_params(
        alpha=_SCREAMING_SNAKE_CASE      ,     beta=_SCREAMING_SNAKE_CASE      ,     unk_score_offset=_SCREAMING_SNAKE_CASE      ,     lm_score_boundary=_SCREAMING_SNAKE_CASE      ,     )

    with get_context("""fork"""						).Pool() as pool:
     UpperCAmelCase	:					int	=	decoder.decode_beams_batch(
         _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     )

    UpperCAmelCase	:					Optional[int]	=	[d[0][0] for d in decoded_decoder_out]

    self.assertListEqual(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)
    self.assertListEqual(["""<s> </s> <s> </s> </s>""", """</s> </s> <s> </s> </s>"""]      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Tuple	=	processor.decoder.model_container[processor.decoder._model_key]
    self.assertEqual(lm_model.alpha      ,     2.0						)
    self.assertEqual(lm_model.beta      ,     5.0						)
    self.assertEqual(lm_model.unk_score_offset      ,     -20.0						)
    self.assertEqual(lm_model.score_boundary      ,     _SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    UpperCAmelCase	:					Dict	=	WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""						)

    UpperCAmelCase	:					List[Any]	=	processor.decoder.model_container[processor.decoder._model_key]
    UpperCAmelCase	:					List[Any]	=	Path(language_model._kenlm_model.path.decode("""utf-8"""						)						).parent.parent.absolute()

    UpperCAmelCase	:					Any	=	os.listdir(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					List[Any]	=	["""alphabet.json""", """language_model"""]

    downloaded_decoder_files.sort()
    expected_decoder_files.sort()

    # test that only decoder relevant files from
    # https://huggingface.co/hf-internal-testing/processor_with_lm/tree/main
    # are downloaded and none of the rest (e.g. README.md, ...)
    self.assertListEqual(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)

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





    '''simple docstring'''




    UpperCAmelCase	:					Optional[int]	=	snapshot_download("""hf-internal-testing/processor_with_lm"""						)

    UpperCAmelCase	:					Union[str, Any]	=	WavaVecaProcessorWithLM.from_pretrained(_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Any	=	processor.decoder.model_container[processor.decoder._model_key]
    UpperCAmelCase	:					Union[str, Any]	=	Path(language_model._kenlm_model.path.decode("""utf-8"""						)						).parent.parent.absolute()

    UpperCAmelCase	:					Union[str, Any]	=	os.listdir(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Optional[Any]	=	os.listdir(_SCREAMING_SNAKE_CASE						)

    local_decoder_files.sort()
    expected_decoder_files.sort()

    # test that both decoder form hub and local files in cache are the same
    self.assertListEqual(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					Tuple	=	WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""						)
    UpperCAmelCase	:					Union[str, Any]	=	AutoProcessor.from_pretrained("""hf-internal-testing/processor_with_lm"""						)

    UpperCAmelCase	:					Any	=	floats_list((3, 1000)						)

    UpperCAmelCase	:					int	=	processor_wavaveca(_SCREAMING_SNAKE_CASE      ,     return_tensors="""np"""						)
    UpperCAmelCase	:					List[str]	=	processor_auto(_SCREAMING_SNAKE_CASE      ,     return_tensors="""np"""						)

    for key in input_wavaveca.keys():
     self.assertAlmostEqual(input_wavaveca[key].sum()      ,     input_auto[key].sum()      ,     delta=1E-2						)

    UpperCAmelCase	:					Any	=	self._get_dummy_logits()

    UpperCAmelCase	:					str	=	processor_wavaveca.batch_decode(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Dict	=	processor_auto.batch_decode(_SCREAMING_SNAKE_CASE						)

    self.assertListEqual(decoded_wavaveca.text      ,     decoded_auto.text						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					Dict	=	self.get_feature_extractor()
    UpperCAmelCase	:					Union[str, Any]	=	self.get_tokenizer()
    UpperCAmelCase	:					str	=	self.get_decoder()

    UpperCAmelCase	:					Tuple	=	WavaVecaProcessorWithLM(tokenizer=_SCREAMING_SNAKE_CASE      ,     feature_extractor=_SCREAMING_SNAKE_CASE      ,     decoder=_SCREAMING_SNAKE_CASE						)

    self.assertListEqual(
        processor.model_input_names      ,     feature_extractor.model_input_names      ,     msg="""`processor` and `feature_extractor` model input names do not match"""      ,     )

   @staticmethod
   def 	SCREAMING_SNAKE_CASE     (   _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)  ->      Union[str, Any]:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[int]	=	[d[key] for d in offsets]
    return retrieved_list

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Tuple:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[int]	=	WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""						)
    UpperCAmelCase	:					Tuple	=	self._get_dummy_logits()[0]

    UpperCAmelCase	:					List[Any]	=	processor.decode(_SCREAMING_SNAKE_CASE      ,     output_word_offsets=_SCREAMING_SNAKE_CASE						)
    # check Wav2Vec2CTCTokenizerOutput keys for word
    self.assertEqual(len(outputs.keys()						)      ,     4						)
    self.assertTrue("""text""" in outputs						)
    self.assertTrue("""word_offsets""" in outputs						)
    self.assertTrue(isinstance(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)						)

    self.assertEqual(""" """.join(self.get_from_offsets(outputs["""word_offsets"""]      ,     """word"""						)						)      ,     outputs.text						)
    self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""]      ,     """word"""						)      ,     ["""<s>""", """<s>""", """</s>"""]						)
    self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""]      ,     """start_offset"""						)      ,     [0, 2, 4]						)
    self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""]      ,     """end_offset"""						)      ,     [1, 3, 5]						)

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





    '''simple docstring'''




    UpperCAmelCase	:					Optional[Any]	=	WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""						)
    UpperCAmelCase	:					int	=	self._get_dummy_logits()

    UpperCAmelCase	:					int	=	processor.batch_decode(_SCREAMING_SNAKE_CASE      ,     output_word_offsets=_SCREAMING_SNAKE_CASE						)

    # check Wav2Vec2CTCTokenizerOutput keys for word
    self.assertEqual(len(outputs.keys()						)      ,     4						)
    self.assertTrue("""text""" in outputs						)
    self.assertTrue("""word_offsets""" in outputs						)
    self.assertTrue(isinstance(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)						)

    self.assertListEqual(
        [""" """.join(self.get_from_offsets(_SCREAMING_SNAKE_CASE      ,     """word"""						)						) for o in outputs["""word_offsets"""]]      ,     outputs.text						)
    self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""][0]      ,     """word"""						)      ,     ["""<s>""", """<s>""", """</s>"""]						)
    self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""][0]      ,     """start_offset"""						)      ,     [0, 2, 4]						)
    self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""][0]      ,     """end_offset"""						)      ,     [1, 3, 5]						)




   @slow
   @require_torch
   @require_torchaudio
   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Tuple:





    '''simple docstring'''




    import torch

    UpperCAmelCase	:					Optional[int]	=	load_dataset("""common_voice"""      ,     """en"""      ,     split="""train"""      ,     streaming=_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					int	=	ds.cast_column("""audio"""      ,     datasets.Audio(sampling_rate=16000						)						)
    UpperCAmelCase	:					int	=	iter(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Any	=	next(_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					str	=	AutoProcessor.from_pretrained("""patrickvonplaten/wav2vec2-base-100h-with-lm"""						)
    UpperCAmelCase	:					str	=	WavaVecaForCTC.from_pretrained("""patrickvonplaten/wav2vec2-base-100h-with-lm"""						)

    # compare to filename `common_voice_en_100038.mp3` of dataset viewer on https://huggingface.co/datasets/common_voice/viewer/en/train
    UpperCAmelCase	:					Tuple	=	processor(sample["""audio"""]["""array"""]      ,     return_tensors="""pt"""						).input_values

    with torch.no_grad():
     UpperCAmelCase	:					str	=	model(_SCREAMING_SNAKE_CASE						).logits.cpu().numpy()

    UpperCAmelCase	:					Tuple	=	processor.decode(logits[0]      ,     output_word_offsets=_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Dict	=	model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
    UpperCAmelCase	:					Optional[Any]	=	[
        {
            """start_time""": d["""start_offset"""] * time_offset,
            """end_time""": d["""end_offset"""] * time_offset,
            """word""": d["""word"""],
        }
        for d in output["""word_offsets"""]
    ]

    UpperCAmelCase	:					Tuple	=	"""WHY DOES MILISANDRA LOOK LIKE SHE WANTS TO CONSUME JOHN SNOW ON THE RIVER AT THE WALL"""

    # output words
    self.assertEqual(""" """.join(self.get_from_offsets(_SCREAMING_SNAKE_CASE      ,     """word"""						)						)      ,     _SCREAMING_SNAKE_CASE						)
    self.assertEqual(""" """.join(self.get_from_offsets(_SCREAMING_SNAKE_CASE      ,     """word"""						)						)      ,     output.text						)

    # output times
    UpperCAmelCase	:					int	=	torch.tensor(self.get_from_offsets(_SCREAMING_SNAKE_CASE      ,     """start_time"""						)						)
    UpperCAmelCase	:					Any	=	torch.tensor(self.get_from_offsets(_SCREAMING_SNAKE_CASE      ,     """end_time"""						)						)

    # fmt: off
    UpperCAmelCase	:					List[str]	=	torch.tensor([1.4199, 1.6599, 2.2599, 3.0, 3.24, 3.5999, 3.7999, 4.0999, 4.26, 4.94, 5.28, 5.6599, 5.78, 5.94, 6.32, 6.5399, 6.6599]						)
    UpperCAmelCase	:					str	=	torch.tensor([1.5399, 1.8999, 2.9, 3.16, 3.5399, 3.72, 4.0199, 4.1799, 4.76, 5.1599, 5.5599, 5.6999, 5.86, 6.1999, 6.38, 6.6199, 6.94]						)
    # fmt: on

    self.assertTrue(torch.allclose(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     atol=0.01						)						)
    self.assertTrue(torch.allclose(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     atol=0.01						)						)





"""simple docstring"""
import argparse

import torch

# Step 1. clone https://github.com/microsoft/unilm
# Step 2. git checkout to https://github.com/microsoft/unilm/commit/b94ec76c36f02fb2b0bf0dcb0b8554a2185173cd
# Step 3. cd unilm
# Step 4. ln -s $(realpath wavlm/modules.py) ./  # create simlink
# import classes
from unilm.wavlm.WavLM import WavLM as WavLMOrig
from unilm.wavlm.WavLM import WavLMConfig as WavLMConfigOrig

from transformers import WavLMConfig, WavLMModel, logging


logging.set_verbosity_info()
A:   Optional[int]								=		logging.get_logger(__name__)

A:   Tuple								=		{
    "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.grep_linear": "encoder.layers.*.attention.gru_rel_pos_linear",
    "self_attn.relative_attention_bias": "encoder.layers.*.attention.rel_attn_embed",
    "self_attn.grep_a": "encoder.layers.*.attention.gru_rel_pos_const",
    "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",
    "quantizer.weight_proj": "quantizer.weight_proj",
    "quantizer.vars": "quantizer.codevectors",
    "project_q": "project_q",
    "final_proj": "project_hid",
    "w2v_encoder.proj": "ctc_proj",
    "mask_emb": "masked_spec_embed",
}
A:   List[str]								=		[
    "ctc_proj",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]

def       _snake_case						(						UpperCamelCase						:		str			,		UpperCamelCase						:		Tuple			,		UpperCamelCase						:		Optional[int]			,		UpperCamelCase						:		Tuple			,		UpperCamelCase						:		Any						):
 for attribute in key.split("""."""						):
  UpperCAmelCase	:					Optional[Any]	=	getattr(UpperCamelCase			,		UpperCamelCase						)

 if weight_type is not None:
  UpperCAmelCase	:					List[Any]	=	getattr(UpperCamelCase			,		UpperCamelCase						).shape
 else:
  UpperCAmelCase	:					str	=	hf_pointer.shape

 assert hf_shape == value.shape, (
     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	:					Optional[Any]	=	value
 elif weight_type == "weight_g":
  UpperCAmelCase	:					str	=	value
 elif weight_type == "weight_v":
  UpperCAmelCase	:					Union[str, Any]	=	value
 elif weight_type == "bias":
  UpperCAmelCase	:					str	=	value
 else:
  UpperCAmelCase	:					Union[str, Any]	=	value

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

def       _snake_case						(						UpperCamelCase						:		List[Any]			,		UpperCamelCase						:		Optional[Any]						):
 UpperCAmelCase	:					Tuple	=	[]
 UpperCAmelCase	:					Any	=	fairseq_model.state_dict()

 UpperCAmelCase	:					Tuple	=	hf_model.feature_extractor

 for name, value in fairseq_dict.items():
  UpperCAmelCase	:					str	=	False
  if "conv_layers" in name:
   load_conv_layer(
       UpperCamelCase			,		UpperCamelCase			,		UpperCamelCase			,		UpperCamelCase			,		hf_model.config.feat_extract_norm == """group"""			,		)
   UpperCAmelCase	:					Union[str, Any]	=	True
  else:
   for key, mapped_key in MAPPING.items():
    if key in name or key.split("""w2v_model."""						)[-1] == name.split("""."""						)[0]:
     UpperCAmelCase	:					Dict	=	True
     if "*" in mapped_key:
      UpperCAmelCase	:					str	=	name.split(UpperCamelCase						)[0].split("""."""						)[-2]
      UpperCAmelCase	:					Tuple	=	mapped_key.replace("""*"""			,		UpperCamelCase						)
     if "weight_g" in name:
      UpperCAmelCase	:					Any	=	"""weight_g"""
     elif "weight_v" in name:
      UpperCAmelCase	:					Optional[Any]	=	"""weight_v"""
     elif "bias" in name and "relative_attention_bias" not in name:
      UpperCAmelCase	:					Union[str, Any]	=	"""bias"""
     elif "weight" in name:
      # TODO: don't match quantizer.weight_proj
      UpperCAmelCase	:					str	=	"""weight"""
     else:
      UpperCAmelCase	:					Optional[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       _snake_case						(						UpperCamelCase						:		Optional[Any]			,		UpperCamelCase						:		Dict			,		UpperCamelCase						:		Tuple			,		UpperCamelCase						:		Any			,		UpperCamelCase						:		Any						):
 UpperCAmelCase	:					str	=	full_name.split("""conv_layers."""						)[-1]
 UpperCAmelCase	:					Dict	=	name.split("""."""						)
 UpperCAmelCase	:					List[str]	=	int(items[0]						)
 UpperCAmelCase	:					Any	=	int(items[1]						)

 if type_id == 0:
  if "bias" in name:
   assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
       F"{full_name} has size {value.shape}, but"
       F" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found."
   )
   UpperCAmelCase	:					Optional[Any]	=	value
   logger.info(F"Feat extract conv layer {layer_id} was initialized from {full_name}."						)
  elif "weight" in name:
   assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
       F"{full_name} has size {value.shape}, but"
       F" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found."
   )
   UpperCAmelCase	:					Tuple	=	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:
   assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
       F"{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was"
       " found."
   )
   UpperCAmelCase	:					str	=	value
   logger.info(F"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."						)
  elif "weight" in name:
   assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
       F"{full_name} has size {value.shape}, but"
       F" {feature_extractor[layer_id].layer_norm.weight.data.shape} was found."
   )
   UpperCAmelCase	:					Optional[Any]	=	value
   logger.info(F"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."						)
 else:
  unused_weights.append(UpperCamelCase						)

@torch.no_grad()
def       _snake_case						(						UpperCamelCase						:		int			,		UpperCamelCase						:		List[Any]			,		UpperCamelCase						:		List[Any]=None						):
 # load the pre-trained checkpoints
 UpperCAmelCase	:					List[Any]	=	torch.load(UpperCamelCase						)
 UpperCAmelCase	:					List[str]	=	WavLMConfigOrig(checkpoint["""cfg"""]						)
 UpperCAmelCase	:					Optional[int]	=	WavLMOrig(UpperCamelCase						)
 model.load_state_dict(checkpoint["""model"""]						)
 model.eval()

 if config_path is not None:
  UpperCAmelCase	:					List[str]	=	WavLMConfig.from_pretrained(UpperCamelCase						)
 else:
  UpperCAmelCase	:					List[Any]	=	WavLMConfig()

 UpperCAmelCase	:					Any	=	WavLMModel(UpperCamelCase						)

 recursively_load_weights(UpperCamelCase			,		UpperCamelCase						)

 hf_wavlm.save_pretrained(UpperCamelCase						)


if __name__ == "__main__":
      A:   int								=		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("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
      A:   Tuple								=		parser.parse_args()
      convert_wavlm_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)
