

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




  def 						A__  (       self      )						->				Union[str, Any]:
        '''simple docstring'''






        _lowercase			 ='| <pad> <unk> <s> </s> a b c d e f g h i j k'.split()
        _lowercase			 =dict(zip(lowerCAmelCase							,     range(len(lowerCAmelCase      )      )      )      )

        _lowercase			 ={
            'unk_token': '<unk>',
            'bos_token': '<s>',
            'eos_token': '</s>',
        }
        _lowercase			 ={
            'feature_size': 1,
            'padding_value': 0.0,
            'sampling_rate': 16_000,
            '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							,     lowerCAmelCase      )
        with open(self.vocab_file							,     'w'							,     encoding='utf-8'      ) as fp:
              fp.write(json.dumps(lowerCAmelCase      ) + '\n'      )

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

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




  def 						A__  (       self							,     **lowerCAmelCase      )						->				Any:
        '''simple docstring'''






        _lowercase			 =self.add_kwargs_tokens_map.copy()
        kwargs.update(lowerCAmelCase      )
        return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname							,     **lowerCAmelCase      )




  def 						A__  (       self							,     **lowerCAmelCase      )						->				Optional[int]:
        '''simple docstring'''






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




  def 						A__  (       self							,     **lowerCAmelCase      )						->				Dict:
        '''simple docstring'''






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




  def 						A__  (       self      )						->				str:
        '''simple docstring'''






        shutil.rmtree(self.tmpdirname      )




  def 						A__  (       self      )						->				Optional[Any]:
        '''simple docstring'''






        _lowercase			 =self.get_tokenizer()
        _lowercase			 =self.get_feature_extractor()
        _lowercase			 =self.get_decoder()

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

        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							,     lowerCAmelCase      )

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

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




  def 						A__  (       self      )						->				Dict:
        '''simple docstring'''






        _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 						A__  (       self      )						->				Any:
        '''simple docstring'''






        _lowercase			 =self.get_tokenizer()
        # add token to trigger raise
        tokenizer.add_tokens(['xx']      )
        with self.assertRaisesRegex(lowerCAmelCase							,     'include'      ):
              WavaVecaProcessorWithLM(
                  tokenizer=lowerCAmelCase							,     feature_extractor=self.get_feature_extractor()							,     decoder=self.get_decoder()      )




  def 						A__  (       self      )						->				Tuple:
        '''simple docstring'''






        _lowercase			 =self.get_feature_extractor()
        _lowercase			 =self.get_tokenizer()
        _lowercase			 =self.get_decoder()

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

        _lowercase			 =floats_list((3, 1_000)      )

        _lowercase			 =feature_extractor(lowerCAmelCase							,     return_tensors='np'      )
        _lowercase			 =processor(lowerCAmelCase							,     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 						A__  (       self      )						->				int:
        '''simple docstring'''






        _lowercase			 =self.get_feature_extractor()
        _lowercase			 =self.get_tokenizer()
        _lowercase			 =self.get_decoder()

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

        _lowercase			 ='This is a test string'

        _lowercase			 =processor(text=lowerCAmelCase      )

        _lowercase			 =tokenizer(lowerCAmelCase      )

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




  def 						A__  (       self							,     lowerCAmelCase=(2, 10, 16)							,     lowerCAmelCase=77      )						->				Optional[Any]:
        '''simple docstring'''






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




  def 						A__  (       self      )						->				List[Any]:
        '''simple docstring'''






        _lowercase			 =self.get_feature_extractor()
        _lowercase			 =self.get_tokenizer()
        _lowercase			 =self.get_decoder()

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

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

        _lowercase			 =processor.decode(lowerCAmelCase      )

        _lowercase			 =decoder.decode_beams(lowerCAmelCase      )[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 						A__  (       self							,     lowerCAmelCase      )						->				int:
        '''simple docstring'''






        _lowercase			 =self.get_feature_extractor()
        _lowercase			 =self.get_tokenizer()
        _lowercase			 =self.get_decoder()

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

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

        _lowercase			 =list(lowerCAmelCase      )

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

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




  def 						A__  (       self      )						->				int:
        '''simple docstring'''






        _lowercase			 =self.get_feature_extractor()
        _lowercase			 =self.get_tokenizer()
        _lowercase			 =self.get_decoder()

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

        _lowercase			 =self._get_dummy_logits()

        _lowercase			 =15
        _lowercase			 =-20.0
        _lowercase			 =-4.0

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

        _lowercase			 =list(lowerCAmelCase      )

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

        _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(lowerCAmelCase							,     lowerCAmelCase      )
        self.assertListEqual(['</s> <s> <s>', '<s> <s> <s>']							,     lowerCAmelCase      )

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

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




  def 						A__  (       self      )						->				Union[str, Any]:
        '''simple docstring'''






        _lowercase			 =self.get_feature_extractor()
        _lowercase			 =self.get_tokenizer()
        _lowercase			 =self.get_decoder()

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

        _lowercase			 =self._get_dummy_logits()

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

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

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

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

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

        self.assertListEqual(lowerCAmelCase							,     lowerCAmelCase      )
        self.assertListEqual(['<s> </s> <s> </s> </s>', '</s> </s> <s> </s> </s>']							,     lowerCAmelCase      )
        _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							,     lowerCAmelCase      )




  def 						A__  (       self      )						->				Any:
        '''simple docstring'''






        _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(lowerCAmelCase      )
        _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(lowerCAmelCase							,     lowerCAmelCase      )




  def 						A__  (       self      )						->				Any:
        '''simple docstring'''






        _lowercase			 =snapshot_download('hf-internal-testing/processor_with_lm'      )

        _lowercase			 =WavaVecaProcessorWithLM.from_pretrained(lowerCAmelCase      )

        _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(lowerCAmelCase      )
        _lowercase			 =os.listdir(lowerCAmelCase      )

        local_decoder_files.sort()
        expected_decoder_files.sort()

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




  def 						A__  (       self      )						->				Union[str, Any]:
        '''simple docstring'''






        _lowercase			 =WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'      )
        _lowercase			 =AutoProcessor.from_pretrained('hf-internal-testing/processor_with_lm'      )

        _lowercase			 =floats_list((3, 1_000)      )

        _lowercase			 =processor_wavaveca(lowerCAmelCase							,     return_tensors='np'      )
        _lowercase			 =processor_auto(lowerCAmelCase							,     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(lowerCAmelCase      )
        _lowercase			 =processor_auto.batch_decode(lowerCAmelCase      )

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




  def 						A__  (       self      )						->				int:
        '''simple docstring'''






        _lowercase			 =self.get_feature_extractor()
        _lowercase			 =self.get_tokenizer()
        _lowercase			 =self.get_decoder()

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

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




  @staticmethod
  def 						A__  (       lowerCAmelCase							,     lowerCAmelCase      )						->				List[str]:
        '''simple docstring'''






        _lowercase			 =[d[key] for d in offsets]
        return retrieved_list




  def 						A__  (       self      )						->				List[Any]:
        '''simple docstring'''






        _lowercase			 =WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'      )
        _lowercase			 =self._get_dummy_logits()[0]

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

        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 						A__  (       self      )						->				Optional[int]:
        '''simple docstring'''






        _lowercase			 =WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'      )
        _lowercase			 =self._get_dummy_logits()

        _lowercase			 =processor.batch_decode(lowerCAmelCase							,     output_word_offsets=lowerCAmelCase      )

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

        self.assertListEqual(
            [' '.join(self.get_from_offsets(lowerCAmelCase							,     '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 						A__  (       self      )						->				List[str]:
        '''simple docstring'''






        import torch

        _lowercase			 =load_dataset('common_voice'							,     'en'							,     split='train'							,     streaming=lowerCAmelCase      )
        _lowercase			 =ds.cast_column('audio'							,     datasets.Audio(sampling_rate=16_000      )      )
        _lowercase			 =iter(lowerCAmelCase      )
        _lowercase			 =next(lowerCAmelCase      )

        _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(lowerCAmelCase      ).logits.cpu().numpy()

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

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

        # output times
        _lowercase			 =torch.tensor(self.get_from_offsets(lowerCAmelCase							,     'start_time'      )      )
        _lowercase			 =torch.tensor(self.get_from_offsets(lowerCAmelCase							,     '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(lowerCAmelCase							,     lowerCAmelCase							,     atol=0.01      )      )
        self.assertTrue(torch.allclose(lowerCAmelCase							,     lowerCAmelCase							,     atol=0.01      )      )





print((lambda quine: quine % quine)('print((lambda quine: quine %% quine)(%r))'))


