






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			):







		'''simple docstring'''



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

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

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

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

							# load decoder from hub
							UpperCAmelCase_    :      List[Any]            =						'hf-internal-testing/ngram-beam-search-decoder'



		def 					__UpperCAmelCase			(							self				,		**_UpperCamelCase   )	->							Union[str, Any]:
							UpperCAmelCase_    :      Optional[Any]            =						self.add_kwargs_tokens_map.copy()
							kwargs.update(_UpperCamelCase   )
							return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname				,		**_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self				,		**_UpperCamelCase   )	->							List[Any]:
							return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname				,		**_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self				,		**_UpperCamelCase   )	->							int:
							return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name				,		**_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							str:
							shutil.rmtree(self.tmpdirname   )



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

							UpperCAmelCase_    :      List[str]            =						WavaVecaProcessorWithLM(tokenizer=_UpperCamelCase				,		feature_extractor=_UpperCamelCase				,		decoder=_UpperCamelCase   )

							processor.save_pretrained(self.tmpdirname   )
							UpperCAmelCase_    :      Any            =						WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname   )

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

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

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



		def 					__UpperCAmelCase			(							self   )	->							Any:
							UpperCAmelCase_    :      Optional[Any]            =						WavaVecaProcessorWithLM(
							    tokenizer=self.get_tokenizer()				,		feature_extractor=self.get_feature_extractor()				,		decoder=self.get_decoder()   )
							processor.save_pretrained(self.tmpdirname   )

							# make sure that error is thrown when decoder alphabet doesn't match
							UpperCAmelCase_    :      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   )	->							Dict:
							UpperCAmelCase_    :      List[Any]            =						self.get_tokenizer()
							# add token to trigger raise
							tokenizer.add_tokens(['xx']   )
							with self.assertRaisesRegex(_UpperCamelCase				,		'include'   ):
												WavaVecaProcessorWithLM(
												    tokenizer=_UpperCamelCase				,		feature_extractor=self.get_feature_extractor()				,		decoder=self.get_decoder()   )



		def 					__UpperCAmelCase			(							self   )	->							int:
							UpperCAmelCase_    :      Union[str, Any]            =						self.get_feature_extractor()
							UpperCAmelCase_    :      Any            =						self.get_tokenizer()
							UpperCAmelCase_    :      Optional[Any]            =						self.get_decoder()

							UpperCAmelCase_    :      List[Any]            =						WavaVecaProcessorWithLM(tokenizer=_UpperCamelCase				,		feature_extractor=_UpperCamelCase				,		decoder=_UpperCamelCase   )

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

							UpperCAmelCase_    :      List[str]            =						feature_extractor(_UpperCamelCase				,		return_tensors='np'   )
							UpperCAmelCase_    :      Any            =						processor(_UpperCamelCase				,		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   )	->							Any:
							UpperCAmelCase_    :      Optional[int]            =						self.get_feature_extractor()
							UpperCAmelCase_    :      Optional[int]            =						self.get_tokenizer()
							UpperCAmelCase_    :      List[Any]            =						self.get_decoder()

							UpperCAmelCase_    :      Optional[Any]            =						WavaVecaProcessorWithLM(tokenizer=_UpperCamelCase				,		feature_extractor=_UpperCamelCase				,		decoder=_UpperCamelCase   )

							UpperCAmelCase_    :      List[Any]            =						'This is a test string'

							UpperCAmelCase_    :      str            =						processor(text=_UpperCamelCase   )

							UpperCAmelCase_    :      List[str]            =						tokenizer(_UpperCamelCase   )

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



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase=(2, 1_0, 1_6)				,		_UpperCamelCase=7_7   )	->							Optional[Any]:
							np.random.seed(_UpperCamelCase   )
							return np.random.rand(*_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Dict:
							UpperCAmelCase_    :      int            =						self.get_feature_extractor()
							UpperCAmelCase_    :      int            =						self.get_tokenizer()
							UpperCAmelCase_    :      Optional[int]            =						self.get_decoder()

							UpperCAmelCase_    :      Union[str, Any]            =						WavaVecaProcessorWithLM(tokenizer=_UpperCamelCase				,		feature_extractor=_UpperCamelCase				,		decoder=_UpperCamelCase   )

							UpperCAmelCase_    :      int            =						self._get_dummy_logits(shape=(1_0, 1_6)				,		seed=1_3   )

							UpperCAmelCase_    :      Union[str, Any]            =						processor.decode(_UpperCamelCase   )

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

							UpperCAmelCase_    :      Dict            =						WavaVecaProcessorWithLM(tokenizer=_UpperCamelCase				,		feature_extractor=_UpperCamelCase				,		decoder=_UpperCamelCase   )

							UpperCAmelCase_    :      int            =						self._get_dummy_logits()

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

							UpperCAmelCase_    :      List[Any]            =						list(_UpperCamelCase   )

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

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



		def 					__UpperCAmelCase			(							self   )	->							int:
							UpperCAmelCase_    :      Optional[int]            =						self.get_feature_extractor()
							UpperCAmelCase_    :      Optional[Any]            =						self.get_tokenizer()
							UpperCAmelCase_    :      Any            =						self.get_decoder()

							UpperCAmelCase_    :      int            =						WavaVecaProcessorWithLM(tokenizer=_UpperCamelCase				,		feature_extractor=_UpperCamelCase				,		decoder=_UpperCamelCase   )

							UpperCAmelCase_    :      Any            =						self._get_dummy_logits()

							UpperCAmelCase_    :      List[str]            =						1_5
							UpperCAmelCase_    :      str            =						-20.0
							UpperCAmelCase_    :      List[Any]            =						-4.0

							UpperCAmelCase_    :      Dict            =						processor.batch_decode(
							    _UpperCamelCase				,		beam_width=_UpperCamelCase				,		beam_prune_logp=_UpperCamelCase				,		token_min_logp=_UpperCamelCase				,		)
							UpperCAmelCase_    :      List[str]            =						decoded_processor_out.text

							UpperCAmelCase_    :      Dict            =						list(_UpperCamelCase   )

							with get_context('fork'   ).Pool() as pool:
												UpperCAmelCase_    :      Optional[int]            =						decoder.decode_beams_batch(
												    _UpperCamelCase				,		_UpperCamelCase				,		beam_width=_UpperCamelCase				,		beam_prune_logp=_UpperCamelCase				,		token_min_logp=_UpperCamelCase				,		)

							UpperCAmelCase_    :      Dict            =						[d[0][0] for d in decoded_decoder_out]
							UpperCAmelCase_    :      Union[str, Any]            =						[d[0][2] for d in decoded_decoder_out]
							UpperCAmelCase_    :      str            =						[d[0][3] for d in decoded_decoder_out]

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

							self.assertTrue(np.array_equal(_UpperCamelCase				,		decoded_processor_out.logit_score   )   )
							self.assertTrue(np.allclose([-20.0_54, -18.4_47]				,		_UpperCamelCase				,		atol=1E-3   )   )

							self.assertTrue(np.array_equal(_UpperCamelCase				,		decoded_processor_out.lm_score   )   )
							self.assertTrue(np.allclose([-15.5_54, -13.94_74]				,		_UpperCamelCase				,		atol=1E-3   )   )



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

							UpperCAmelCase_    :      Tuple            =						WavaVecaProcessorWithLM(tokenizer=_UpperCamelCase				,		feature_extractor=_UpperCamelCase				,		decoder=_UpperCamelCase   )

							UpperCAmelCase_    :      Optional[Any]            =						self._get_dummy_logits()

							UpperCAmelCase_    :      List[Any]            =						2.0
							UpperCAmelCase_    :      List[str]            =						5.0
							UpperCAmelCase_    :      Optional[Any]            =						-20.0
							UpperCAmelCase_    :      Union[str, Any]            =						True

							UpperCAmelCase_    :      Any            =						processor.batch_decode(
							    _UpperCamelCase				,		alpha=_UpperCamelCase				,		beta=_UpperCamelCase				,		unk_score_offset=_UpperCamelCase				,		lm_score_boundary=_UpperCamelCase				,		)
							UpperCAmelCase_    :      Dict            =						decoded_processor_out.text

							UpperCAmelCase_    :      Any            =						list(_UpperCamelCase   )
							decoder.reset_params(
							    alpha=_UpperCamelCase				,		beta=_UpperCamelCase				,		unk_score_offset=_UpperCamelCase				,		lm_score_boundary=_UpperCamelCase				,		)

							with get_context('fork'   ).Pool() as pool:
												UpperCAmelCase_    :      Tuple            =						decoder.decode_beams_batch(
												    _UpperCamelCase				,		_UpperCamelCase				,		)

							UpperCAmelCase_    :      List[Any]            =						[d[0][0] for d in decoded_decoder_out]

							self.assertListEqual(_UpperCamelCase				,		_UpperCamelCase   )
							self.assertListEqual(['<s> </s> <s> </s> </s>', '</s> </s> <s> </s> </s>']				,		_UpperCamelCase   )
							UpperCAmelCase_    :      int            =						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				,		_UpperCamelCase   )



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

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

							UpperCAmelCase_    :      List[Any]            =						os.listdir(_UpperCamelCase   )
							UpperCAmelCase_    :      List[str]            =						['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(_UpperCamelCase				,		_UpperCamelCase   )



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

							UpperCAmelCase_    :      Union[str, Any]            =						WavaVecaProcessorWithLM.from_pretrained(_UpperCamelCase   )

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

							UpperCAmelCase_    :      int            =						os.listdir(_UpperCamelCase   )
							UpperCAmelCase_    :      Dict            =						os.listdir(_UpperCamelCase   )

							local_decoder_files.sort()
							expected_decoder_files.sort()

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



		def 					__UpperCAmelCase			(							self   )	->							Dict:
							UpperCAmelCase_    :      Any            =						WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'   )
							UpperCAmelCase_    :      Any            =						AutoProcessor.from_pretrained('hf-internal-testing/processor_with_lm'   )

							UpperCAmelCase_    :      Union[str, Any]            =						floats_list((3, 1_0_0_0)   )

							UpperCAmelCase_    :      Tuple            =						processor_wavaveca(_UpperCamelCase				,		return_tensors='np'   )
							UpperCAmelCase_    :      List[str]            =						processor_auto(_UpperCamelCase				,		return_tensors='np'   )

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

							UpperCAmelCase_    :      Optional[Any]            =						self._get_dummy_logits()

							UpperCAmelCase_    :      List[Any]            =						processor_wavaveca.batch_decode(_UpperCamelCase   )
							UpperCAmelCase_    :      List[str]            =						processor_auto.batch_decode(_UpperCamelCase   )

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



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

							UpperCAmelCase_    :      Dict            =						WavaVecaProcessorWithLM(tokenizer=_UpperCamelCase				,		feature_extractor=_UpperCamelCase				,		decoder=_UpperCamelCase   )

							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			(							_UpperCamelCase				,		_UpperCamelCase   )	->							str:
							UpperCAmelCase_    :      Optional[int]            =						[d[key] for d in offsets]
							return retrieved_list



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

							UpperCAmelCase_    :      Union[str, Any]            =						processor.decode(_UpperCamelCase				,		output_word_offsets=_UpperCamelCase   )
							# 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(_UpperCamelCase				,		_UpperCamelCase   )   )

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

							UpperCAmelCase_    :      Any            =						processor.batch_decode(_UpperCamelCase				,		output_word_offsets=_UpperCamelCase   )

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

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

							UpperCAmelCase_    :      Optional[Any]            =						load_dataset('common_voice'				,		'en'				,		split='train'				,		streaming=_UpperCamelCase   )
							UpperCAmelCase_    :      Dict            =						ds.cast_column('audio'				,		datasets.Audio(sampling_rate=1_6_0_0_0   )   )
							UpperCAmelCase_    :      Any            =						iter(_UpperCamelCase   )
							UpperCAmelCase_    :      int            =						next(_UpperCamelCase   )

							UpperCAmelCase_    :      List[str]            =						AutoProcessor.from_pretrained('patrickvonplaten/wav2vec2-base-100h-with-lm'   )
							UpperCAmelCase_    :      List[Any]            =						WavaVecaForCTC.from_pretrained('patrickvonplaten/wav2vec2-base-100h-with-lm'   )

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

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

							UpperCAmelCase_    :      List[Any]            =						processor.decode(logits[0]				,		output_word_offsets=_UpperCamelCase   )

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

							UpperCAmelCase_    :      List[str]            =						'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(_UpperCamelCase				,		'word'   )   )				,		_UpperCamelCase   )
							self.assertEqual(' '.join(self.get_from_offsets(_UpperCamelCase				,		'word'   )   )				,		output.text   )

							# output times
							UpperCAmelCase_    :      Tuple            =						torch.tensor(self.get_from_offsets(_UpperCamelCase				,		'start_time'   )   )
							UpperCAmelCase_    :      List[Any]            =						torch.tensor(self.get_from_offsets(_UpperCamelCase				,		'end_time'   )   )

							# fmt: off
							UpperCAmelCase_    :      Tuple            =						torch.tensor([1.41_99, 1.65_99, 2.25_99, 3.0, 3.24, 3.59_99, 3.79_99, 4.09_99, 4.26, 4.94, 5.28, 5.65_99, 5.78, 5.94, 6.32, 6.53_99, 6.65_99]   )
							UpperCAmelCase_    :      Optional[int]            =						torch.tensor([1.53_99, 1.89_99, 2.9, 3.16, 3.53_99, 3.72, 4.01_99, 4.17_99, 4.76, 5.15_99, 5.55_99, 5.69_99, 5.86, 6.19_99, 6.38, 6.61_99, 6.94]   )
							# fmt: on

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














from dataclasses import dataclass
from typing import Dict, Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .attention_processor import AttentionProcessor, AttnProcessor
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder



@dataclass
class   lowerCamelCase      (_snake_case			):







		'''simple docstring'''



		_snake_case    :    "DiagonalGaussianDistribution"





class   lowerCamelCase      (_snake_case ,  _snake_case			):







		'''simple docstring'''



		_snake_case    :    Optional[int]       						=  True



		@register_to_config
		def __init__(							self				,		_UpperCamelCase = 3				,		_UpperCamelCase = 3				,		_UpperCamelCase = ("DownEncoderBlock2D",)				,		_UpperCamelCase = ("UpDecoderBlock2D",)				,		_UpperCamelCase = (6_4,)				,		_UpperCamelCase = 1				,		_UpperCamelCase = "silu"				,		_UpperCamelCase = 4				,		_UpperCamelCase = 3_2				,		_UpperCamelCase = 3_2				,		_UpperCamelCase = 0.1_82_15				,		)	->							List[Any]:
							super().__init__()

							# pass init params to Encoder
							UpperCAmelCase_    :      List[str]            =						Encoder(
							    in_channels=_UpperCamelCase				,		out_channels=_UpperCamelCase				,		down_block_types=_UpperCamelCase				,		block_out_channels=_UpperCamelCase				,		layers_per_block=_UpperCamelCase				,		act_fn=_UpperCamelCase				,		norm_num_groups=_UpperCamelCase				,		double_z=_UpperCamelCase				,		)

							# pass init params to Decoder
							UpperCAmelCase_    :      Dict            =						Decoder(
							    in_channels=_UpperCamelCase				,		out_channels=_UpperCamelCase				,		up_block_types=_UpperCamelCase				,		block_out_channels=_UpperCamelCase				,		layers_per_block=_UpperCamelCase				,		norm_num_groups=_UpperCamelCase				,		act_fn=_UpperCamelCase				,		)

							UpperCAmelCase_    :      Any            =						nn.Convad(2 * latent_channels				,		2 * latent_channels				,		1   )
							UpperCAmelCase_    :      List[Any]            =						nn.Convad(_UpperCamelCase				,		_UpperCamelCase				,		1   )

							UpperCAmelCase_    :      Any            =						False
							UpperCAmelCase_    :      int            =						False

							# only relevant if vae tiling is enabled
							UpperCAmelCase_    :      Optional[int]            =						self.config.sample_size
							UpperCAmelCase_    :      int            =						(
							    self.config.sample_size[0]
							    if isinstance(self.config.sample_size				,		(list, tuple)   )
							    else self.config.sample_size
							)
							UpperCAmelCase_    :      Union[str, Any]            =						int(sample_size / (2 ** (len(self.config.block_out_channels   ) - 1))   )
							UpperCAmelCase_    :      Optional[Any]            =						0.25



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase=False   )	->							List[str]:
							if isinstance(_UpperCamelCase				,		(Encoder, Decoder)   ):
												UpperCAmelCase_    :      Union[str, Any]            =						value



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase = True   )	->							int:
							UpperCAmelCase_    :      Tuple            =						use_tiling



		def 					__UpperCAmelCase			(							self   )	->							Dict:
							self.enable_tiling(_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							UpperCAmelCase_    :      str            =						True



		def 					__UpperCAmelCase			(							self   )	->							List[Any]:
							UpperCAmelCase_    :      Optional[int]            =						False



		@property
		# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
		def 					__UpperCAmelCase			(							self   )	->							Dict[str, AttentionProcessor]:
							UpperCAmelCase_    :      Optional[int]            =						{}

							def fn_recursive_add_processors(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   ):
												if hasattr(_UpperCamelCase				,		'set_processor'   ):
																	UpperCAmelCase_    :      Optional[int]            =						module.processor

												for sub_name, child in module.named_children():
																	fn_recursive_add_processors(f"{name}.{sub_name}"				,		_UpperCamelCase				,		_UpperCamelCase   )

												return processors

							for name, module in self.named_children():
												fn_recursive_add_processors(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )

							return processors



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase   )	->							List[Any]:
							UpperCAmelCase_    :      Union[str, Any]            =						len(self.attn_processors.keys()   )

							if isinstance(_UpperCamelCase				,		_UpperCamelCase   ) and len(_UpperCamelCase   ) != count:
												raise ValueError(
												    f"A dict of processors was passed, but the number of processors {len(_UpperCamelCase   )} does not match the"
												    f" number of attention layers: {count}. Please make sure to pass {count} processor classes."   )

							def fn_recursive_attn_processor(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   ):
												if hasattr(_UpperCamelCase				,		'set_processor'   ):
																	if not isinstance(_UpperCamelCase				,		_UpperCamelCase   ):
																						module.set_processor(_UpperCamelCase   )
																	else:
																						module.set_processor(processor.pop(f"{name}.processor"   )   )

												for sub_name, child in module.named_children():
																	fn_recursive_attn_processor(f"{name}.{sub_name}"				,		_UpperCamelCase				,		_UpperCamelCase   )

							for name, module in self.named_children():
												fn_recursive_attn_processor(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Union[str, Any]:
							self.set_attn_processor(AttnProcessor()   )



		@apply_forward_hook
		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = True   )	->							AutoencoderKLOutput:
							if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
												return self.tiled_encode(_UpperCamelCase				,		return_dict=_UpperCamelCase   )

							if self.use_slicing and x.shape[0] > 1:
												UpperCAmelCase_    :      Union[str, Any]            =						[self.encoder(_UpperCamelCase   ) for x_slice in x.split(1   )]
												UpperCAmelCase_    :      Tuple            =						torch.cat(_UpperCamelCase   )
							else:
												UpperCAmelCase_    :      List[Any]            =						self.encoder(_UpperCamelCase   )

							UpperCAmelCase_    :      Optional[Any]            =						self.quant_conv(_UpperCamelCase   )
							UpperCAmelCase_    :      Tuple            =						DiagonalGaussianDistribution(_UpperCamelCase   )

							if not return_dict:
												return (posterior,)

							return AutoencoderKLOutput(latent_dist=_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = True   )	->							Union[DecoderOutput, torch.FloatTensor]:
							if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
												return self.tiled_decode(_UpperCamelCase				,		return_dict=_UpperCamelCase   )

							UpperCAmelCase_    :      str            =						self.post_quant_conv(_UpperCamelCase   )
							UpperCAmelCase_    :      List[str]            =						self.decoder(_UpperCamelCase   )

							if not return_dict:
												return (dec,)

							return DecoderOutput(sample=_UpperCamelCase   )



		@apply_forward_hook
		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = True   )	->							Union[DecoderOutput, torch.FloatTensor]:
							if self.use_slicing and z.shape[0] > 1:
												UpperCAmelCase_    :      List[str]            =						[self._decode(_UpperCamelCase   ).sample for z_slice in z.split(1   )]
												UpperCAmelCase_    :      Dict            =						torch.cat(_UpperCamelCase   )
							else:
												UpperCAmelCase_    :      Any            =						self._decode(_UpperCamelCase   ).sample

							if not return_dict:
												return (decoded,)

							return DecoderOutput(sample=_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Any:
							UpperCAmelCase_    :      Tuple            =						min(a.shape[2]				,		b.shape[2]				,		_UpperCamelCase   )
							for y in range(_UpperCamelCase   ):
												UpperCAmelCase_    :      str            =						a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
							return b



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Dict:
							UpperCAmelCase_    :      Tuple            =						min(a.shape[3]				,		b.shape[3]				,		_UpperCamelCase   )
							for x in range(_UpperCamelCase   ):
												UpperCAmelCase_    :      int            =						a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
							return b



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = True   )	->							AutoencoderKLOutput:
							UpperCAmelCase_    :      Any            =						int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)   )
							UpperCAmelCase_    :      Tuple            =						int(self.tile_latent_min_size * self.tile_overlap_factor   )
							UpperCAmelCase_    :      Optional[int]            =						self.tile_latent_min_size - blend_extent

							# Split the image into 512x512 tiles and encode them separately.
							UpperCAmelCase_    :      List[str]            =						[]
							for i in range(0				,		x.shape[2]				,		_UpperCamelCase   ):
												UpperCAmelCase_    :      Any            =						[]
												for j in range(0				,		x.shape[3]				,		_UpperCamelCase   ):
																	UpperCAmelCase_    :      Any            =						x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
																	UpperCAmelCase_    :      Dict            =						self.encoder(_UpperCamelCase   )
																	UpperCAmelCase_    :      List[str]            =						self.quant_conv(_UpperCamelCase   )
																	row.append(_UpperCamelCase   )
												rows.append(_UpperCamelCase   )
							UpperCAmelCase_    :      str            =						[]
							for i, row in enumerate(_UpperCamelCase   ):
												UpperCAmelCase_    :      List[Any]            =						[]
												for j, tile in enumerate(_UpperCamelCase   ):
																	# blend the above tile and the left tile
																	# to the current tile and add the current tile to the result row
																	if i > 0:
																						UpperCAmelCase_    :      Dict            =						self.blend_v(rows[i - 1][j]				,		_UpperCamelCase				,		_UpperCamelCase   )
																	if j > 0:
																						UpperCAmelCase_    :      List[str]            =						self.blend_h(row[j - 1]				,		_UpperCamelCase				,		_UpperCamelCase   )
																	result_row.append(tile[:, :, :row_limit, :row_limit]   )
												result_rows.append(torch.cat(_UpperCamelCase				,		dim=3   )   )

							UpperCAmelCase_    :      Union[str, Any]            =						torch.cat(_UpperCamelCase				,		dim=2   )
							UpperCAmelCase_    :      List[Any]            =						DiagonalGaussianDistribution(_UpperCamelCase   )

							if not return_dict:
												return (posterior,)

							return AutoencoderKLOutput(latent_dist=_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = True   )	->							Union[DecoderOutput, torch.FloatTensor]:
							UpperCAmelCase_    :      str            =						int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)   )
							UpperCAmelCase_    :      Dict            =						int(self.tile_sample_min_size * self.tile_overlap_factor   )
							UpperCAmelCase_    :      Dict            =						self.tile_sample_min_size - blend_extent

							# Split z into overlapping 64x64 tiles and decode them separately.
							# The tiles have an overlap to avoid seams between tiles.
							UpperCAmelCase_    :      Union[str, Any]            =						[]
							for i in range(0				,		z.shape[2]				,		_UpperCamelCase   ):
												UpperCAmelCase_    :      List[str]            =						[]
												for j in range(0				,		z.shape[3]				,		_UpperCamelCase   ):
																	UpperCAmelCase_    :      List[str]            =						z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
																	UpperCAmelCase_    :      Optional[Any]            =						self.post_quant_conv(_UpperCamelCase   )
																	UpperCAmelCase_    :      Tuple            =						self.decoder(_UpperCamelCase   )
																	row.append(_UpperCamelCase   )
												rows.append(_UpperCamelCase   )
							UpperCAmelCase_    :      Optional[Any]            =						[]
							for i, row in enumerate(_UpperCamelCase   ):
												UpperCAmelCase_    :      List[Any]            =						[]
												for j, tile in enumerate(_UpperCamelCase   ):
																	# blend the above tile and the left tile
																	# to the current tile and add the current tile to the result row
																	if i > 0:
																						UpperCAmelCase_    :      Union[str, Any]            =						self.blend_v(rows[i - 1][j]				,		_UpperCamelCase				,		_UpperCamelCase   )
																	if j > 0:
																						UpperCAmelCase_    :      Optional[Any]            =						self.blend_h(row[j - 1]				,		_UpperCamelCase				,		_UpperCamelCase   )
																	result_row.append(tile[:, :, :row_limit, :row_limit]   )
												result_rows.append(torch.cat(_UpperCamelCase				,		dim=3   )   )

							UpperCAmelCase_    :      Dict            =						torch.cat(_UpperCamelCase				,		dim=2   )
							if not return_dict:
												return (dec,)

							return DecoderOutput(sample=_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = False				,		_UpperCamelCase = True				,		_UpperCamelCase = None				,		)	->							Union[DecoderOutput, torch.FloatTensor]:
							UpperCAmelCase_    :      Optional[Any]            =						sample
							UpperCAmelCase_    :      Union[str, Any]            =						self.encode(_UpperCamelCase   ).latent_dist
							if sample_posterior:
												UpperCAmelCase_    :      str            =						posterior.sample(generator=_UpperCamelCase   )
							else:
												UpperCAmelCase_    :      int            =						posterior.mode()
							UpperCAmelCase_    :      Dict            =						self.decode(_UpperCamelCase   ).sample

							if not return_dict:
												return (dec,)

							return DecoderOutput(sample=_UpperCamelCase   )






