




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

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

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

from ..wavaveca.test_feature_extraction_wavaveca import floats_list


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

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

if is_torch_available():
		from transformers import WavaVecaForCTC




@require_pyctcdecode
class     snake_case__							(unittest.TestCase   ):
				"""simple docstring"""



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

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

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





				def 							__UpperCAmelCase							(      self				:     str					,      **__lowerCamelCase				:     str							)					->     int:
							a    =			self.add_kwargs_tokens_map.copy()
							kwargs.update(__lowerCamelCase							)
							return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname					,      **__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     str					,      **__lowerCamelCase				:     Dict							)					->     int:
							return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname					,      **__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Union[str, Any]					,      **__lowerCamelCase				:     Optional[int]							)					->     Union[str, Any]:
							return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name					,      **__lowerCamelCase							)





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





				def 							__UpperCAmelCase							(      self				:     Union[str, Any]							)					->     List[str]:
							a    =			self.get_tokenizer()
							a    =			self.get_feature_extractor()
							a    =			self.get_decoder()

							a    =			WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,      feature_extractor=__lowerCamelCase					,      decoder=__lowerCamelCase							)

							processor.save_pretrained(self.tmpdirname							)
							a    =			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 							__UpperCAmelCase							(      self				:     Dict							)					->     Any:
							a    =			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
							a    =			WavaVecaProcessorWithLM.from_pretrained(
							    self.tmpdirname					,      alpha=5.0					,      beta=3.0					,      score_boundary=-7.0					,      unk_score_offset=3							)

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





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

							a    =			WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,      feature_extractor=__lowerCamelCase					,      decoder=__lowerCamelCase							)

							a    =			floats_list((3, 10_00)							)

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

							a    =			WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,      feature_extractor=__lowerCamelCase					,      decoder=__lowerCamelCase							)

							a    =			"This is a test string"

							a    =			processor(text=__lowerCamelCase							)

							a    =			tokenizer(__lowerCamelCase							)

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





				def 							__UpperCAmelCase							(      self				:     int					,      __lowerCamelCase				:     List[str]=(2, 10, 16)					,      __lowerCamelCase				:     Any=77							)					->     List[Any]:
							np.random.seed(__lowerCamelCase							)
							return np.random.rand(*__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Dict							)					->     int:
							a    =			self.get_feature_extractor()
							a    =			self.get_tokenizer()
							a    =			self.get_decoder()

							a    =			WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,      feature_extractor=__lowerCamelCase					,      decoder=__lowerCamelCase							)

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

							a    =			processor.decode(__lowerCamelCase							)

							a    =			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 							__UpperCAmelCase							(      self				:     Optional[Any]					,      __lowerCamelCase				:     Any							)					->     str:
							a    =			self.get_feature_extractor()
							a    =			self.get_tokenizer()
							a    =			self.get_decoder()

							a    =			WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,      feature_extractor=__lowerCamelCase					,      decoder=__lowerCamelCase							)

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

							a    =			list(__lowerCamelCase							)

							with get_context("fork"							).Pool() as p:
										a    =			decoder.decode_beams_batch(__lowerCamelCase					,      __lowerCamelCase							)

							a    ,		a    ,		a    =			[], [], []
							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 							__UpperCAmelCase							(      self				:     Dict							)					->     Any:
							a    =			self.get_feature_extractor()
							a    =			self.get_tokenizer()
							a    =			self.get_decoder()

							a    =			WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,      feature_extractor=__lowerCamelCase					,      decoder=__lowerCamelCase							)

							a    =			self._get_dummy_logits()

							a    =			15
							a    =			-20.0
							a    =			-4.0

							a    =			processor.batch_decode(
							    __lowerCamelCase					,      beam_width=__lowerCamelCase					,      beam_prune_logp=__lowerCamelCase					,      token_min_logp=__lowerCamelCase					,      )
							a    =			decoded_processor_out.text

							a    =			list(__lowerCamelCase							)

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

							a    =			[d[0][0] for d in decoded_decoder_out]
							a    =			[d[0][2] for d in decoded_decoder_out]
							a    =			[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 							__UpperCAmelCase							(      self				:     Optional[int]							)					->     Optional[Any]:
							a    =			self.get_feature_extractor()
							a    =			self.get_tokenizer()
							a    =			self.get_decoder()

							a    =			WavaVecaProcessorWithLM(tokenizer=__lowerCamelCase					,      feature_extractor=__lowerCamelCase					,      decoder=__lowerCamelCase							)

							a    =			self._get_dummy_logits()

							a    =			2.0
							a    =			5.0
							a    =			-20.0
							a    =			True

							a    =			processor.batch_decode(
							    __lowerCamelCase					,      alpha=__lowerCamelCase					,      beta=__lowerCamelCase					,      unk_score_offset=__lowerCamelCase					,      lm_score_boundary=__lowerCamelCase					,      )
							a    =			decoded_processor_out.text

							a    =			list(__lowerCamelCase							)
							decoder.reset_params(
							    alpha=__lowerCamelCase					,      beta=__lowerCamelCase					,      unk_score_offset=__lowerCamelCase					,      lm_score_boundary=__lowerCamelCase					,      )

							with get_context("fork"							).Pool() as pool:
										a    =			decoder.decode_beams_batch(
										    __lowerCamelCase					,      __lowerCamelCase					,      )

							a    =			[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							)
							a    =			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 							__UpperCAmelCase							(      self				:     Tuple							)					->     str:
							a    =			WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"							)

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

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

							a    =			WavaVecaProcessorWithLM.from_pretrained(__lowerCamelCase							)

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

							a    =			os.listdir(__lowerCamelCase							)
							a    =			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 							__UpperCAmelCase							(      self				:     int							)					->     Dict:
							a    =			WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"							)
							a    =			AutoProcessor.from_pretrained("hf-internal-testing/processor_with_lm"							)

							a    =			floats_list((3, 10_00)							)

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

							a    =			self._get_dummy_logits()

							a    =			processor_wavaveca.batch_decode(__lowerCamelCase							)
							a    =			processor_auto.batch_decode(__lowerCamelCase							)

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





				def 							__UpperCAmelCase							(      self				:     List[str]							)					->     Union[str, Any]:
							a    =			self.get_feature_extractor()
							a    =			self.get_tokenizer()
							a    =			self.get_decoder()

							a    =			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 							__UpperCAmelCase							(      __lowerCamelCase				:     Any					,      __lowerCamelCase				:     int							)					->     int:
							a    =			[d[key] for d in offsets]
							return retrieved_list





				def 							__UpperCAmelCase							(      self				:     List[Any]							)					->     int:
							a    =			WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"							)
							a    =			self._get_dummy_logits()[0]

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

							a    =			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 							__UpperCAmelCase							(      self				:     List[Any]							)					->     List[Any]:
							import torch

							a    =			load_dataset("common_voice"					,      "en"					,      split="train"					,      streaming=__lowerCamelCase							)
							a    =			ds.cast_column("audio"					,      datasets.Audio(sampling_rate=1_60_00							)							)
							a    =			iter(__lowerCamelCase							)
							a    =			next(__lowerCamelCase							)

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

							with torch.no_grad():
										a    =			model(__lowerCamelCase							).logits.cpu().numpy()

							a    =			processor.decode(logits[0]					,      output_word_offsets=__lowerCamelCase							)

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

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

							# fmt: off
							a    =			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]							)
							a    =			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 random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
		import os

		import jax.numpy as jnp
		from jax import jit

		from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
		from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

		__lowerCAmelCase		:				List[Any]       			=		'0.12'  # assumed parallelism: 8


