





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    __UpperCAmelCase (    self	:	Union[str, Any]  )						->					Any:
         _lowercase							= '| <pad> <unk> <s> </s> a b c d e f g h i j k'.split()
         _lowercase							= dict(zip(__A						,range(len(__A  )  )  )  )

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

         _lowercase							= tempfile.mkdtemp()
         _lowercase							= os.path.join(self.tmpdirname						,VOCAB_FILES_NAMES['vocab_file']  )
         _lowercase							= 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
         _lowercase							= 'hf-internal-testing/ngram-beam-search-decoder'







       def    __UpperCAmelCase (    self	:	List[Any]						,**__A	:	Dict  )						->					Tuple:
         _lowercase							= self.add_kwargs_tokens_map.copy()
         kwargs.update(__A  )
         return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname						,**__A  )







       def    __UpperCAmelCase (    self	:	Union[str, Any]						,**__A	:	int  )						->					Dict:
         return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname						,**__A  )







       def    __UpperCAmelCase (    self	:	Union[str, Any]						,**__A	:	int  )						->					Any:
         return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name						,**__A  )







       def    __UpperCAmelCase (    self	:	str  )						->					Union[str, Any]:
         shutil.rmtree(self.tmpdirname  )







       def    __UpperCAmelCase (    self	:	Any  )						->					List[str]:
         _lowercase							= self.get_tokenizer()
         _lowercase							= self.get_feature_extractor()
         _lowercase							= self.get_decoder()

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

         processor.save_pretrained(self.tmpdirname  )
         _lowercase							= 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    __UpperCAmelCase (    self	:	List[str]  )						->					Union[str, Any]:
         _lowercase							= 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
         _lowercase							= 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    __UpperCAmelCase (    self	:	Optional[Any]  )						->					Dict:
         _lowercase							= 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    __UpperCAmelCase (    self	:	Tuple  )						->					str:
         _lowercase							= self.get_feature_extractor()
         _lowercase							= self.get_tokenizer()
         _lowercase							= self.get_decoder()

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

         _lowercase							= floats_list((3, 1000)  )

         _lowercase							= feature_extractor(__A						,return_tensors='np'  )
         _lowercase							= 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    __UpperCAmelCase (    self	:	Optional[int]  )						->					Dict:
         _lowercase							= self.get_feature_extractor()
         _lowercase							= self.get_tokenizer()
         _lowercase							= self.get_decoder()

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

         _lowercase							= 'This is a test string'

         _lowercase							= processor(text=__A  )

         _lowercase							= tokenizer(__A  )

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







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







       def    __UpperCAmelCase (    self	:	List[str]  )						->					str:
         _lowercase							= self.get_feature_extractor()
         _lowercase							= self.get_tokenizer()
         _lowercase							= self.get_decoder()

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

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

         _lowercase							= processor.decode(__A  )

         _lowercase							= 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    __UpperCAmelCase (    self	:	Tuple						,__A	:	int  )						->					List[Any]:
         _lowercase							= self.get_feature_extractor()
         _lowercase							= self.get_tokenizer()
         _lowercase							= self.get_decoder()

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

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

         _lowercase							= list(__A  )

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

         _lowercase						,		_lowercase						,		_lowercase							= [], [], []
         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    __UpperCAmelCase (    self	:	Union[str, Any]  )						->					Any:
         _lowercase							= self.get_feature_extractor()
         _lowercase							= self.get_tokenizer()
         _lowercase							= self.get_decoder()

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

         _lowercase							= self._get_dummy_logits()

         _lowercase							= 15
         _lowercase							= -20.0
         _lowercase							= -4.0

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

         _lowercase							= list(__A  )

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

         _lowercase							= [d[0][0] for d in decoded_decoder_out]
         _lowercase							= [d[0][2] for d in decoded_decoder_out]
         _lowercase							= [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.9474]						,__A						,atol=1e-3  )  )







       def    __UpperCAmelCase (    self	:	Union[str, Any]  )						->					Dict:
         _lowercase							= self.get_feature_extractor()
         _lowercase							= self.get_tokenizer()
         _lowercase							= self.get_decoder()

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

         _lowercase							= self._get_dummy_logits()

         _lowercase							= 2.0
         _lowercase							= 5.0
         _lowercase							= -20.0
         _lowercase							= True

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

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

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

         _lowercase							= [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  )
         _lowercase							= 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    __UpperCAmelCase (    self	:	int  )						->					Dict:
         _lowercase							= WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'  )

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

         _lowercase							= os.listdir(__A  )
         _lowercase							= ['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    __UpperCAmelCase (    self	:	Optional[int]  )						->					Optional[Any]:
         _lowercase							= snapshot_download('hf-internal-testing/processor_with_lm'  )

         _lowercase							= WavaVecaProcessorWithLM.from_pretrained(__A  )

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

         _lowercase							= os.listdir(__A  )
         _lowercase							= 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    __UpperCAmelCase (    self	:	Optional[int]  )						->					Any:
         _lowercase							= WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'  )
         _lowercase							= AutoProcessor.from_pretrained('hf-internal-testing/processor_with_lm'  )

         _lowercase							= floats_list((3, 1000)  )

         _lowercase							= processor_wavaveca(__A						,return_tensors='np'  )
         _lowercase							= 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  )

         _lowercase							= self._get_dummy_logits()

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

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







       def    __UpperCAmelCase (    self	:	List[str]  )						->					List[Any]:
         _lowercase							= self.get_feature_extractor()
         _lowercase							= self.get_tokenizer()
         _lowercase							= self.get_decoder()

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







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

         _lowercase							= 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    __UpperCAmelCase (    self	:	List[str]  )						->					Union[str, Any]:
         _lowercase							= WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'  )
         _lowercase							= self._get_dummy_logits()

         _lowercase							= 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    __UpperCAmelCase (    self	:	List[str]  )						->					Any:
         import torch

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

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

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

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

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

         _lowercase							= '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
         _lowercase							= torch.tensor(self.get_from_offsets(__A						,'start_time'  )  )
         _lowercase							= torch.tensor(self.get_from_offsets(__A						,'end_time'  )  )

         # fmt: off
         _lowercase							= 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]  )
         _lowercase							= 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(__A						,__A						,atol=0.01  )  )
         self.assertTrue(torch.allclose(__A						,__A						,atol=0.01  )  )






from ....utils import logging


snake_case 				=       logging.get_logger(__name__)





class 				A_				(      UpperCAmelCase	):







       """simple docstring"""





       def __init__(    self	:	Union[str, Any]						,__A	:	int						,__A	:	List[Any]=None						,__A	:	List[str]=2048  )						->					List[Any]:
         _lowercase							= config.__dict__
         _lowercase							= modal_hidden_size
         if num_labels:
           _lowercase							= num_labels