





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





					"""simple docstring"""


					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       "| <pad> <unk> <s> </s> a b c d e f g h i j k".split()
										lowerCamelCase__       =       dict(zip(_lowerCAmelCase	,  range(len(_lowerCAmelCase	)	)	)	)

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

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





					def 					__magic_name__ ( self	,  **_lowerCAmelCase	):
										lowerCamelCase__       =       self.add_kwargs_tokens_map.copy()
										kwargs.update(_lowerCAmelCase	)
										return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname	,  **_lowerCAmelCase	)





					def 					__magic_name__ ( self	,  **_lowerCAmelCase	):
										return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname	,  **_lowerCAmelCase	)





					def 					__magic_name__ ( self	,  **_lowerCAmelCase	):
										return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name	,  **_lowerCAmelCase	)





					def 					__magic_name__ ( self	):
										shutil.rmtree(self.tmpdirname	)





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.get_tokenizer()
										lowerCamelCase__       =       self.get_feature_extractor()
										lowerCamelCase__       =       self.get_decoder()

										lowerCamelCase__       =       WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase	,  feature_extractor=_lowerCAmelCase	,  decoder=_lowerCAmelCase	)

										processor.save_pretrained(self.tmpdirname	)
										lowerCamelCase__       =       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 					__magic_name__ ( self	):
										lowerCamelCase__       =       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
										lowerCamelCase__       =       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 					__magic_name__ ( self	):
										lowerCamelCase__       =       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 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.get_feature_extractor()
										lowerCamelCase__       =       self.get_tokenizer()
										lowerCamelCase__       =       self.get_decoder()

										lowerCamelCase__       =       WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase	,  feature_extractor=_lowerCAmelCase	,  decoder=_lowerCAmelCase	)

										lowerCamelCase__       =       floats_list((3, 1000)	)

										lowerCamelCase__       =       feature_extractor(_lowerCAmelCase	,  return_tensors="np"	)
										lowerCamelCase__       =       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 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.get_feature_extractor()
										lowerCamelCase__       =       self.get_tokenizer()
										lowerCamelCase__       =       self.get_decoder()

										lowerCamelCase__       =       WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase	,  feature_extractor=_lowerCAmelCase	,  decoder=_lowerCAmelCase	)

										lowerCamelCase__       =       "This is a test string"

										lowerCamelCase__       =       processor(text=_lowerCAmelCase	)

										lowerCamelCase__       =       tokenizer(_lowerCAmelCase	)

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





					def 					__magic_name__ ( self	,  _lowerCAmelCase=(2, 10, 16)	,  _lowerCAmelCase=77	):
										np.random.seed(_lowerCAmelCase	)
										return np.random.rand(*_lowerCAmelCase	)





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.get_feature_extractor()
										lowerCamelCase__       =       self.get_tokenizer()
										lowerCamelCase__       =       self.get_decoder()

										lowerCamelCase__       =       WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase	,  feature_extractor=_lowerCAmelCase	,  decoder=_lowerCAmelCase	)

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

										lowerCamelCase__       =       processor.decode(_lowerCAmelCase	)

										lowerCamelCase__       =       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 					__magic_name__ ( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       self.get_feature_extractor()
										lowerCamelCase__       =       self.get_tokenizer()
										lowerCamelCase__       =       self.get_decoder()

										lowerCamelCase__       =       WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase	,  feature_extractor=_lowerCAmelCase	,  decoder=_lowerCAmelCase	)

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

										lowerCamelCase__       =       list(_lowerCAmelCase	)

										with get_context("fork"	).Pool() as p:
															lowerCamelCase__       =       decoder.decode_beams_batch(_lowerCAmelCase	,  _lowerCAmelCase	)

										lowerCamelCase__    ,		lowerCamelCase__    ,		lowerCamelCase__       =       [], [], []
										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 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.get_feature_extractor()
										lowerCamelCase__       =       self.get_tokenizer()
										lowerCamelCase__       =       self.get_decoder()

										lowerCamelCase__       =       WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase	,  feature_extractor=_lowerCAmelCase	,  decoder=_lowerCAmelCase	)

										lowerCamelCase__       =       self._get_dummy_logits()

										lowerCamelCase__       =       15
										lowerCamelCase__       =       -20.0
										lowerCamelCase__       =       -4.0

										lowerCamelCase__       =       processor.batch_decode(
										    _lowerCAmelCase	,  beam_width=_lowerCAmelCase	,  beam_prune_logp=_lowerCAmelCase	,  token_min_logp=_lowerCAmelCase	,  )
										lowerCamelCase__       =       decoded_processor_out.text

										lowerCamelCase__       =       list(_lowerCAmelCase	)

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

										lowerCamelCase__       =       [d[0][0] for d in decoded_decoder_out]
										lowerCamelCase__       =       [d[0][2] for d in decoded_decoder_out]
										lowerCamelCase__       =       [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 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.get_feature_extractor()
										lowerCamelCase__       =       self.get_tokenizer()
										lowerCamelCase__       =       self.get_decoder()

										lowerCamelCase__       =       WavaVecaProcessorWithLM(tokenizer=_lowerCAmelCase	,  feature_extractor=_lowerCAmelCase	,  decoder=_lowerCAmelCase	)

										lowerCamelCase__       =       self._get_dummy_logits()

										lowerCamelCase__       =       2.0
										lowerCamelCase__       =       5.0
										lowerCamelCase__       =       -20.0
										lowerCamelCase__       =       True

										lowerCamelCase__       =       processor.batch_decode(
										    _lowerCAmelCase	,  alpha=_lowerCAmelCase	,  beta=_lowerCAmelCase	,  unk_score_offset=_lowerCAmelCase	,  lm_score_boundary=_lowerCAmelCase	,  )
										lowerCamelCase__       =       decoded_processor_out.text

										lowerCamelCase__       =       list(_lowerCAmelCase	)
										decoder.reset_params(
										    alpha=_lowerCAmelCase	,  beta=_lowerCAmelCase	,  unk_score_offset=_lowerCAmelCase	,  lm_score_boundary=_lowerCAmelCase	,  )

										with get_context("fork"	).Pool() as pool:
															lowerCamelCase__       =       decoder.decode_beams_batch(
															    _lowerCAmelCase	,  _lowerCAmelCase	,  )

										lowerCamelCase__       =       [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	)
										lowerCamelCase__       =       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 					__magic_name__ ( self	):
										lowerCamelCase__       =       WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"	)

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

										lowerCamelCase__       =       os.listdir(_lowerCAmelCase	)
										lowerCamelCase__       =       ["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 					__magic_name__ ( self	):
										lowerCamelCase__       =       snapshot_download("hf-internal-testing/processor_with_lm"	)

										lowerCamelCase__       =       WavaVecaProcessorWithLM.from_pretrained(_lowerCAmelCase	)

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

										lowerCamelCase__       =       os.listdir(_lowerCAmelCase	)
										lowerCamelCase__       =       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 					__magic_name__ ( self	):
										lowerCamelCase__       =       WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"	)
										lowerCamelCase__       =       AutoProcessor.from_pretrained("hf-internal-testing/processor_with_lm"	)

										lowerCamelCase__       =       floats_list((3, 1000)	)

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

										lowerCamelCase__       =       self._get_dummy_logits()

										lowerCamelCase__       =       processor_wavaveca.batch_decode(_lowerCAmelCase	)
										lowerCamelCase__       =       processor_auto.batch_decode(_lowerCAmelCase	)

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





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.get_feature_extractor()
										lowerCamelCase__       =       self.get_tokenizer()
										lowerCamelCase__       =       self.get_decoder()

										lowerCamelCase__       =       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 					__magic_name__ ( _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       [d[key] for d in offsets]
										return retrieved_list





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"	)
										lowerCamelCase__       =       self._get_dummy_logits()[0]

										lowerCamelCase__       =       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 					__magic_name__ ( self	):
										lowerCamelCase__       =       WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"	)
										lowerCamelCase__       =       self._get_dummy_logits()

										lowerCamelCase__       =       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 					__magic_name__ ( self	):
										import torch

										lowerCamelCase__       =       load_dataset("common_voice"	,  "en"	,  split="train"	,  streaming=_lowerCAmelCase	)
										lowerCamelCase__       =       ds.cast_column("audio"	,  datasets.Audio(sampling_rate=1_6000	)	)
										lowerCamelCase__       =       iter(_lowerCAmelCase	)
										lowerCamelCase__       =       next(_lowerCAmelCase	)

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

										with torch.no_grad():
															lowerCamelCase__       =       model(_lowerCAmelCase	).logits.cpu().numpy()

										lowerCamelCase__       =       processor.decode(logits[0]	,  output_word_offsets=_lowerCAmelCase	)

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

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

										# fmt: off
										lowerCamelCase__       =       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]	)
										lowerCamelCase__       =       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(_lowerCAmelCase	,  _lowerCAmelCase	,  atol=0.01	)	)
										self.assertTrue(torch.allclose(_lowerCAmelCase	,  _lowerCAmelCase	,  atol=0.01	)	)








