


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




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

              snake_case__ :       str	   =							{
                  'unk_token': '<unk>',
                  'bos_token': '<s>',
                  'eos_token': '</s>',
              }
              snake_case__ :       List[Any]	   =							{
                  'feature_size': 1,
                  'padding_value': 0.0,
                  'sampling_rate': 1_6_0_0_0,
                  'return_attention_mask': False,
                  'do_normalize': True,
              }

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




       def     UpperCAmelCase__    ( self				:    List[Any]    ,				**_lowerCamelCase				:    Tuple			):
              snake_case__ :       Dict	   =							self.add_kwargs_tokens_map.copy()
              kwargs.update(SCREAMING_SNAKE_CASE_			)
              return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname    ,				**SCREAMING_SNAKE_CASE_			)




       def     UpperCAmelCase__    ( self				:    Tuple    ,				**_lowerCamelCase				:    Tuple			):
              return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname    ,				**SCREAMING_SNAKE_CASE_			)




       def     UpperCAmelCase__    ( self				:    Union[str, Any]    ,				**_lowerCamelCase				:    Any			):
              return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name    ,				**SCREAMING_SNAKE_CASE_			)




       def     UpperCAmelCase__    ( self				:    str			):
              shutil.rmtree(self.tmpdirname			)




       def     UpperCAmelCase__    ( self				:    Tuple			):
              snake_case__ :       str	   =							self.get_tokenizer()
              snake_case__ :       Dict	   =							self.get_feature_extractor()
              snake_case__ :       List[Any]	   =							self.get_decoder()

              snake_case__ :       List[Any]	   =							WavaVecaProcessorWithLM(tokenizer=SCREAMING_SNAKE_CASE_    ,				feature_extractor=SCREAMING_SNAKE_CASE_    ,				decoder=SCREAMING_SNAKE_CASE_			)

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

              snake_case__ :       Optional[int]	   =							WavaVecaProcessorWithLM(tokenizer=SCREAMING_SNAKE_CASE_    ,				feature_extractor=SCREAMING_SNAKE_CASE_    ,				decoder=SCREAMING_SNAKE_CASE_			)

              snake_case__ :       Dict	   =							floats_list((3, 1_0_0_0)			)

              snake_case__ :       Union[str, Any]	   =							feature_extractor(SCREAMING_SNAKE_CASE_    ,				return_tensors='np'			)
              snake_case__ :       Union[str, Any]	   =							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     UpperCAmelCase__    ( self				:    str			):
              snake_case__ :       Union[str, Any]	   =							self.get_feature_extractor()
              snake_case__ :       Optional[int]	   =							self.get_tokenizer()
              snake_case__ :       Optional[Any]	   =							self.get_decoder()

              snake_case__ :       str	   =							WavaVecaProcessorWithLM(tokenizer=SCREAMING_SNAKE_CASE_    ,				feature_extractor=SCREAMING_SNAKE_CASE_    ,				decoder=SCREAMING_SNAKE_CASE_			)

              snake_case__ :       int	   =							'This is a test string'

              snake_case__ :       Any	   =							processor(text=SCREAMING_SNAKE_CASE_			)

              snake_case__ :       List[Any]	   =							tokenizer(SCREAMING_SNAKE_CASE_			)

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




       def     UpperCAmelCase__    ( self				:    Union[str, Any]    ,				_lowerCamelCase				:    Union[str, Any]=(2, 1_0, 1_6)    ,				_lowerCamelCase				:    Any=7_7			):
              np.random.seed(SCREAMING_SNAKE_CASE_			)
              return np.random.rand(*SCREAMING_SNAKE_CASE_			)




       def     UpperCAmelCase__    ( self				:    Any			):
              snake_case__ :       str	   =							self.get_feature_extractor()
              snake_case__ :       Union[str, Any]	   =							self.get_tokenizer()
              snake_case__ :       Optional[Any]	   =							self.get_decoder()

              snake_case__ :       Union[str, Any]	   =							WavaVecaProcessorWithLM(tokenizer=SCREAMING_SNAKE_CASE_    ,				feature_extractor=SCREAMING_SNAKE_CASE_    ,				decoder=SCREAMING_SNAKE_CASE_			)

              snake_case__ :       Optional[int]	   =							self._get_dummy_logits(shape=(1_0, 1_6)    ,				seed=1_3			)

              snake_case__ :       Optional[int]	   =							processor.decode(SCREAMING_SNAKE_CASE_			)

              snake_case__ :       List[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     UpperCAmelCase__    ( self				:    Dict    ,				_lowerCamelCase				:    Union[str, Any]			):
              snake_case__ :       Optional[int]	   =							self.get_feature_extractor()
              snake_case__ :       Any	   =							self.get_tokenizer()
              snake_case__ :       List[Any]	   =							self.get_decoder()

              snake_case__ :       str	   =							WavaVecaProcessorWithLM(tokenizer=SCREAMING_SNAKE_CASE_    ,				feature_extractor=SCREAMING_SNAKE_CASE_    ,				decoder=SCREAMING_SNAKE_CASE_			)

              snake_case__ :       Optional[Any]	   =							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:
                     snake_case__ :       List[str]	   =							processor.batch_decode(SCREAMING_SNAKE_CASE_			)
              else:
                     with get_context(SCREAMING_SNAKE_CASE_			).Pool() as pool:
                            snake_case__ :       str	   =							processor.batch_decode(SCREAMING_SNAKE_CASE_    ,				SCREAMING_SNAKE_CASE_			)

              snake_case__ :       Optional[int]	   =							list(SCREAMING_SNAKE_CASE_			)

              with get_context('fork'			).Pool() as p:
                     snake_case__ :       Dict	   =							decoder.decode_beams_batch(SCREAMING_SNAKE_CASE_    ,				SCREAMING_SNAKE_CASE_			)

              snake_case__ :       List[Any]	   =							[], [], []
              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     UpperCAmelCase__    ( self				:    int			):
              snake_case__ :       str	   =							self.get_feature_extractor()
              snake_case__ :       str	   =							self.get_tokenizer()
              snake_case__ :       Any	   =							self.get_decoder()

              snake_case__ :       str	   =							WavaVecaProcessorWithLM(tokenizer=SCREAMING_SNAKE_CASE_    ,				feature_extractor=SCREAMING_SNAKE_CASE_    ,				decoder=SCREAMING_SNAKE_CASE_			)

              snake_case__ :       Any	   =							self._get_dummy_logits()

              snake_case__ :       int	   =							1_5
              snake_case__ :       Dict	   =							-2_0.0
              snake_case__ :       Optional[Any]	   =							-4.0

              snake_case__ :       List[Any]	   =							processor.batch_decode(
                  SCREAMING_SNAKE_CASE_    ,				beam_width=SCREAMING_SNAKE_CASE_    ,				beam_prune_logp=SCREAMING_SNAKE_CASE_    ,				token_min_logp=SCREAMING_SNAKE_CASE_    ,				)
              snake_case__ :       List[Any]	   =							decoded_processor_out.text

              snake_case__ :       Optional[int]	   =							list(SCREAMING_SNAKE_CASE_			)

              with get_context('fork'			).Pool() as pool:
                     snake_case__ :       List[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_    ,				)

              snake_case__ :       Union[str, Any]	   =							[d[0][0] for d in decoded_decoder_out]
              snake_case__ :       Tuple	   =							[d[0][2] for d in decoded_decoder_out]
              snake_case__ :       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([-2_0.0_5_4, -1_8.4_4_7]    ,				SCREAMING_SNAKE_CASE_    ,				atol=1E-3			)			)

              self.assertTrue(np.array_equal(SCREAMING_SNAKE_CASE_    ,				decoded_processor_out.lm_score			)			)
              self.assertTrue(np.allclose([-1_5.5_5_4, -1_3.9_4_7_4]    ,				SCREAMING_SNAKE_CASE_    ,				atol=1E-3			)			)




       def     UpperCAmelCase__    ( self				:    Optional[Any]			):
              snake_case__ :       List[Any]	   =							self.get_feature_extractor()
              snake_case__ :       int	   =							self.get_tokenizer()
              snake_case__ :       int	   =							self.get_decoder()

              snake_case__ :       Dict	   =							WavaVecaProcessorWithLM(tokenizer=SCREAMING_SNAKE_CASE_    ,				feature_extractor=SCREAMING_SNAKE_CASE_    ,				decoder=SCREAMING_SNAKE_CASE_			)

              snake_case__ :       Dict	   =							self._get_dummy_logits()

              snake_case__ :       Union[str, Any]	   =							2.0
              snake_case__ :       str	   =							5.0
              snake_case__ :       int	   =							-2_0.0
              snake_case__ :       Optional[Any]	   =							True

              snake_case__ :       str	   =							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_    ,				)
              snake_case__ :       Any	   =							decoded_processor_out.text

              snake_case__ :       int	   =							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:
                     snake_case__ :       Union[str, Any]	   =							decoder.decode_beams_batch(
                         SCREAMING_SNAKE_CASE_    ,				SCREAMING_SNAKE_CASE_    ,				)

              snake_case__ :       Tuple	   =							[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_			)
              snake_case__ :       Union[str, Any]	   =							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    ,				-2_0.0			)
              self.assertEqual(lm_model.score_boundary    ,				SCREAMING_SNAKE_CASE_			)




       def     UpperCAmelCase__    ( self				:    Any			):
              snake_case__ :       Any	   =							WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'			)

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

              snake_case__ :       Optional[int]	   =							os.listdir(SCREAMING_SNAKE_CASE_			)
              snake_case__ :       Union[str, 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     UpperCAmelCase__    ( self				:    List[Any]			):
              snake_case__ :       Union[str, Any]	   =							snapshot_download('hf-internal-testing/processor_with_lm'			)

              snake_case__ :       Optional[Any]	   =							WavaVecaProcessorWithLM.from_pretrained(SCREAMING_SNAKE_CASE_			)

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

              snake_case__ :       str	   =							os.listdir(SCREAMING_SNAKE_CASE_			)
              snake_case__ :       Dict	   =							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     UpperCAmelCase__    ( self				:    Optional[int]			):
              snake_case__ :       Optional[int]	   =							WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'			)
              snake_case__ :       Dict	   =							AutoProcessor.from_pretrained('hf-internal-testing/processor_with_lm'			)

              snake_case__ :       Tuple	   =							floats_list((3, 1_0_0_0)			)

              snake_case__ :       Optional[Any]	   =							processor_wavaveca(SCREAMING_SNAKE_CASE_    ,				return_tensors='np'			)
              snake_case__ :       Optional[int]	   =							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			)

              snake_case__ :       int	   =							self._get_dummy_logits()

              snake_case__ :       Tuple	   =							processor_wavaveca.batch_decode(SCREAMING_SNAKE_CASE_			)
              snake_case__ :       Optional[Any]	   =							processor_auto.batch_decode(SCREAMING_SNAKE_CASE_			)

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




       def     UpperCAmelCase__    ( self				:    Tuple			):
              snake_case__ :       Optional[int]	   =							self.get_feature_extractor()
              snake_case__ :       List[str]	   =							self.get_tokenizer()
              snake_case__ :       List[Any]	   =							self.get_decoder()

              snake_case__ :       Union[str, Any]	   =							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     UpperCAmelCase__    ( _lowerCamelCase				:    List[str]    ,				_lowerCamelCase				:    List[str]			):
              snake_case__ :       int	   =							[d[key] for d in offsets]
              return retrieved_list




       def     UpperCAmelCase__    ( self				:    int			):
              snake_case__ :       str	   =							WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'			)
              snake_case__ :       Optional[Any]	   =							self._get_dummy_logits()[0]

              snake_case__ :       Dict	   =							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     UpperCAmelCase__    ( self				:    Optional[int]			):
              snake_case__ :       Tuple	   =							WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'			)
              snake_case__ :       Union[str, Any]	   =							self._get_dummy_logits()

              snake_case__ :       Tuple	   =							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     UpperCAmelCase__    ( self				:    Optional[Any]			):
              import torch

              snake_case__ :       Optional[Any]	   =							load_dataset('common_voice'    ,				'en'    ,				split='train'    ,				streaming=SCREAMING_SNAKE_CASE_			)
              snake_case__ :       Union[str, Any]	   =							ds.cast_column('audio'    ,				datasets.Audio(sampling_rate=1_6_0_0_0			)			)
              snake_case__ :       List[str]	   =							iter(SCREAMING_SNAKE_CASE_			)
              snake_case__ :       List[Any]	   =							next(SCREAMING_SNAKE_CASE_			)

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

              with torch.no_grad():
                     snake_case__ :       Optional[int]	   =							model(SCREAMING_SNAKE_CASE_			).logits.cpu().numpy()

              snake_case__ :       Optional[int]	   =							processor.decode(logits[0]    ,				output_word_offsets=SCREAMING_SNAKE_CASE_			)

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

              snake_case__ :       Optional[Any]	   =							'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
              snake_case__ :       str	   =							torch.tensor(self.get_from_offsets(SCREAMING_SNAKE_CASE_    ,				'start_time'			)			)
              snake_case__ :       Tuple	   =							torch.tensor(self.get_from_offsets(SCREAMING_SNAKE_CASE_    ,				'end_time'			)			)

              # fmt: off
              snake_case__ :       Any	   =							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]			)
              snake_case__ :       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 flax.linen as nn
