




"""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 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    str             =					"| <pad> <unk> <s> </s> a b c d e f g h i j k".split()
					snake_case		:    Optional[int]             =					dict(zip(lowerCamelCase_					,			range(len(lowerCamelCase_				)				)				)				)

					snake_case		:    str             =					{
					    "unk_token": "<unk>",
					    "bos_token": "<s>",
					    "eos_token": "</s>",
					}
					snake_case		:    Union[str, Any]             =					{
					    "feature_size": 1,
					    "padding_value": 0.0,
					    "sampling_rate": 16_000,
					    "return_attention_mask": False,
					    "do_normalize": True,
					}

					snake_case		:    Any             =					tempfile.mkdtemp()
					snake_case		:    Tuple             =					os.path.join(self.tmpdirname					,			VOCAB_FILES_NAMES["vocab_file"]				)
					snake_case		:    List[Any]             =					os.path.join(self.tmpdirname					,			lowerCamelCase_				)
					with open(self.vocab_file					,			"w"					,			encoding="utf-8"				) as fp:
							fp.write(json.dumps(lowerCamelCase_				) + "\n"				)

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

					# load decoder from hub
					snake_case		:    Optional[Any]             =					"hf-internal-testing/ngram-beam-search-decoder"


			def 				lowerCamelCase_					(  self					,			**SCREAMING_SNAKE_CASE				):



					"""simple docstring"""







					snake_case		:    Dict             =					self.add_kwargs_tokens_map.copy()
					kwargs.update(lowerCamelCase_				)
					return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname					,			**lowerCamelCase_				)


			def 				lowerCamelCase_					(  self					,			**SCREAMING_SNAKE_CASE				):



					"""simple docstring"""







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


			def 				lowerCamelCase_					(  self					,			**SCREAMING_SNAKE_CASE				):



					"""simple docstring"""







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


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					shutil.rmtree(self.tmpdirname				)


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    str             =					self.get_tokenizer()
					snake_case		:    Optional[int]             =					self.get_feature_extractor()
					snake_case		:    Any             =					self.get_decoder()

					snake_case		:    Optional[int]             =					WavaVecaProcessorWithLM(tokenizer=lowerCamelCase_					,			feature_extractor=lowerCamelCase_					,			decoder=lowerCamelCase_				)

					processor.save_pretrained(self.tmpdirname				)
					snake_case		:    Dict             =					WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname				)

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

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

					# decoder
					self.assertEqual(processor.decoder._alphabet.labels					,			decoder._alphabet.labels				)
					self.assertEqual(
					    processor.decoder.model_container[decoder._model_key]._unigram_set					,			decoder.model_container[decoder._model_key]._unigram_set					,			)
					self.assertIsInstance(processor.decoder					,			lowerCamelCase_				)


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    List[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		:    Union[str, Any]             =					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 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    Dict             =					self.get_tokenizer()
					# add token to trigger raise
					tokenizer.add_tokens(["xx"]				)
					with self.assertRaisesRegex(lowerCamelCase_					,			"include"				):
							WavaVecaProcessorWithLM(
							    tokenizer=lowerCamelCase_					,			feature_extractor=self.get_feature_extractor()					,			decoder=self.get_decoder()				)


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    Any             =					self.get_feature_extractor()
					snake_case		:    List[Any]             =					self.get_tokenizer()
					snake_case		:    Union[str, Any]             =					self.get_decoder()

					snake_case		:    Optional[Any]             =					WavaVecaProcessorWithLM(tokenizer=lowerCamelCase_					,			feature_extractor=lowerCamelCase_					,			decoder=lowerCamelCase_				)

					snake_case		:    Union[str, Any]             =					floats_list((3, 1_000)				)

					snake_case		:    Optional[Any]             =					feature_extractor(lowerCamelCase_					,			return_tensors="np"				)
					snake_case		:    List[Any]             =					processor(lowerCamelCase_					,			return_tensors="np"				)

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


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    List[Any]             =					self.get_feature_extractor()
					snake_case		:    Optional[Any]             =					self.get_tokenizer()
					snake_case		:    str             =					self.get_decoder()

					snake_case		:    Any             =					WavaVecaProcessorWithLM(tokenizer=lowerCamelCase_					,			feature_extractor=lowerCamelCase_					,			decoder=lowerCamelCase_				)

					snake_case		:    int             =					"This is a test string"

					snake_case		:    Any             =					processor(text=lowerCamelCase_				)

					snake_case		:    int             =					tokenizer(lowerCamelCase_				)

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


			def 				lowerCamelCase_					(  self					,			SCREAMING_SNAKE_CASE=(2, 10, 16)					,			SCREAMING_SNAKE_CASE=77				):



					"""simple docstring"""







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


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    int             =					self.get_feature_extractor()
					snake_case		:    str             =					self.get_tokenizer()
					snake_case		:    Tuple             =					self.get_decoder()

					snake_case		:    Optional[Any]             =					WavaVecaProcessorWithLM(tokenizer=lowerCamelCase_					,			feature_extractor=lowerCamelCase_					,			decoder=lowerCamelCase_				)

					snake_case		:    List[str]             =					self._get_dummy_logits(shape=(10, 16)					,			seed=13				)

					snake_case		:    str             =					processor.decode(lowerCamelCase_				)

					snake_case		:    Union[str, Any]             =					decoder.decode_beams(lowerCamelCase_				)[0]

					self.assertEqual(decoded_decoder[0]					,			decoded_processor.text				)
					self.assertEqual("</s> <s> </s>"					,			decoded_processor.text				)
					self.assertEqual(decoded_decoder[-2]					,			decoded_processor.logit_score				)
					self.assertEqual(decoded_decoder[-1]					,			decoded_processor.lm_score				)


			@parameterized.expand([[None], ["fork"], ["spawn"]]				)
			def 				lowerCamelCase_					(  self					,			SCREAMING_SNAKE_CASE				):



					"""simple docstring"""







					snake_case		:    Union[str, Any]             =					self.get_feature_extractor()
					snake_case		:    Dict             =					self.get_tokenizer()
					snake_case		:    Dict             =					self.get_decoder()

					snake_case		:    Optional[int]             =					WavaVecaProcessorWithLM(tokenizer=lowerCamelCase_					,			feature_extractor=lowerCamelCase_					,			decoder=lowerCamelCase_				)

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

					snake_case		:    Dict             =					list(lowerCamelCase_				)

					with get_context("fork"				).Pool() as p:
							snake_case		:    int             =					decoder.decode_beams_batch(lowerCamelCase_					,			lowerCamelCase_				)

					snake_case      ,					snake_case      ,					snake_case		:    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(lowerCamelCase_					,			decoded_processor.text				)
					self.assertListEqual(["<s> <s> </s>", "<s> <s> <s>"]					,			decoded_processor.text				)
					self.assertListEqual(lowerCamelCase_					,			decoded_processor.logit_score				)
					self.assertListEqual(lowerCamelCase_					,			decoded_processor.lm_score				)


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    Union[str, Any]             =					self.get_feature_extractor()
					snake_case		:    Any             =					self.get_tokenizer()
					snake_case		:    Optional[int]             =					self.get_decoder()

					snake_case		:    List[Any]             =					WavaVecaProcessorWithLM(tokenizer=lowerCamelCase_					,			feature_extractor=lowerCamelCase_					,			decoder=lowerCamelCase_				)

					snake_case		:    Optional[int]             =					self._get_dummy_logits()

					snake_case		:    Any             =					15
					snake_case		:    Union[str, Any]             =					-20.0
					snake_case		:    Optional[Any]             =					-4.0

					snake_case		:    Tuple             =					processor.batch_decode(
					    lowerCamelCase_					,			beam_width=lowerCamelCase_					,			beam_prune_logp=lowerCamelCase_					,			token_min_logp=lowerCamelCase_					,			)
					snake_case		:    Any             =					decoded_processor_out.text

					snake_case		:    Union[str, Any]             =					list(lowerCamelCase_				)

					with get_context("fork"				).Pool() as pool:
							snake_case		:    Tuple             =					decoder.decode_beams_batch(
							    lowerCamelCase_					,			lowerCamelCase_					,			beam_width=lowerCamelCase_					,			beam_prune_logp=lowerCamelCase_					,			token_min_logp=lowerCamelCase_					,			)

					snake_case		:    str             =					[d[0][0] for d in decoded_decoder_out]
					snake_case		:    Optional[int]             =					[d[0][2] for d in decoded_decoder_out]
					snake_case		:    Tuple             =					[d[0][3] for d in decoded_decoder_out]

					self.assertListEqual(lowerCamelCase_					,			lowerCamelCase_				)
					self.assertListEqual(["</s> <s> <s>", "<s> <s> <s>"]					,			lowerCamelCase_				)

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

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


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    Union[str, Any]             =					self.get_feature_extractor()
					snake_case		:    Dict             =					self.get_tokenizer()
					snake_case		:    List[str]             =					self.get_decoder()

					snake_case		:    Any             =					WavaVecaProcessorWithLM(tokenizer=lowerCamelCase_					,			feature_extractor=lowerCamelCase_					,			decoder=lowerCamelCase_				)

					snake_case		:    Any             =					self._get_dummy_logits()

					snake_case		:    Union[str, Any]             =					2.0
					snake_case		:    int             =					5.0
					snake_case		:    List[str]             =					-20.0
					snake_case		:    Union[str, Any]             =					True

					snake_case		:    Optional[Any]             =					processor.batch_decode(
					    lowerCamelCase_					,			alpha=lowerCamelCase_					,			beta=lowerCamelCase_					,			unk_score_offset=lowerCamelCase_					,			lm_score_boundary=lowerCamelCase_					,			)
					snake_case		:    Union[str, Any]             =					decoded_processor_out.text

					snake_case		:    List[str]             =					list(lowerCamelCase_				)
					decoder.reset_params(
					    alpha=lowerCamelCase_					,			beta=lowerCamelCase_					,			unk_score_offset=lowerCamelCase_					,			lm_score_boundary=lowerCamelCase_					,			)

					with get_context("fork"				).Pool() as pool:
							snake_case		:    str             =					decoder.decode_beams_batch(
							    lowerCamelCase_					,			lowerCamelCase_					,			)

					snake_case		:    List[str]             =					[d[0][0] for d in decoded_decoder_out]

					self.assertListEqual(lowerCamelCase_					,			lowerCamelCase_				)
					self.assertListEqual(["<s> </s> <s> </s> </s>", "</s> </s> <s> </s> </s>"]					,			lowerCamelCase_				)
					snake_case		:    List[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					,			-20.0				)
					self.assertEqual(lm_model.score_boundary					,			lowerCamelCase_				)


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    str             =					WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"				)

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


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    str             =					snapshot_download("hf-internal-testing/processor_with_lm"				)

					snake_case		:    Union[str, Any]             =					WavaVecaProcessorWithLM.from_pretrained(lowerCamelCase_				)

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

					snake_case		:    int             =					os.listdir(lowerCamelCase_				)
					snake_case		:    Optional[Any]             =					os.listdir(lowerCamelCase_				)

					local_decoder_files.sort()
					expected_decoder_files.sort()

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


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    Any             =					WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"				)
					snake_case		:    Any             =					AutoProcessor.from_pretrained("hf-internal-testing/processor_with_lm"				)

					snake_case		:    str             =					floats_list((3, 1_000)				)

					snake_case		:    Any             =					processor_wavaveca(lowerCamelCase_					,			return_tensors="np"				)
					snake_case		:    str             =					processor_auto(lowerCamelCase_					,			return_tensors="np"				)

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

					snake_case		:    Dict             =					self._get_dummy_logits()

					snake_case		:    str             =					processor_wavaveca.batch_decode(lowerCamelCase_				)
					snake_case		:    Any             =					processor_auto.batch_decode(lowerCamelCase_				)

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


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    int             =					self.get_feature_extractor()
					snake_case		:    Union[str, Any]             =					self.get_tokenizer()
					snake_case		:    str             =					self.get_decoder()

					snake_case		:    Any             =					WavaVecaProcessorWithLM(tokenizer=lowerCamelCase_					,			feature_extractor=lowerCamelCase_					,			decoder=lowerCamelCase_				)

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


			@staticmethod
			def 				lowerCamelCase_					(  SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE				):



					"""simple docstring"""







					snake_case		:    List[Any]             =					[d[key] for d in offsets]
					return retrieved_list


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    List[str]             =					WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"				)
					snake_case		:    Dict             =					self._get_dummy_logits()[0]

					snake_case		:    List[Any]             =					processor.decode(lowerCamelCase_					,			output_word_offsets=lowerCamelCase_				)
					# check Wav2Vec2CTCTokenizerOutput keys for word
					self.assertEqual(len(outputs.keys()				)					,			4				)
					self.assertTrue("text" in outputs				)
					self.assertTrue("word_offsets" in outputs				)
					self.assertTrue(isinstance(lowerCamelCase_					,			lowerCamelCase_				)				)

					self.assertEqual(" ".join(self.get_from_offsets(outputs["word_offsets"]					,			"word"				)				)					,			outputs.text				)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"]					,			"word"				)					,			["<s>", "<s>", "</s>"]				)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"]					,			"start_offset"				)					,			[0, 2, 4]				)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"]					,			"end_offset"				)					,			[1, 3, 5]				)


			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					snake_case		:    List[str]             =					WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"				)
					snake_case		:    Union[str, Any]             =					self._get_dummy_logits()

					snake_case		:    int             =					processor.batch_decode(lowerCamelCase_					,			output_word_offsets=lowerCamelCase_				)

					# check Wav2Vec2CTCTokenizerOutput keys for word
					self.assertEqual(len(outputs.keys()				)					,			4				)
					self.assertTrue("text" in outputs				)
					self.assertTrue("word_offsets" in outputs				)
					self.assertTrue(isinstance(lowerCamelCase_					,			lowerCamelCase_				)				)

					self.assertListEqual(
					    [" ".join(self.get_from_offsets(lowerCamelCase_					,			"word"				)				) for o in outputs["word_offsets"]]					,			outputs.text				)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"][0]					,			"word"				)					,			["<s>", "<s>", "</s>"]				)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"][0]					,			"start_offset"				)					,			[0, 2, 4]				)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"][0]					,			"end_offset"				)					,			[1, 3, 5]				)



			@slow
			@require_torch
			@require_torchaudio
			def 				lowerCamelCase_					(  self				):



					"""simple docstring"""







					import torch

					snake_case		:    Union[str, Any]             =					load_dataset("common_voice"					,			"en"					,			split="train"					,			streaming=lowerCamelCase_				)
					snake_case		:    Dict             =					ds.cast_column("audio"					,			datasets.Audio(sampling_rate=16_000				)				)
					snake_case		:    Union[str, Any]             =					iter(lowerCamelCase_				)
					snake_case		:    List[str]             =					next(lowerCamelCase_				)

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

					with torch.no_grad():
							snake_case		:    Dict             =					model(lowerCamelCase_				).logits.cpu().numpy()

					snake_case		:    Optional[Any]             =					processor.decode(logits[0]					,			output_word_offsets=lowerCamelCase_				)

					snake_case		:    Any             =					model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
					snake_case		:    Any             =					[
					    {
					        "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		:    Dict             =					"WHY DOES MILISANDRA LOOK LIKE SHE WANTS TO CONSUME JOHN SNOW ON THE RIVER AT THE WALL"

					# output words
					self.assertEqual(" ".join(self.get_from_offsets(lowerCamelCase_					,			"word"				)				)					,			lowerCamelCase_				)
					self.assertEqual(" ".join(self.get_from_offsets(lowerCamelCase_					,			"word"				)				)					,			output.text				)

					# output times
					snake_case		:    Dict             =					torch.tensor(self.get_from_offsets(lowerCamelCase_					,			"start_time"				)				)
					snake_case		:    Any             =					torch.tensor(self.get_from_offsets(lowerCamelCase_					,			"end_time"				)				)

					# fmt: off
					snake_case		:    Optional[Any]             =					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]				)
					snake_case		:    Tuple             =					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(lowerCamelCase_					,			lowerCamelCase_					,			atol=0.01				)				)
					self.assertTrue(torch.allclose(lowerCamelCase_					,			lowerCamelCase_					,			atol=0.01				)				)


import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin
class SCREAMING_SNAKE_CASE_		(      __lowerCAmelCase  ,     __lowerCAmelCase		):







   @register_to_config
   def __init__(     self						:						List[Any]   ,   lowerCamelCase_						:						int = 128   ,   lowerCamelCase_						:						int = 256   ,   lowerCamelCase_						:						float = 2_0_0_0.0   ,   lowerCamelCase_						:						int = 768   ,   lowerCamelCase_						:						int = 12   ,   lowerCamelCase_						:						int = 12   ,   lowerCamelCase_						:						int = 64   ,   lowerCamelCase_						:						int = 2048   ,   lowerCamelCase_						:						float = 0.1   ,   ):





    """simple docstring"""


    super().__init__()

    UpperCamelCase			  =      nn.Sequential(
        nn.Linear(lowerCamelCase_   ,   d_model * 4   ,   bias=lowerCamelCase_				)   ,   nn.SiLU()   ,   nn.Linear(d_model * 4   ,   d_model * 4   ,   bias=lowerCamelCase_				)   ,   nn.SiLU()   ,   )

    UpperCamelCase			  =      nn.Embedding(lowerCamelCase_   ,   lowerCamelCase_				)
    UpperCamelCase			  =      False

    UpperCamelCase			  =      nn.Linear(lowerCamelCase_   ,   lowerCamelCase_   ,   bias=lowerCamelCase_				)

    UpperCamelCase			  =      nn.Dropout(p=lowerCamelCase_				)

    UpperCamelCase			  =      nn.ModuleList()
    for lyr_num in range(lowerCamelCase_				):
     # FiLM conditional T5 decoder
     UpperCamelCase			  =      DecoderLayer(d_model=lowerCamelCase_   ,   d_kv=lowerCamelCase_   ,   num_heads=lowerCamelCase_   ,   d_ff=lowerCamelCase_   ,   dropout_rate=lowerCamelCase_				)
     self.decoders.append(lowerCamelCase_				)

    UpperCamelCase			  =      TaLayerNorm(lowerCamelCase_				)

    UpperCamelCase			  =      nn.Dropout(p=lowerCamelCase_				)
    UpperCamelCase			  =      nn.Linear(lowerCamelCase_   ,   lowerCamelCase_   ,   bias=lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						Union[str, Any]   ,   lowerCamelCase_						:						int   ,   lowerCamelCase_						:						List[str]				):





    """simple docstring"""


    UpperCamelCase			  =      torch.mul(query_input.unsqueeze(-1				)   ,   key_input.unsqueeze(-2				)				)
    return mask.unsqueeze(-3				)







   def 						lowerCamelCase_	(     self						:						List[Any]   ,   lowerCamelCase_						:						List[str]   ,   lowerCamelCase_						:						Optional[int]   ,   lowerCamelCase_						:						Union[str, Any]				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      decoder_input_tokens.shape
    assert decoder_noise_time.shape == (batch,)

    # decoder_noise_time is in [0, 1), so rescale to expected timing range.
    UpperCamelCase			  =      get_timestep_embedding(
        decoder_noise_time * self.config.max_decoder_noise_time   ,   embedding_dim=self.config.d_model   ,   max_period=self.config.max_decoder_noise_time   ,   ).to(dtype=self.dtype				)

    UpperCamelCase			  =      self.conditioning_emb(lowerCamelCase_				).unsqueeze(1				)

    assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

    UpperCamelCase			  =      decoder_input_tokens.shape[1]

    # If we want to use relative positions for audio context, we can just offset
    # this sequence by the length of encodings_and_masks.
    UpperCamelCase			  =      torch.broadcast_to(
        torch.arange(lowerCamelCase_   ,   device=decoder_input_tokens.device				)   ,   (batch, seq_length)   ,   )

    UpperCamelCase			  =      self.position_encoding(lowerCamelCase_				)

    UpperCamelCase			  =      self.continuous_inputs_projection(lowerCamelCase_				)
    inputs += position_encodings
    UpperCamelCase			  =      self.dropout(lowerCamelCase_				)

    # decoder: No padding present.
    UpperCamelCase			  =      torch.ones(
        decoder_input_tokens.shape[:2]   ,   device=decoder_input_tokens.device   ,   dtype=inputs.dtype				)

    # Translate encoding masks to encoder-decoder masks.
    UpperCamelCase			  =      [(x, self.encoder_decoder_mask(lowerCamelCase_   ,   lowerCamelCase_				)) for x, y in encodings_and_masks]

    # cross attend style: concat encodings
    UpperCamelCase			  =      torch.cat([x[0] for x in encodings_and_encdec_masks]   ,   dim=1				)
    UpperCamelCase			  =      torch.cat([x[1] for x in encodings_and_encdec_masks]   ,   dim=-1				)

    for lyr in self.decoders:
     UpperCamelCase			  =      lyr(
         lowerCamelCase_   ,   conditioning_emb=lowerCamelCase_   ,   encoder_hidden_states=lowerCamelCase_   ,   encoder_attention_mask=lowerCamelCase_   ,   )[0]

    UpperCamelCase			  =      self.decoder_norm(lowerCamelCase_				)
    UpperCamelCase			  =      self.post_dropout(lowerCamelCase_				)

    UpperCamelCase			  =      self.spec_out(lowerCamelCase_				)
    return spec_out





class SCREAMING_SNAKE_CASE_		(      nn.Module		):







   def __init__(     self						:						List[str]   ,   lowerCamelCase_						:						int   ,   lowerCamelCase_						:						Tuple   ,   lowerCamelCase_						:						List[str]   ,   lowerCamelCase_						:						List[str]   ,   lowerCamelCase_						:						Tuple   ,   lowerCamelCase_						:						Any=1E-6				):





    """simple docstring"""


    super().__init__()
    UpperCamelCase			  =      nn.ModuleList()

    # cond self attention: layer 0
    self.layer.append(
        TaLayerSelfAttentionCond(d_model=lowerCamelCase_   ,   d_kv=lowerCamelCase_   ,   num_heads=lowerCamelCase_   ,   dropout_rate=lowerCamelCase_				)				)

    # cross attention: layer 1
    self.layer.append(
        TaLayerCrossAttention(
            d_model=lowerCamelCase_   ,   d_kv=lowerCamelCase_   ,   num_heads=lowerCamelCase_   ,   dropout_rate=lowerCamelCase_   ,   layer_norm_epsilon=lowerCamelCase_   ,   )				)

    # Film Cond MLP + dropout: last layer
    self.layer.append(
        TaLayerFFCond(d_model=lowerCamelCase_   ,   d_ff=lowerCamelCase_   ,   dropout_rate=lowerCamelCase_   ,   layer_norm_epsilon=lowerCamelCase_				)				)







   def 						lowerCamelCase_	(     self						:						List[str]   ,   lowerCamelCase_						:						Tuple   ,   lowerCamelCase_						:						List[Any]=None   ,   lowerCamelCase_						:						int=None   ,   lowerCamelCase_						:						Optional[int]=None   ,   lowerCamelCase_						:						Any=None   ,   lowerCamelCase_						:						Union[str, Any]=None   ,   ):





    """simple docstring"""


    UpperCamelCase			  =      self.layer[0](
        lowerCamelCase_   ,   conditioning_emb=lowerCamelCase_   ,   attention_mask=lowerCamelCase_   ,   )

    if encoder_hidden_states is not None:
     UpperCamelCase			  =      torch.where(encoder_attention_mask > 0   ,   0   ,   -1E10				).to(
         encoder_hidden_states.dtype				)

     UpperCamelCase			  =      self.layer[1](
         lowerCamelCase_   ,   key_value_states=lowerCamelCase_   ,   attention_mask=lowerCamelCase_   ,   )

    # Apply Film Conditional Feed Forward layer
    UpperCamelCase			  =      self.layer[-1](lowerCamelCase_   ,   lowerCamelCase_				)

    return (hidden_states,)





class SCREAMING_SNAKE_CASE_		(      nn.Module		):







   def __init__(     self						:						Dict   ,   lowerCamelCase_						:						Any   ,   lowerCamelCase_						:						Tuple   ,   lowerCamelCase_						:						str   ,   lowerCamelCase_						:						List[Any]				):





    """simple docstring"""


    super().__init__()
    UpperCamelCase			  =      TaLayerNorm(lowerCamelCase_				)
    UpperCamelCase			  =      TaFiLMLayer(in_features=d_model * 4   ,   out_features=lowerCamelCase_				)
    UpperCamelCase			  =      Attention(query_dim=lowerCamelCase_   ,   heads=lowerCamelCase_   ,   dim_head=lowerCamelCase_   ,   out_bias=lowerCamelCase_   ,   scale_qk=lowerCamelCase_				)
    UpperCamelCase			  =      nn.Dropout(lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						int   ,   lowerCamelCase_						:						Optional[Any]   ,   lowerCamelCase_						:						Tuple=None   ,   lowerCamelCase_						:						Optional[Any]=None   ,   ):





    """simple docstring"""


    UpperCamelCase			  =      self.layer_norm(lowerCamelCase_				)

    if conditioning_emb is not None:
     UpperCamelCase			  =      self.FiLMLayer(lowerCamelCase_   ,   lowerCamelCase_				)

    # Self-attention block
    UpperCamelCase			  =      self.attention(lowerCamelCase_				)

    UpperCamelCase			  =      hidden_states + self.dropout(lowerCamelCase_				)

    return hidden_states





class SCREAMING_SNAKE_CASE_		(      nn.Module		):







   def __init__(     self						:						Dict   ,   lowerCamelCase_						:						str   ,   lowerCamelCase_						:						int   ,   lowerCamelCase_						:						Optional[Any]   ,   lowerCamelCase_						:						str   ,   lowerCamelCase_						:						Dict				):





    """simple docstring"""


    super().__init__()
    UpperCamelCase			  =      Attention(query_dim=lowerCamelCase_   ,   heads=lowerCamelCase_   ,   dim_head=lowerCamelCase_   ,   out_bias=lowerCamelCase_   ,   scale_qk=lowerCamelCase_				)
    UpperCamelCase			  =      TaLayerNorm(lowerCamelCase_   ,   eps=lowerCamelCase_				)
    UpperCamelCase			  =      nn.Dropout(lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						Optional[Any]   ,   lowerCamelCase_						:						List[str]   ,   lowerCamelCase_						:						Dict=None   ,   lowerCamelCase_						:						int=None   ,   ):





    """simple docstring"""


    UpperCamelCase			  =      self.layer_norm(lowerCamelCase_				)
    UpperCamelCase			  =      self.attention(
        lowerCamelCase_   ,   encoder_hidden_states=lowerCamelCase_   ,   attention_mask=attention_mask.squeeze(1				)   ,   )
    UpperCamelCase			  =      hidden_states + self.dropout(lowerCamelCase_				)
    return layer_output





class SCREAMING_SNAKE_CASE_		(      nn.Module		):







   def __init__(     self						:						Union[str, Any]   ,   lowerCamelCase_						:						Any   ,   lowerCamelCase_						:						str   ,   lowerCamelCase_						:						Optional[int]   ,   lowerCamelCase_						:						Dict				):





    """simple docstring"""


    super().__init__()
    UpperCamelCase			  =      TaDenseGatedActDense(d_model=lowerCamelCase_   ,   d_ff=lowerCamelCase_   ,   dropout_rate=lowerCamelCase_				)
    UpperCamelCase			  =      TaFiLMLayer(in_features=d_model * 4   ,   out_features=lowerCamelCase_				)
    UpperCamelCase			  =      TaLayerNorm(lowerCamelCase_   ,   eps=lowerCamelCase_				)
    UpperCamelCase			  =      nn.Dropout(lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						List[Any]   ,   lowerCamelCase_						:						List[Any]   ,   lowerCamelCase_						:						int=None				):





    """simple docstring"""


    UpperCamelCase			  =      self.layer_norm(lowerCamelCase_				)
    if conditioning_emb is not None:
     UpperCamelCase			  =      self.film(lowerCamelCase_   ,   lowerCamelCase_				)

    UpperCamelCase			  =      self.DenseReluDense(lowerCamelCase_				)
    UpperCamelCase			  =      hidden_states + self.dropout(lowerCamelCase_				)
    return hidden_states





class SCREAMING_SNAKE_CASE_		(      nn.Module		):







   def __init__(     self						:						List[Any]   ,   lowerCamelCase_						:						int   ,   lowerCamelCase_						:						List[Any]   ,   lowerCamelCase_						:						List[Any]				):





    """simple docstring"""


    super().__init__()
    UpperCamelCase			  =      nn.Linear(lowerCamelCase_   ,   lowerCamelCase_   ,   bias=lowerCamelCase_				)
    UpperCamelCase			  =      nn.Linear(lowerCamelCase_   ,   lowerCamelCase_   ,   bias=lowerCamelCase_				)
    UpperCamelCase			  =      nn.Linear(lowerCamelCase_   ,   lowerCamelCase_   ,   bias=lowerCamelCase_				)
    UpperCamelCase			  =      nn.Dropout(lowerCamelCase_				)
    UpperCamelCase			  =      NewGELUActivation()







   def 						lowerCamelCase_	(     self						:						Union[str, Any]   ,   lowerCamelCase_						:						Optional[int]				):





    """simple docstring"""


    UpperCamelCase			  =      self.act(self.wi_a(lowerCamelCase_				)				)
    UpperCamelCase			  =      self.wi_a(lowerCamelCase_				)
    UpperCamelCase			  =      hidden_gelu * hidden_linear
    UpperCamelCase			  =      self.dropout(lowerCamelCase_				)

    UpperCamelCase			  =      self.wo(lowerCamelCase_				)
    return hidden_states





class SCREAMING_SNAKE_CASE_		(      nn.Module		):







   def __init__(     self						:						Optional[int]   ,   lowerCamelCase_						:						Tuple   ,   lowerCamelCase_						:						Dict=1E-6				):





    """simple docstring"""


    super().__init__()
    UpperCamelCase			  =      nn.Parameter(torch.ones(lowerCamelCase_				)				)
    UpperCamelCase			  =      eps







   def 						lowerCamelCase_	(     self						:						List[Any]   ,   lowerCamelCase_						:						Any				):





    """simple docstring"""


    UpperCamelCase			  =      hidden_states.to(torch.floataa				).pow(2				).mean(-1   ,   keepdim=lowerCamelCase_				)
    UpperCamelCase			  =      hidden_states * torch.rsqrt(variance + self.variance_epsilon				)

    # convert into half-precision if necessary
    if self.weight.dtype in [torch.floataa, torch.bfloataa]:
     UpperCamelCase			  =      hidden_states.to(self.weight.dtype				)

    return self.weight * hidden_states





class SCREAMING_SNAKE_CASE_		(      nn.Module		):







   def 						lowerCamelCase_	(     self						:						List[Any]   ,   lowerCamelCase_						:						torch.Tensor				):





    """simple docstring"""


    return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi				) * (input + 0.0_4_4_7_1_5 * torch.pow(lowerCamelCase_   ,   3.0				))				))





class SCREAMING_SNAKE_CASE_		(      nn.Module		):







   def __init__(     self						:						Optional[Any]   ,   lowerCamelCase_						:						List[str]   ,   lowerCamelCase_						:						str				):





    """simple docstring"""


    super().__init__()
    UpperCamelCase			  =      nn.Linear(lowerCamelCase_   ,   out_features * 2   ,   bias=lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						List[str]   ,   lowerCamelCase_						:						Optional[Any]   ,   lowerCamelCase_						:						Dict				):





    """simple docstring"""


    UpperCamelCase			  =      self.scale_bias(lowerCamelCase_				)
    UpperCamelCase  ,     UpperCamelCase			  =      torch.chunk(lowerCamelCase_   ,   2   ,   -1				)
    UpperCamelCase			  =      x * (1 + scale) + shift
    return x