if is_torch_available():
		import torch




def 			__magic_name__				( A		:					Dict,  A		:					Union[str, Any],  A		:					Optional[int]=None   ):






			'''simple docstring'''







			if rng is None:
						a    =			random.Random()

			a    =			1
			for dim in shape:
						total_dims *= dim

			a    =			[]
			for _ in range(A   ):
						values.append(rng.randint(0,  vocab_size - 1   )   )

			a    =			np.array(A,  dtype=jnp.intaa   ).reshape(A   )

			return output




def 			__magic_name__				( A		:					Dict,  A		:					Union[str, Any]=None   ):






			'''simple docstring'''







			a    =			ids_tensor(A,  vocab_size=2,  rng=A   )
			# make sure that at least one token is attended to for each batch
			a    =			1
			return attn_mask




@require_flax
class     snake_case__							:
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_       :							List[str]							 =					None
				SCREAMING_SNAKE_CASE_       :							Any							 =					()





				def 							__UpperCAmelCase							(      self				:     int							)					->     List[str]:
							a    ,		a    =			self.model_tester.prepare_config_and_inputs_for_common()

							# cut to half length & take max batch_size 3
							a    =			2
							a    =			inputs["input_ids"].shape[-1] // 2
							a    =			inputs["input_ids"][:max_batch_size, :sequence_length]

							a    =			jnp.ones_like(__lowerCamelCase							)
							a    =			attention_mask[:max_batch_size, :sequence_length]

							# generate max 5 tokens
							a    =			input_ids.shape[-1] + 5
							if config.eos_token_id is not None and config.pad_token_id is None:
										# hack to allow generate for models such as GPT2 as is done in `generate()`
										a    =			config.eos_token_id
							return config, input_ids, attention_mask, max_length





				@is_pt_flax_cross_test
				def 							__UpperCAmelCase							(      self				:     Optional[Any]							)					->     int:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()
							a    =			False
							a    =			max_length
							a    =			0

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model_class.__name__[4:]  # Skip the "Flax" at the beginning
										a    =			getattr(__lowerCamelCase					,      __lowerCamelCase							)
										a    =			pt_model_class(__lowerCamelCase							).eval()
										a    =			load_flax_weights_in_pytorch_model(__lowerCamelCase					,      flax_model.params							)

										a    =			flax_model.generate(__lowerCamelCase							).sequences
										a    =			pt_model.generate(torch.tensor(__lowerCamelCase					,      dtype=torch.long							)							)

										if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
													a    =			flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

										self.assertListEqual(pt_generation_outputs.numpy().tolist()					,      flax_generation_outputs.tolist()							)





				def 							__UpperCAmelCase							(      self				:     List[str]							)					->     Optional[int]:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()
							a    =			False
							a    =			max_length

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model.generate(__lowerCamelCase							).sequences
										self.assertEqual(generation_outputs.shape[-1]					,      __lowerCamelCase							)

										a    =			jit(model.generate							)
										a    =			jit_generate(__lowerCamelCase							).sequences

										self.assertListEqual(generation_outputs.tolist()					,      jit_generation_outputs.tolist()							)





				def 							__UpperCAmelCase							(      self				:     Optional[int]							)					->     Any:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()
							a    =			True
							a    =			max_length

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model.generate(__lowerCamelCase							).sequences
										self.assertEqual(generation_outputs.shape[-1]					,      __lowerCamelCase							)

										a    =			jit(model.generate							)
										a    =			jit_generate(__lowerCamelCase							).sequences

										self.assertListEqual(generation_outputs.tolist()					,      jit_generation_outputs.tolist()							)





				def 							__UpperCAmelCase							(      self				:     int							)					->     Dict:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()
							a    =			False
							a    =			max_length
							a    =			2

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model.generate(__lowerCamelCase							).sequences
										self.assertEqual(generation_outputs.shape[-1]					,      __lowerCamelCase							)

										a    =			jit(model.generate							)
										a    =			jit_generate(__lowerCamelCase							).sequences

										self.assertListEqual(generation_outputs.tolist()					,      jit_generation_outputs.tolist()							)





				def 							__UpperCAmelCase							(      self				:     Any							)					->     Union[str, Any]:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()
							a    =			False
							a    =			max_length
							a    =			2
							a    =			2

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model.generate(__lowerCamelCase							).sequences
										self.assertEqual(generation_outputs.shape[0]					,      input_ids.shape[0] * config.num_return_sequences							)





				def 							__UpperCAmelCase							(      self				:     Optional[Any]							)					->     Dict:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()
							a    =			True
							a    =			max_length
							a    =			0.8
							a    =			10
							a    =			0.3
							a    =			1
							a    =			8
							a    =			9

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model.generate(__lowerCamelCase							).sequences
										self.assertEqual(generation_outputs.shape[-1]					,      __lowerCamelCase							)

										a    =			jit(model.generate							)
										a    =			jit_generate(__lowerCamelCase							).sequences

										self.assertListEqual(generation_outputs.tolist()					,      jit_generation_outputs.tolist()							)





				def 							__UpperCAmelCase							(      self				:     Optional[Any]							)					->     Optional[Any]:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()
							a    =			max_length
							a    =			1
							a    =			8
							a    =			9

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model.generate(__lowerCamelCase							).sequences
										self.assertEqual(generation_outputs.shape[-1]					,      __lowerCamelCase							)

										a    =			jit(model.generate							)
										a    =			jit_generate(__lowerCamelCase							).sequences

										self.assertListEqual(generation_outputs.tolist()					,      jit_generation_outputs.tolist()							)





				def 							__UpperCAmelCase							(      self				:     List[str]							)					->     Union[str, Any]:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()
							a    =			max_length
							a    =			2
							a    =			1
							a    =			8
							a    =			9

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model.generate(__lowerCamelCase							).sequences
										self.assertEqual(generation_outputs.shape[-1]					,      __lowerCamelCase							)

										a    =			jit(model.generate							)
										a    =			jit_generate(__lowerCamelCase							).sequences

										self.assertListEqual(generation_outputs.tolist()					,      jit_generation_outputs.tolist()							)





				def 							__UpperCAmelCase							(      self				:     Union[str, Any]							)					->     Dict:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()

							# pad attention mask on the left
							a    =			attention_mask.at[(0, 0)].set(0							)

							a    =			False
							a    =			max_length

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model.generate(__lowerCamelCase					,      attention_mask=__lowerCamelCase							).sequences
										self.assertEqual(generation_outputs.shape[-1]					,      __lowerCamelCase							)

										a    =			jit(model.generate							)
										a    =			jit_generate(__lowerCamelCase					,      attention_mask=__lowerCamelCase							).sequences

										self.assertListEqual(generation_outputs.tolist()					,      jit_generation_outputs.tolist()							)





				def 							__UpperCAmelCase							(      self				:     Tuple							)					->     Tuple:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()

							# pad attention mask on the left
							a    =			attention_mask.at[(0, 0)].set(0							)

							a    =			True
							a    =			max_length

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model.generate(__lowerCamelCase					,      attention_mask=__lowerCamelCase							).sequences
										self.assertEqual(generation_outputs.shape[-1]					,      __lowerCamelCase							)

										a    =			jit(model.generate							)
										a    =			jit_generate(__lowerCamelCase					,      attention_mask=__lowerCamelCase							).sequences

										self.assertListEqual(generation_outputs.tolist()					,      jit_generation_outputs.tolist()							)





				def 							__UpperCAmelCase							(      self				:     Optional[int]							)					->     List[Any]:
							a    ,		a    ,		a    ,		a    =			self._get_input_ids_and_config()

							# pad attention mask on the left
							a    =			attention_mask.at[(0, 0)].set(0							)

							a    =			2
							a    =			max_length

							for model_class in self.all_generative_model_classes:
										a    =			model_class(__lowerCamelCase							)

										a    =			model.generate(__lowerCamelCase					,      attention_mask=__lowerCamelCase							).sequences
										self.assertEqual(generation_outputs.shape[-1]					,      __lowerCamelCase							)

										a    =			jit(model.generate							)
										a    =			jit_generate(__lowerCamelCase					,      attention_mask=__lowerCamelCase							).sequences

										self.assertListEqual(generation_outputs.tolist()					,      jit_generation_outputs.tolist()							)




@require_flax
class     snake_case__							(unittest.TestCase   ):
				"""simple docstring"""



				def 							__UpperCAmelCase							(      self				:     Dict							)					->     Optional[Any]:
							a    =			AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"							)
							a    =			FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"							)

							a    =			"Hello world"
							a    =			tokenizer(__lowerCamelCase					,      return_tensors="np"							).input_ids

							# typos are quickly detected (the correct argument is `do_sample`)
							with self.assertRaisesRegex(__lowerCamelCase					,      "do_samples"							):
										model.generate(__lowerCamelCase					,      do_samples=__lowerCamelCase							)

							# arbitrary arguments that will not be used anywhere are also not accepted
							with self.assertRaisesRegex(__lowerCamelCase					,      "foo"							):
										a    =			{"foo": "bar"}
										model.generate(__lowerCamelCase					,      **__lowerCamelCase							)