import jax
import jax.numpy as jnp





class   UpperCAmelCase_    (nn.Module    ):





    """simple docstring"""






    lowerCamelCase    :   int
    lowerCamelCase    :   jnp.dtype					=     jnp.floataa







    def  lowercase_							(				self      )							->	Union[str, Any]:
         __lowerCamelCase  :      Optional[Any]														=						nn.Conv(
             self.out_channels     ,	kernel_size=(3, 3)     ,	strides=(1, 1)     ,	padding=((1, 1), (1, 1))     ,	dtype=self.dtype     ,	)







    def __call__(				self     ,	SCREAMING_SNAKE_CASE_      )							->	Any:
         __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase  :      Optional[int]														=						hidden_states.shape
         __lowerCamelCase  :      Dict														=						jax.image.resize(
             SCREAMING_SNAKE_CASE_     ,	shape=(batch, height * 2, width * 2, channels)     ,	method='nearest'     ,	)
         __lowerCamelCase  :      Optional[Any]														=						self.conv(SCREAMING_SNAKE_CASE_      )
         return hidden_states







class   UpperCAmelCase_    (nn.Module    ):





    """simple docstring"""






    lowerCamelCase    :   int
    lowerCamelCase    :   jnp.dtype					=     jnp.floataa







    def  lowercase_							(				self      )							->	List[str]:
         __lowerCamelCase  :      str														=						nn.Conv(
             self.out_channels     ,	kernel_size=(3, 3)     ,	strides=(2, 2)     ,	padding=((1, 1), (1, 1))     ,	dtype=self.dtype     ,	)







    def __call__(				self     ,	SCREAMING_SNAKE_CASE_      )							->	List[str]:
         # pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
         # hidden_states = jnp.pad(hidden_states, pad_width=pad)
         __lowerCamelCase  :      str														=						self.conv(SCREAMING_SNAKE_CASE_      )
         return hidden_states







