

"""simple docstring"""




import json
import os
import shutil
import tempfile
import unittest
from multiprocessing import get_context
from pathlib import Path

import datasets
import numpy as np
from datasets import load_dataset
from parameterized import parameterized

from transformers import AutoProcessor
from transformers.models.wavaveca import WavaVecaCTCTokenizer, WavaVecaFeatureExtractor
from transformers.models.wavaveca.tokenization_wavaveca import VOCAB_FILES_NAMES
from transformers.testing_utils import require_pyctcdecode, require_torch, require_torchaudio, slow
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_pyctcdecode_available, is_torch_available

from ..wavaveca.test_feature_extraction_wavaveca import floats_list


if is_pyctcdecode_available():
						from huggingface_hub import snapshot_download
						from pyctcdecode import BeamSearchDecoderCTC

						from transformers.models.wavaveca_with_lm import WavaVecaProcessorWithLM
						from transformers.models.wavaveca_with_lm.processing_wavaveca_with_lm import WavaVecaDecoderWithLMOutput

if is_torch_available():
						from transformers import WavaVecaForCTC





@require_pyctcdecode
class    lowerCamelCase						(unittest.TestCase			):

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

						SCREAMING_SNAKE_CASE__								=      {
						    """unk_token""": """<unk>""",
						    """bos_token""": """<s>""",
						    """eos_token""": """</s>""",
						}
						SCREAMING_SNAKE_CASE__								=      {
						    """feature_size""": 1,
						    """padding_value""": 0.0,
						    """sampling_rate""": 1_6_0_0_0,
						    """return_attention_mask""": False,
						    """do_normalize""": True,
						}

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

			def 					SCREAMING_SNAKE_CASE				(					self				:						List[Any]    ,				**__UpperCAmelCase				:						Optional[int]       )   ->    Tuple:
						SCREAMING_SNAKE_CASE__								=      self.add_kwargs_tokens_map.copy()
						kwargs.update(__UpperCAmelCase       )
						return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname    ,				**__UpperCAmelCase       )

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

			def 					SCREAMING_SNAKE_CASE				(					self				:						str    ,				**__UpperCAmelCase				:						List[Any]       )   ->    List[str]:
						return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name    ,				**__UpperCAmelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						int       )   ->    Optional[int]:
						shutil.rmtree(self.tmpdirname       )

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

						SCREAMING_SNAKE_CASE__								=      WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase    ,				feature_extractor=__UpperCAmelCase    ,				decoder=__UpperCAmelCase       )

						processor.save_pretrained(self.tmpdirname       )
						SCREAMING_SNAKE_CASE__								=      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 					SCREAMING_SNAKE_CASE				(					self				:						Optional[Any]       )   ->    Optional[Any]:
						SCREAMING_SNAKE_CASE__								=      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
						SCREAMING_SNAKE_CASE__								=      WavaVecaProcessorWithLM.from_pretrained(
						    self.tmpdirname    ,				alpha=5.0    ,				beta=3.0    ,				score_boundary=-7.0    ,				unk_score_offset=3       )

						# decoder
						self.assertEqual(processor.language_model.alpha    ,				5.0       )
						self.assertEqual(processor.language_model.beta    ,				3.0       )
						self.assertEqual(processor.language_model.score_boundary    ,				-7.0       )
						self.assertEqual(processor.language_model.unk_score_offset    ,				3       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						int       )   ->    Any:
						SCREAMING_SNAKE_CASE__								=      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 					SCREAMING_SNAKE_CASE				(					self				:						List[str]       )   ->    str:
						SCREAMING_SNAKE_CASE__								=      self.get_feature_extractor()
						SCREAMING_SNAKE_CASE__								=      self.get_tokenizer()
						SCREAMING_SNAKE_CASE__								=      self.get_decoder()

						SCREAMING_SNAKE_CASE__								=      WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase    ,				feature_extractor=__UpperCAmelCase    ,				decoder=__UpperCAmelCase       )

						SCREAMING_SNAKE_CASE__								=      floats_list((3, 1_0_0_0)       )

						SCREAMING_SNAKE_CASE__								=      feature_extractor(__UpperCAmelCase    ,				return_tensors="""np"""       )
						SCREAMING_SNAKE_CASE__								=      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 					SCREAMING_SNAKE_CASE				(					self				:						Any       )   ->    List[Any]:
						SCREAMING_SNAKE_CASE__								=      self.get_feature_extractor()
						SCREAMING_SNAKE_CASE__								=      self.get_tokenizer()
						SCREAMING_SNAKE_CASE__								=      self.get_decoder()

						SCREAMING_SNAKE_CASE__								=      WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase    ,				feature_extractor=__UpperCAmelCase    ,				decoder=__UpperCAmelCase       )

						SCREAMING_SNAKE_CASE__								=      """This is a test string"""

						SCREAMING_SNAKE_CASE__								=      processor(text=__UpperCAmelCase       )

						SCREAMING_SNAKE_CASE__								=      tokenizer(__UpperCAmelCase       )

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

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

			def 					SCREAMING_SNAKE_CASE				(					self				:						Optional[Any]       )   ->    Dict:
						SCREAMING_SNAKE_CASE__								=      self.get_feature_extractor()
						SCREAMING_SNAKE_CASE__								=      self.get_tokenizer()
						SCREAMING_SNAKE_CASE__								=      self.get_decoder()

						SCREAMING_SNAKE_CASE__								=      WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase    ,				feature_extractor=__UpperCAmelCase    ,				decoder=__UpperCAmelCase       )

						SCREAMING_SNAKE_CASE__								=      self._get_dummy_logits(shape=(1_0, 1_6)    ,				seed=1_3       )

						SCREAMING_SNAKE_CASE__								=      processor.decode(__UpperCAmelCase       )

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

						SCREAMING_SNAKE_CASE__								=      WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase    ,				feature_extractor=__UpperCAmelCase    ,				decoder=__UpperCAmelCase       )

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

						SCREAMING_SNAKE_CASE__								=      list(__UpperCAmelCase       )

						with get_context("""fork"""       ).Pool() as p:
									SCREAMING_SNAKE_CASE__								=      decoder.decode_beams_batch(__UpperCAmelCase    ,				__UpperCAmelCase       )

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

						SCREAMING_SNAKE_CASE__								=      WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase    ,				feature_extractor=__UpperCAmelCase    ,				decoder=__UpperCAmelCase       )

						SCREAMING_SNAKE_CASE__								=      self._get_dummy_logits()

						SCREAMING_SNAKE_CASE__								=      1_5
						SCREAMING_SNAKE_CASE__								=      -20.0
						SCREAMING_SNAKE_CASE__								=      -4.0

						SCREAMING_SNAKE_CASE__								=      processor.batch_decode(
						    __UpperCAmelCase    ,				beam_width=__UpperCAmelCase    ,				beam_prune_logp=__UpperCAmelCase    ,				token_min_logp=__UpperCAmelCase    ,				)
						SCREAMING_SNAKE_CASE__								=      decoded_processor_out.text

						SCREAMING_SNAKE_CASE__								=      list(__UpperCAmelCase       )

						with get_context("""fork"""       ).Pool() as pool:
									SCREAMING_SNAKE_CASE__								=      decoder.decode_beams_batch(
									    __UpperCAmelCase    ,				__UpperCAmelCase    ,				beam_width=__UpperCAmelCase    ,				beam_prune_logp=__UpperCAmelCase    ,				token_min_logp=__UpperCAmelCase    ,				)

						SCREAMING_SNAKE_CASE__								=      [d[0][0] for d in decoded_decoder_out]
						SCREAMING_SNAKE_CASE__								=      [d[0][2] for d in decoded_decoder_out]
						SCREAMING_SNAKE_CASE__								=      [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.054, -18.447]    ,				__UpperCAmelCase    ,				atol=1e-3       )       )

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

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

						SCREAMING_SNAKE_CASE__								=      WavaVecaProcessorWithLM(tokenizer=__UpperCAmelCase    ,				feature_extractor=__UpperCAmelCase    ,				decoder=__UpperCAmelCase       )

						SCREAMING_SNAKE_CASE__								=      self._get_dummy_logits()

						SCREAMING_SNAKE_CASE__								=      2.0
						SCREAMING_SNAKE_CASE__								=      5.0
						SCREAMING_SNAKE_CASE__								=      -20.0
						SCREAMING_SNAKE_CASE__								=      True

						SCREAMING_SNAKE_CASE__								=      processor.batch_decode(
						    __UpperCAmelCase    ,				alpha=__UpperCAmelCase    ,				beta=__UpperCAmelCase    ,				unk_score_offset=__UpperCAmelCase    ,				lm_score_boundary=__UpperCAmelCase    ,				)
						SCREAMING_SNAKE_CASE__								=      decoded_processor_out.text

						SCREAMING_SNAKE_CASE__								=      list(__UpperCAmelCase       )
						decoder.reset_params(
						    alpha=__UpperCAmelCase    ,				beta=__UpperCAmelCase    ,				unk_score_offset=__UpperCAmelCase    ,				lm_score_boundary=__UpperCAmelCase    ,				)

						with get_context("""fork"""       ).Pool() as pool:
									SCREAMING_SNAKE_CASE__								=      decoder.decode_beams_batch(
									    __UpperCAmelCase    ,				__UpperCAmelCase    ,				)

						SCREAMING_SNAKE_CASE__								=      [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       )
						SCREAMING_SNAKE_CASE__								=      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 					SCREAMING_SNAKE_CASE				(					self				:						List[Any]       )   ->    int:
						SCREAMING_SNAKE_CASE__								=      WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""       )

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

						SCREAMING_SNAKE_CASE__								=      os.listdir(__UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      ["""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 					SCREAMING_SNAKE_CASE				(					self				:						List[str]       )   ->    Optional[Any]:
						SCREAMING_SNAKE_CASE__								=      snapshot_download("""hf-internal-testing/processor_with_lm"""       )

						SCREAMING_SNAKE_CASE__								=      WavaVecaProcessorWithLM.from_pretrained(__UpperCAmelCase       )

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

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

						SCREAMING_SNAKE_CASE__								=      floats_list((3, 1_0_0_0)       )

						SCREAMING_SNAKE_CASE__								=      processor_wavaveca(__UpperCAmelCase    ,				return_tensors="""np"""       )
						SCREAMING_SNAKE_CASE__								=      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       )

						SCREAMING_SNAKE_CASE__								=      self._get_dummy_logits()

						SCREAMING_SNAKE_CASE__								=      processor_wavaveca.batch_decode(__UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      processor_auto.batch_decode(__UpperCAmelCase       )

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

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

						SCREAMING_SNAKE_CASE__								=      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 					SCREAMING_SNAKE_CASE				(					__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						Dict       )   ->    Optional[int]:
						SCREAMING_SNAKE_CASE__								=      [d[key] for d in offsets]
						return retrieved_list

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

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

						SCREAMING_SNAKE_CASE__								=      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 					SCREAMING_SNAKE_CASE				(					self				:						List[Any]       )   ->    Optional[Any]:
						import torch

						SCREAMING_SNAKE_CASE__								=      load_dataset("""common_voice"""    ,				"""en"""    ,				split="""train"""    ,				streaming=__UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      ds.cast_column("""audio"""    ,				datasets.Audio(sampling_rate=1_6_0_0_0       )       )
						SCREAMING_SNAKE_CASE__								=      iter(__UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      next(__UpperCAmelCase       )

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

						with torch.no_grad():
									SCREAMING_SNAKE_CASE__								=      model(__UpperCAmelCase       ).logits.cpu().numpy()

						SCREAMING_SNAKE_CASE__								=      processor.decode(logits[0]    ,				output_word_offsets=__UpperCAmelCase       )

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

						SCREAMING_SNAKE_CASE__								=      """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
						SCREAMING_SNAKE_CASE__								=      torch.tensor(self.get_from_offsets(__UpperCAmelCase    ,				"""start_time"""       )       )
						SCREAMING_SNAKE_CASE__								=      torch.tensor(self.get_from_offsets(__UpperCAmelCase    ,				"""end_time"""       )       )

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

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





"""simple docstring"""




import unittest

import numpy as np
import torch

from diffusers import KarrasVePipeline, KarrasVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()





class    lowerCamelCase						(unittest.TestCase			):

			@property
			def 					SCREAMING_SNAKE_CASE				(					self				:						Dict       )   ->    Any:
						torch.manual_seed(0       )
						SCREAMING_SNAKE_CASE__								=      UNetaDModel(
						    block_out_channels=(3_2, 6_4)    ,				layers_per_block=2    ,				sample_size=3_2    ,				in_channels=3    ,				out_channels=3    ,				down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")    ,				up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")    ,				)
						return model





			def 					SCREAMING_SNAKE_CASE				(					self				:						Any       )   ->    Tuple:
						SCREAMING_SNAKE_CASE__								=      self.dummy_uncond_unet
						SCREAMING_SNAKE_CASE__								=      KarrasVeScheduler()

						SCREAMING_SNAKE_CASE__								=      KarrasVePipeline(unet=__UpperCAmelCase    ,				scheduler=__UpperCAmelCase       )
						pipe.to(__UpperCAmelCase       )
						pipe.set_progress_bar_config(disable=__UpperCAmelCase       )

						SCREAMING_SNAKE_CASE__								=      torch.manual_seed(0       )
						SCREAMING_SNAKE_CASE__								=      pipe(num_inference_steps=2    ,				generator=__UpperCAmelCase    ,				output_type="""numpy"""       ).images

						SCREAMING_SNAKE_CASE__								=      torch.manual_seed(0       )
						SCREAMING_SNAKE_CASE__								=      pipe(num_inference_steps=2    ,				generator=__UpperCAmelCase    ,				output_type="""numpy"""    ,				return_dict=__UpperCAmelCase       )[0]

						SCREAMING_SNAKE_CASE__								=      image[0, -3:, -3:, -1]
						SCREAMING_SNAKE_CASE__								=      image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 3_2, 3_2, 3)
						SCREAMING_SNAKE_CASE__								=      np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]       )

						assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2
						assert np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 1e-2





