


"""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       __A				(       unittest.TestCase			):






 """simple docstring"""




 def       SCREAMING_SNAKE_CASE							(		self						)					->     Optional[Any]:
     a 		='''| <pad> <unk> <s> </s> a b c d e f g h i j k'''.split()
     a 		=dict(zip(__A ,    range(len(__A						)						)						)						)

     a 		={
         '''unk_token''': '''<unk>''',
         '''bos_token''': '''<s>''',
         '''eos_token''': '''</s>''',
     }
     a 		={
         '''feature_size''': 1,
         '''padding_value''': 0.0,
         '''sampling_rate''': 1_6000,
         '''return_attention_mask''': False,
         '''do_normalize''': True,
     }

     a 		=tempfile.mkdtemp()
     a 		=os.path.join(self.tmpdirname ,    VOCAB_FILES_NAMES['''vocab_file''']						)
     a 		=os.path.join(self.tmpdirname ,    __A						)
     with open(self.vocab_file ,    '''w''' ,    encoding='''utf-8'''						) as fp:
         fp.write(json.dumps(__A						) + '''\n'''						)

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

     # load decoder from hub
     a 		='''hf-internal-testing/ngram-beam-search-decoder'''


 def       SCREAMING_SNAKE_CASE							(		self ,    **__A						)					->     Dict:
     a 		=self.add_kwargs_tokens_map.copy()
     kwargs.update(__A						)
     return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname ,    **__A						)


 def       SCREAMING_SNAKE_CASE							(		self ,    **__A						)					->     Optional[Any]:
     return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname ,    **__A						)


 def       SCREAMING_SNAKE_CASE							(		self ,    **__A						)					->     Dict:
     return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name ,    **__A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Tuple:
     shutil.rmtree(self.tmpdirname						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     int:
     a 		=self.get_tokenizer()
     a 		=self.get_feature_extractor()
     a 		=self.get_decoder()

     a 		=WavaVecaProcessorWithLM(tokenizer=__A ,    feature_extractor=__A ,    decoder=__A						)

     processor.save_pretrained(self.tmpdirname						)
     a 		=WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname						)

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

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

     # 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 ,    __A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Dict:
     a 		=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
     a 		=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						)					->     List[Any]:
     a 		=self.get_tokenizer()
     # add token to trigger raise
     tokenizer.add_tokens(['''xx''']						)
     with self.assertRaisesRegex(__A ,    '''include'''						):
         WavaVecaProcessorWithLM(
             tokenizer=__A ,    feature_extractor=self.get_feature_extractor() ,    decoder=self.get_decoder()						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Union[str, Any]:
     a 		=self.get_feature_extractor()
     a 		=self.get_tokenizer()
     a 		=self.get_decoder()

     a 		=WavaVecaProcessorWithLM(tokenizer=__A ,    feature_extractor=__A ,    decoder=__A						)

     a 		=floats_list((3, 1000)						)

     a 		=feature_extractor(__A ,    return_tensors='''np'''						)
     a 		=processor(__A ,    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						)					->     str:
     a 		=self.get_feature_extractor()
     a 		=self.get_tokenizer()
     a 		=self.get_decoder()

     a 		=WavaVecaProcessorWithLM(tokenizer=__A ,    feature_extractor=__A ,    decoder=__A						)

     a 		='''This is a test string'''

     a 		=processor(text=__A						)

     a 		=tokenizer(__A						)

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


 def       SCREAMING_SNAKE_CASE							(		self ,    __A=(2, 10, 16) ,    __A=77						)					->     Optional[int]:
     np.random.seed(__A						)
     return np.random.rand(*__A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     int:
     a 		=self.get_feature_extractor()
     a 		=self.get_tokenizer()
     a 		=self.get_decoder()

     a 		=WavaVecaProcessorWithLM(tokenizer=__A ,    feature_extractor=__A ,    decoder=__A						)

     a 		=self._get_dummy_logits(shape=(10, 16) ,    seed=13						)

     a 		=processor.decode(__A						)

     a 		=decoder.decode_beams(__A						)[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 ,    __A						)					->     List[str]:
     a 		=self.get_feature_extractor()
     a 		=self.get_tokenizer()
     a 		=self.get_decoder()

     a 		=WavaVecaProcessorWithLM(tokenizer=__A ,    feature_extractor=__A ,    decoder=__A						)

     a 		=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:
         a 		=processor.batch_decode(__A						)
     else:
         with get_context(__A						).Pool() as pool:
             a 		=processor.batch_decode(__A ,    __A						)

     a 		=list(__A						)

     with get_context('''fork'''						).Pool() as p:
         a 		=decoder.decode_beams_batch(__A ,    __A						)

     a		,      a		,      a 		=[], [], []
     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(__A ,    decoded_processor.text						)
     self.assertListEqual(['''<s> <s> </s>''', '''<s> <s> <s>'''] ,    decoded_processor.text						)
     self.assertListEqual(__A ,    decoded_processor.logit_score						)
     self.assertListEqual(__A ,    decoded_processor.lm_score						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     List[Any]:
     a 		=self.get_feature_extractor()
     a 		=self.get_tokenizer()
     a 		=self.get_decoder()

     a 		=WavaVecaProcessorWithLM(tokenizer=__A ,    feature_extractor=__A ,    decoder=__A						)

     a 		=self._get_dummy_logits()

     a 		=15
     a 		=-20.0
     a 		=-4.0

     a 		=processor.batch_decode(
         __A ,    beam_width=__A ,    beam_prune_logp=__A ,    token_min_logp=__A ,    )
     a 		=decoded_processor_out.text

     a 		=list(__A						)

     with get_context('''fork'''						).Pool() as pool:
         a 		=decoder.decode_beams_batch(
             __A ,    __A ,    beam_width=__A ,    beam_prune_logp=__A ,    token_min_logp=__A ,    )

     a 		=[d[0][0] for d in decoded_decoder_out]
     a 		=[d[0][2] for d in decoded_decoder_out]
     a 		=[d[0][3] for d in decoded_decoder_out]

     self.assertListEqual(__A ,    __A						)
     self.assertListEqual(['''</s> <s> <s>''', '''<s> <s> <s>'''] ,    __A						)

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

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


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Union[str, Any]:
     a 		=self.get_feature_extractor()
     a 		=self.get_tokenizer()
     a 		=self.get_decoder()

     a 		=WavaVecaProcessorWithLM(tokenizer=__A ,    feature_extractor=__A ,    decoder=__A						)

     a 		=self._get_dummy_logits()

     a 		=2.0
     a 		=5.0
     a 		=-20.0
     a 		=True

     a 		=processor.batch_decode(
         __A ,    alpha=__A ,    beta=__A ,    unk_score_offset=__A ,    lm_score_boundary=__A ,    )
     a 		=decoded_processor_out.text

     a 		=list(__A						)
     decoder.reset_params(
         alpha=__A ,    beta=__A ,    unk_score_offset=__A ,    lm_score_boundary=__A ,    )

     with get_context('''fork'''						).Pool() as pool:
         a 		=decoder.decode_beams_batch(
             __A ,    __A ,    )

     a 		=[d[0][0] for d in decoded_decoder_out]

     self.assertListEqual(__A ,    __A						)
     self.assertListEqual(['''<s> </s> <s> </s> </s>''', '''</s> </s> <s> </s> </s>'''] ,    __A						)
     a 		=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 ,    __A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Union[str, Any]:
     a 		=WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''						)

     a 		=processor.decoder.model_container[processor.decoder._model_key]
     a 		=Path(language_model._kenlm_model.path.decode('''utf-8'''						)						).parent.parent.absolute()

     a 		=os.listdir(__A						)
     a 		=['''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(__A ,    __A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     List[Any]:
     a 		=snapshot_download('''hf-internal-testing/processor_with_lm'''						)

     a 		=WavaVecaProcessorWithLM.from_pretrained(__A						)

     a 		=processor.decoder.model_container[processor.decoder._model_key]
     a 		=Path(language_model._kenlm_model.path.decode('''utf-8'''						)						).parent.parent.absolute()

     a 		=os.listdir(__A						)
     a 		=os.listdir(__A						)

     local_decoder_files.sort()
     expected_decoder_files.sort()

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


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Any:
     a 		=WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''						)
     a 		=AutoProcessor.from_pretrained('''hf-internal-testing/processor_with_lm'''						)

     a 		=floats_list((3, 1000)						)

     a 		=processor_wavaveca(__A ,    return_tensors='''np'''						)
     a 		=processor_auto(__A ,    return_tensors='''np'''						)

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

     a 		=self._get_dummy_logits()

     a 		=processor_wavaveca.batch_decode(__A						)
     a 		=processor_auto.batch_decode(__A						)

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


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Any:
     a 		=self.get_feature_extractor()
     a 		=self.get_tokenizer()
     a 		=self.get_decoder()

     a 		=WavaVecaProcessorWithLM(tokenizer=__A ,    feature_extractor=__A ,    decoder=__A						)

     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							(		__A ,    __A						)					->     List[Any]:
     a 		=[d[key] for d in offsets]
     return retrieved_list


 def       SCREAMING_SNAKE_CASE							(		self						)					->     List[str]:
     a 		=WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''						)
     a 		=self._get_dummy_logits()[0]

     a 		=processor.decode(__A ,    output_word_offsets=__A						)
     # 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(__A ,    __A						)						)

     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						)					->     Tuple:
     a 		=WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''						)
     a 		=self._get_dummy_logits()

     a 		=processor.batch_decode(__A ,    output_word_offsets=__A						)

     # 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(__A ,    __A						)						)

     self.assertListEqual(
         [''' '''.join(self.get_from_offsets(__A ,    '''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						)					->     Optional[Any]:
     import torch

     a 		=load_dataset('''common_voice''' ,    '''en''' ,    split='''train''' ,    streaming=__A						)
     a 		=ds.cast_column('''audio''' ,    datasets.Audio(sampling_rate=1_6000						)						)
     a 		=iter(__A						)
     a 		=next(__A						)

     a 		=AutoProcessor.from_pretrained('''patrickvonplaten/wav2vec2-base-100h-with-lm'''						)
     a 		=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
     a 		=processor(sample['''audio''']['''array'''] ,    return_tensors='''pt'''						).input_values

     with torch.no_grad():
         a 		=model(__A						).logits.cpu().numpy()

     a 		=processor.decode(logits[0] ,    output_word_offsets=__A						)

     a 		=model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
     a 		=[
         {
             '''start_time''': d['''start_offset'''] * time_offset,
             '''end_time''': d['''end_offset'''] * time_offset,
             '''word''': d['''word'''],
         }
         for d in output['''word_offsets''']
     ]

     a 		='''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(__A ,    '''word'''						)						) ,    __A						)
     self.assertEqual(''' '''.join(self.get_from_offsets(__A ,    '''word'''						)						) ,    output.text						)

     # output times
     a 		=torch.tensor(self.get_from_offsets(__A ,    '''start_time'''						)						)
     a 		=torch.tensor(self.get_from_offsets(__A ,    '''end_time'''						)						)

     # fmt: off
     a 		=torch.tensor([1.4_199, 1.6_599, 2.2_599, 3.0, 3.24, 3.5_999, 3.7_999, 4.0_999, 4.26, 4.94, 5.28, 5.6_599, 5.78, 5.94, 6.32, 6.5_399, 6.6_599]						)
     a 		=torch.tensor([1.5_399, 1.8_999, 2.9, 3.16, 3.5_399, 3.72, 4.0_199, 4.1_799, 4.76, 5.1_599, 5.5_599, 5.6_999, 5.86, 6.1_999, 6.38, 6.6_199, 6.94]						)
     # fmt: on

     self.assertTrue(torch.allclose(__A ,    __A ,    atol=0.01						)						)
     self.assertTrue(torch.allclose(__A ,    __A ,    atol=0.01						)						)
'''simple docstring'''

import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__a: List[str]          =   logging.get_logger(__name__)



class 			UpperCAmelCase  (						a__					):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE							=    "encoder-decoder"
       SCREAMING_SNAKE_CASE							=    True

       def __init__(						self	,   **__lowerCAmelCase )  ->  int:
        super().__init__(**__lowerCAmelCase )
        assert (
            "encoder" in kwargs and "decoder" in kwargs
        ), "Config has to be initialized with encoder and decoder config"
        lowercase__				:    Optional[int]  =					kwargs.pop('''encoder''' )
        lowercase__				:    Union[str, Any]  =					encoder_config.pop('''model_type''' )
        lowercase__				:    Any  =					kwargs.pop('''decoder''' )
        lowercase__				:    Any  =					decoder_config.pop('''model_type''' )

        from ..auto.configuration_auto import AutoConfig

        lowercase__				:    Union[str, Any]  =					AutoConfig.for_model(__lowerCAmelCase	,   **__lowerCAmelCase )
        lowercase__				:    Optional[Any]  =					AutoConfig.for_model(__lowerCAmelCase	,   **__lowerCAmelCase )
        lowercase__				:    Tuple  =					True

       @classmethod
       def _lowerCAmelCase(						cls	,   __lowerCAmelCase	,   __lowerCAmelCase	,   **__lowerCAmelCase )  ->  PretrainedConfig:
        logger.info('''Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config''' )
        lowercase__				:    Union[str, Any]  =					True
        lowercase__				:    Any  =					True

        return cls(encoder=encoder_config.to_dict()	,   decoder=decoder_config.to_dict()	,   **__lowerCAmelCase )


       def _lowerCAmelCase(						self )  ->  Any:
        lowercase__				:    Any  =					copy.deepcopy(self.__dict__ )
        lowercase__				:    Optional[Any]  =					self.encoder.to_dict()
        lowercase__				:    Tuple  =					self.decoder.to_dict()
        lowercase__				:    Dict  =					self.__class__.model_type
        return output