class   UpperCAmelCase_    (nn.Module    ):





    """simple docstring"""






    lowerCamelCase    :   int
    lowerCamelCase    :   int					=     None
    lowerCamelCase    :   float					=     0.0
    lowerCamelCase    :   bool					=     None
    lowerCamelCase    :   jnp.dtype					=     jnp.floataa







    def  lowercase_							(				self      )							->	Optional[int]:
         __lowerCamelCase  :      Optional[Any]														=						self.in_channels if self.out_channels is None else self.out_channels

         __lowerCamelCase  :      Optional[Any]														=						nn.GroupNorm(num_groups=32     ,	epsilon=1E-5      )
         __lowerCamelCase  :      Tuple														=						nn.Conv(
             SCREAMING_SNAKE_CASE_     ,	kernel_size=(3, 3)     ,	strides=(1, 1)     ,	padding=((1, 1), (1, 1))     ,	dtype=self.dtype     ,	)

         __lowerCamelCase  :      List[str]														=						nn.Dense(SCREAMING_SNAKE_CASE_     ,	dtype=self.dtype      )

         __lowerCamelCase  :      Dict														=						nn.GroupNorm(num_groups=32     ,	epsilon=1E-5      )
         __lowerCamelCase  :      int														=						nn.Dropout(self.dropout_prob      )
         __lowerCamelCase  :      Union[str, Any]														=						nn.Conv(
             SCREAMING_SNAKE_CASE_     ,	kernel_size=(3, 3)     ,	strides=(1, 1)     ,	padding=((1, 1), (1, 1))     ,	dtype=self.dtype     ,	)

         __lowerCamelCase  :      Optional[int]														=						self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

         __lowerCamelCase  :      List[Any]														=						None
         if use_nin_shortcut:
              __lowerCamelCase  :      Any														=						nn.Conv(
                  SCREAMING_SNAKE_CASE_     ,	kernel_size=(1, 1)     ,	strides=(1, 1)     ,	padding='VALID'     ,	dtype=self.dtype     ,	)







    def __call__(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_=True      )							->	Tuple:
         __lowerCamelCase  :      List[Any]														=						hidden_states
         __lowerCamelCase  :      str														=						self.norma(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      Union[str, Any]														=						nn.swish(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      Union[str, Any]														=						self.conva(SCREAMING_SNAKE_CASE_      )

         __lowerCamelCase  :      str														=						self.time_emb_proj(nn.swish(SCREAMING_SNAKE_CASE_      )      )
         __lowerCamelCase  :      List[str]														=						jnp.expand_dims(jnp.expand_dims(SCREAMING_SNAKE_CASE_     ,	1      )     ,	1      )
         __lowerCamelCase  :      Optional[int]														=						hidden_states + temb

         __lowerCamelCase  :      List[Any]														=						self.norma(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      Any														=						nn.swish(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      List[Any]														=						self.dropout(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      List[Any]														=						self.conva(SCREAMING_SNAKE_CASE_      )

         if self.conv_shortcut is not None:
              __lowerCamelCase  :      List[str]														=						self.conv_shortcut(SCREAMING_SNAKE_CASE_      )

         return hidden_states + residual