import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


A_												=     get_logger(__name__)


A_												=     r"\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n"

class SCREAMING_SNAKE_CASE_ :





					"""simple docstring"""


					@add_start_docstrings(_lowerCAmelCase	)
					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										raise NotImplementedError(
										    F"{self.__class__} is an abstract class. Only classes inheriting this class can be called."	)

class SCREAMING_SNAKE_CASE_ :





					"""simple docstring"""


					@add_start_docstrings(_lowerCAmelCase	)
					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										raise NotImplementedError(
										    F"{self.__class__} is an abstract class. Only classes inheriting this class can be called."	)

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					@add_start_docstrings(_lowerCAmelCase	)
					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  **_lowerCAmelCase	):
										for processor in self:
															lowerCamelCase__       =       inspect.signature(processor.__call__	).parameters
															if len(_lowerCAmelCase	) > 3:
																				if not all(arg in kwargs for arg in list(function_args.keys()	)[2:]	):
																									raise ValueError(
																									    F"Make sure that all the required parameters: {list(function_args.keys()	)} for "
																									    F"{processor.__class__} are passed to the logits processor."	)
																				lowerCamelCase__       =       processor(_lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  **_lowerCAmelCase	)
															else:
																				lowerCamelCase__       =       processor(_lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	)
										return scores

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	):
										if not isinstance(_lowerCAmelCase	,  _lowerCAmelCase	) or not (temperature > 0):
															raise ValueError(F"`temperature` has to be a strictly positive float, but is {temperature}"	)

										lowerCamelCase__       =       temperature






					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       scores / self.temperature
										return scores

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase = -float("Inf"	)	,  _lowerCAmelCase = 1	):
										if not isinstance(_lowerCAmelCase	,  _lowerCAmelCase	) or (top_p < 0 or top_p > 1.0):
															raise ValueError(F"`top_p` has to be a float > 0 and < 1, but is {top_p}"	)
										if not isinstance(_lowerCAmelCase	,  _lowerCAmelCase	) or (min_tokens_to_keep < 1):
															raise ValueError(F"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}"	)

										lowerCamelCase__       =       top_p
										lowerCamelCase__       =       filter_value
										lowerCamelCase__       =       min_tokens_to_keep






					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__    ,		lowerCamelCase__       =       lax.top_k(_lowerCAmelCase	,  scores.shape[-1]	)

										lowerCamelCase__       =       jnp.full_like(_lowerCAmelCase	,  self.filter_value	)
										lowerCamelCase__       =       jax.nn.softmax(_lowerCAmelCase	,  axis=-1	).cumsum(axis=-1	)
										lowerCamelCase__       =       cumulative_probs < self.top_p

										# include the token that is higher than top_p as well
										lowerCamelCase__       =       jnp.roll(_lowerCAmelCase	,  1	)
										score_mask |= score_mask.at[:, 0].set(_lowerCAmelCase	)

										# min tokens to keep
										lowerCamelCase__       =       score_mask.at[:, : self.min_tokens_to_keep].set(_lowerCAmelCase	)

										lowerCamelCase__       =       jnp.where(_lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	)
										lowerCamelCase__       =       jax.lax.sort_key_val(_lowerCAmelCase	,  _lowerCAmelCase	)[-1]

										return next_scores

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase = -float("Inf"	)	,  _lowerCAmelCase = 1	):
										if not isinstance(_lowerCAmelCase	,  _lowerCAmelCase	) or top_k <= 0:
															raise ValueError(F"`top_k` has to be a strictly positive integer, but is {top_k}"	)

										lowerCamelCase__       =       max(_lowerCAmelCase	,  _lowerCAmelCase	)
										lowerCamelCase__       =       filter_value






					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__    ,		lowerCamelCase__       =       scores.shape
										lowerCamelCase__       =       jnp.full(batch_size * vocab_size	,  self.filter_value	)

										lowerCamelCase__       =       min(self.top_k	,  scores.shape[-1]	)  # Safety check
										lowerCamelCase__    ,		lowerCamelCase__       =       lax.top_k(_lowerCAmelCase	,  _lowerCAmelCase	)
										lowerCamelCase__       =       jnp.broadcast_to((jnp.arange(_lowerCAmelCase	) * vocab_size)[:, None]	,  (batch_size, topk)	).flatten()
										lowerCamelCase__       =       topk_scores.flatten()
										lowerCamelCase__       =       topk_indices.flatten() + shift

										lowerCamelCase__       =       next_scores_flat.at[topk_indices_flat].set(_lowerCAmelCase	)
										lowerCamelCase__       =       next_scores_flat.reshape(_lowerCAmelCase	,  _lowerCAmelCase	)
										return next_scores

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       bos_token_id






					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       jnp.full(scores.shape	,  -float("inf"	)	)

										lowerCamelCase__       =       1 - jnp.bool_(cur_len - 1	)

										lowerCamelCase__       =       jnp.where(_lowerCAmelCase	,  new_scores.at[:, self.bos_token_id].set(0	)	,  _lowerCAmelCase	)

										return scores

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       max_length
										lowerCamelCase__       =       eos_token_id






					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       jnp.full(scores.shape	,  -float("inf"	)	)

										lowerCamelCase__       =       1 - jnp.bool_(cur_len - self.max_length + 1	)

										lowerCamelCase__       =       jnp.where(_lowerCAmelCase	,  new_scores.at[:, self.eos_token_id].set(0	)	,  _lowerCAmelCase	)

										return scores

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										if not isinstance(_lowerCAmelCase	,  _lowerCAmelCase	) or min_length < 0:
															raise ValueError(F"`min_length` has to be a positive integer, but is {min_length}"	)

										if not isinstance(_lowerCAmelCase	,  _lowerCAmelCase	) or eos_token_id < 0:
															raise ValueError(F"`eos_token_id` has to be a positive integer, but is {eos_token_id}"	)

										lowerCamelCase__       =       min_length
										lowerCamelCase__       =       eos_token_id






					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										# create boolean flag to decide if min length penalty should be applied
										lowerCamelCase__       =       1 - jnp.clip(cur_len - self.min_length	,  0	,  1	)

										lowerCamelCase__       =       jnp.where(_lowerCAmelCase	,  scores.at[:, self.eos_token_id].set(-float("inf"	)	)	,  _lowerCAmelCase	)

										return scores

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       list(_lowerCAmelCase	)
										lowerCamelCase__       =       begin_index






					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       1 - jnp.bool_(cur_len - self.begin_index	)

										lowerCamelCase__       =       jnp.where(_lowerCAmelCase	,  scores.at[:, self.begin_suppress_tokens].set(-float("inf"	)	)	,  _lowerCAmelCase	)

										return scores

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       list(_lowerCAmelCase	)






					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       scores.at[..., self.suppress_tokens].set(-float("inf"	)	)

										return scores

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       dict(_lowerCAmelCase	)
										# Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
										# index of the array corresponds to the index of the token to be forced, for XLA compatibility.
										# Indexes without forced tokens will have a negative value.
										lowerCamelCase__       =       jnp.ones((max(force_token_map.keys()	) + 1)	,  dtype=jnp.intaa	) * -1
										for index, token in force_token_map.items():
															if token is not None:
																				lowerCamelCase__       =       force_token_array.at[index].set(_lowerCAmelCase	)
										lowerCamelCase__       =       jnp.intaa(_lowerCAmelCase	)






					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										def _force_token(_lowerCAmelCase	):
															lowerCamelCase__       =       scores.shape[0]
															lowerCamelCase__       =       self.force_token_array[generation_idx]

															lowerCamelCase__       =       jnp.ones_like(_lowerCAmelCase	,  dtype=scores.dtype	) * -float("inf"	)
															lowerCamelCase__       =       jnp.zeros((batch_size, 1)	,  dtype=scores.dtype	)
															lowerCamelCase__       =       lax.dynamic_update_slice(_lowerCAmelCase	,  _lowerCAmelCase	,  (0, current_token)	)
															return new_scores

										lowerCamelCase__       =       lax.cond(
										    cur_len >= self.force_token_array.shape[0]	,  lambda: scores	,  lambda: lax.cond(
										        self.force_token_array[cur_len] >= 0	,  lambda: _force_token(_lowerCAmelCase	)	,  lambda: scores	,  )	,  )
										return scores

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       generate_config.eos_token_id
										lowerCamelCase__       =       generate_config.no_timestamps_token_id
										lowerCamelCase__       =       generate_config.no_timestamps_token_id + 1

										lowerCamelCase__       =       decoder_input_length + 1

										if generate_config.is_multilingual:
															# room for language token and task token
															self.begin_index += 2
										if hasattr(_lowerCAmelCase	,  "max_initial_timestamp_index"	):
															lowerCamelCase__       =       generate_config.max_initial_timestamp_index
										else:
															lowerCamelCase__       =       model_config.vocab_size
										if self.max_initial_timestamp_index is None:
															lowerCamelCase__       =       model_config.vocab_size






					def __call__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										# suppress <|notimestamps|> which is handled by without_timestamps
										lowerCamelCase__       =       scores.at[:, self.no_timestamps_token_id].set(-float("inf"	)	)

										def handle_pairs(_lowerCAmelCase	,  _lowerCAmelCase	):
															lowerCamelCase__       =       jnp.where((cur_len - self.begin_index) >= 1	,  _lowerCAmelCase	,  _lowerCAmelCase	)
															lowerCamelCase__       =       jnp.where(
															    input_ids_k[cur_len - 1] >= self.timestamp_begin	,  True and last_was_timestamp	,  _lowerCAmelCase	,  )

															lowerCamelCase__       =       jnp.where((cur_len - self.begin_index) < 2	,  _lowerCAmelCase	,  _lowerCAmelCase	)
															lowerCamelCase__       =       jnp.where(
															    input_ids_k[cur_len - 2] >= self.timestamp_begin	,  _lowerCAmelCase	,  _lowerCAmelCase	,  )

															return jnp.where(
															    _lowerCAmelCase	,  jnp.where(
															        penultimate_was_timestamp > 0	,  scores_k.at[self.timestamp_begin :].set(-float("inf"	)	)	,  scores_k.at[: self.eos_token_id].set(-float("inf"	)	)	,  )	,  _lowerCAmelCase	,  )

										lowerCamelCase__       =       jax.vmap(_lowerCAmelCase	)(_lowerCAmelCase	,  _lowerCAmelCase	)

										lowerCamelCase__       =       jnp.where(cur_len == self.begin_index	,  _lowerCAmelCase	,  _lowerCAmelCase	)
										lowerCamelCase__       =       jnp.where(
										    self.max_initial_timestamp_index is not None	,  True and apply_max_initial_timestamp	,  _lowerCAmelCase	,  )

										lowerCamelCase__       =       self.timestamp_begin + self.max_initial_timestamp_index

										lowerCamelCase__       =       jnp.where(
										    _lowerCAmelCase	,  scores.at[:, last_allowed + 1 :].set(-float("inf"	)	)	,  _lowerCAmelCase	,  )

										# if sum of probability over timestamps is above any other token, sample timestamp
										lowerCamelCase__       =       jax.nn.log_softmax(_lowerCAmelCase	,  axis=-1	)

										def handle_cumulative_probs(_lowerCAmelCase	,  _lowerCAmelCase	):
															lowerCamelCase__       =       jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]	,  axis=-1	)
															lowerCamelCase__       =       jnp.max(logprobs_k[: self.timestamp_begin]	)
															return jnp.where(
															    timestamp_logprob > max_text_token_logprob	,  scores_k.at[: self.timestamp_begin].set(-float("inf"	)	)	,  _lowerCAmelCase	,  )

										lowerCamelCase__       =       jax.vmap(_lowerCAmelCase	)(_lowerCAmelCase	,  _lowerCAmelCase	)

										return scores