@slow
@require_torch
class    lowerCamelCase						(unittest.TestCase			):

			def 					SCREAMING_SNAKE_CASE				(					self				:						Optional[Any]       )   ->    Optional[int]:
						SCREAMING_SNAKE_CASE__								=      """google/ncsnpp-celebahq-256"""
						SCREAMING_SNAKE_CASE__								=      UNetaDModel.from_pretrained(__UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      KarrasVeScheduler()

						SCREAMING_SNAKE_CASE__								=      KarrasVePipeline(unet=__UpperCAmelCase    ,				scheduler=__UpperCAmelCase       )
						pipe.to(__UpperCAmelCase       )
						pipe.set_progress_bar_config(disable=__UpperCAmelCase       )

						SCREAMING_SNAKE_CASE__								=      torch.manual_seed(0       )
						SCREAMING_SNAKE_CASE__								=      pipe(num_inference_steps=2_0    ,				generator=__UpperCAmelCase    ,				output_type="""numpy"""       ).images

						SCREAMING_SNAKE_CASE__								=      image[0, -3:, -3:, -1]
						assert image.shape == (1, 2_5_6, 2_5_6, 3)
						SCREAMING_SNAKE_CASE__								=      np.array([0.578, 0.5_811, 0.5_924, 0.5_809, 0.587, 0.5_886, 0.5_861, 0.5_802, 0.586]       )

						assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2


