





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

			_a 						=      "0.12"  # assumed parallelism: 8


if is_torch_available():
			import torch





def lowerCAmelCase__(__snake_case       ,__snake_case       ,__snake_case=None							)				->      Tuple:







	'''simple docstring'''







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

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

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

	lowerCamelCase__								=		np.array(__snake_case       ,dtype=jnp.intaa							).reshape(__snake_case							)

	return output





def lowerCAmelCase__(__snake_case       ,__snake_case=None							)				->      Any:







	'''simple docstring'''







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






@require_flax
class 					__A    :

			'''simple docstring'''







			lowerCAmelCase_							     =						None
			lowerCAmelCase_							     =						()





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__								=		self.model_tester.prepare_config_and_inputs_for_common()

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

				lowerCamelCase__								=		jnp.ones_like(__lowerCAmelCase							)
				lowerCamelCase__								=		attention_mask[:max_batch_size, :sequence_length]

				# generate max 5 tokens
				lowerCamelCase__								=		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()`
					lowerCamelCase__								=		config.eos_token_id
				return config, input_ids, attention_mask, max_length





			@is_pt_flax_cross_test
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()
				lowerCamelCase__								=		False
				lowerCamelCase__								=		max_length
				lowerCamelCase__								=		0

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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

					lowerCamelCase__								=		flax_model.generate(__lowerCAmelCase							).sequences
					lowerCamelCase__								=		pt_model.generate(torch.tensor(__lowerCAmelCase       ,  dtype=torch.long							)							)

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

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()
				lowerCamelCase__								=		False
				lowerCamelCase__								=		max_length

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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

					lowerCamelCase__								=		jit(model.generate							)
					lowerCamelCase__								=		jit_generate(__lowerCAmelCase							).sequences

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()
				lowerCamelCase__								=		True
				lowerCamelCase__								=		max_length

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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

					lowerCamelCase__								=		jit(model.generate							)
					lowerCamelCase__								=		jit_generate(__lowerCAmelCase							).sequences

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()
				lowerCamelCase__								=		False
				lowerCamelCase__								=		max_length
				lowerCamelCase__								=		2

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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

					lowerCamelCase__								=		jit(model.generate							)
					lowerCamelCase__								=		jit_generate(__lowerCAmelCase							).sequences

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()
				lowerCamelCase__								=		False
				lowerCamelCase__								=		max_length
				lowerCamelCase__								=		2
				lowerCamelCase__								=		2

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()
				lowerCamelCase__								=		True
				lowerCamelCase__								=		max_length
				lowerCamelCase__								=		0.8
				lowerCamelCase__								=		1_0
				lowerCamelCase__								=		0.3
				lowerCamelCase__								=		1
				lowerCamelCase__								=		8
				lowerCamelCase__								=		9

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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

					lowerCamelCase__								=		jit(model.generate							)
					lowerCamelCase__								=		jit_generate(__lowerCAmelCase							).sequences

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()
				lowerCamelCase__								=		max_length
				lowerCamelCase__								=		1
				lowerCamelCase__								=		8
				lowerCamelCase__								=		9

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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

					lowerCamelCase__								=		jit(model.generate							)
					lowerCamelCase__								=		jit_generate(__lowerCAmelCase							).sequences

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()
				lowerCamelCase__								=		max_length
				lowerCamelCase__								=		2
				lowerCamelCase__								=		1
				lowerCamelCase__								=		8
				lowerCamelCase__								=		9

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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

					lowerCamelCase__								=		jit(model.generate							)
					lowerCamelCase__								=		jit_generate(__lowerCAmelCase							).sequences

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()

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

				lowerCamelCase__								=		False
				lowerCamelCase__								=		max_length

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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

					lowerCamelCase__								=		jit(model.generate							)
					lowerCamelCase__								=		jit_generate(__lowerCAmelCase       ,  attention_mask=__lowerCAmelCase							).sequences

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()

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

				lowerCamelCase__								=		True
				lowerCamelCase__								=		max_length

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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

					lowerCamelCase__								=		jit(model.generate							)
					lowerCamelCase__								=		jit_generate(__lowerCAmelCase       ,  attention_mask=__lowerCAmelCase							).sequences

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__				,	lowerCamelCase__								=		self._get_input_ids_and_config()

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

				lowerCamelCase__								=		2
				lowerCamelCase__								=		max_length

				for model_class in self.all_generative_model_classes:
					lowerCamelCase__								=		model_class(__lowerCAmelCase							)

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

					lowerCamelCase__								=		jit(model.generate							)
					lowerCamelCase__								=		jit_generate(__lowerCAmelCase       ,  attention_mask=__lowerCAmelCase							).sequences

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






@require_flax
class 					__A    (    unittest.TestCase ):

			'''simple docstring'''







			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''							)
				lowerCamelCase__								=		FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''							)

				lowerCamelCase__								=		'''Hello world'''
				lowerCamelCase__								=		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'''							):
					lowerCamelCase__								=		{'''foo''': '''bar'''}
					model.generate(__lowerCAmelCase       ,  **__lowerCAmelCase							)










import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin






class 					__A    (    lowerCAmelCase ):

			'''simple docstring'''







			lowerCAmelCase_							     =						"""Speech2TextFeatureExtractor"""
			lowerCAmelCase_							     =						"""Speech2TextTokenizer"""





			def __init__(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase							):



				'''simple docstring'''







				super().__init__(__lowerCAmelCase       ,  __lowerCAmelCase							)
				lowerCamelCase__								=		self.feature_extractor
				lowerCamelCase__								=		False





			def __call__(    self       ,  *__lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				if self._in_target_context_manager:
					return self.current_processor(*__lowerCAmelCase       ,  **__lowerCAmelCase							)

				if "raw_speech" in kwargs:
					warnings.warn('''Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'''							)
					lowerCamelCase__								=		kwargs.pop('''raw_speech'''							)
				else:
					lowerCamelCase__								=		kwargs.pop('''audio'''       ,  __lowerCAmelCase							)
				lowerCamelCase__								=		kwargs.pop('''sampling_rate'''       ,  __lowerCAmelCase							)
				lowerCamelCase__								=		kwargs.pop('''text'''       ,  __lowerCAmelCase							)
				if len(__lowerCAmelCase							) > 0:
					lowerCamelCase__								=		args[0]
					lowerCamelCase__								=		args[1:]

				if audio is None and text is None:
					raise ValueError('''You need to specify either an `audio` or `text` input to process.'''							)

				if audio is not None:
					lowerCamelCase__								=		self.feature_extractor(__lowerCAmelCase       ,  *__lowerCAmelCase       ,  sampling_rate=__lowerCAmelCase       ,  **__lowerCAmelCase							)
				if text is not None:
					lowerCamelCase__								=		self.tokenizer(__lowerCAmelCase       ,  **__lowerCAmelCase							)

				if text is None:
					return inputs
				elif audio is None:
					return encodings
				else:
					lowerCamelCase__								=		encodings['''input_ids''']
					return inputs





			def 		__lowerCamelCase				(    self       ,  *__lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				return self.tokenizer.batch_decode(*__lowerCAmelCase       ,  **__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self       ,  *__lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				return self.tokenizer.decode(*__lowerCAmelCase       ,  **__lowerCAmelCase							)





			@contextmanager
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				warnings.warn(
				    '''`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '''
				    '''labels by using the argument `text` of the regular `__call__` method (either in the same call as '''
				    '''your audio inputs, or in a separate call.'''							)
				lowerCamelCase__								=		True
				lowerCamelCase__								=		self.tokenizer
				yield
				lowerCamelCase__								=		self.feature_extractor
				lowerCamelCase__								=		False



