"""simple docstring"""



import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
			import os

			from flax.core.frozen_dict import unfreeze
			from flax.traverse_util import flatten_dict

			from transformers import FlaxBertModel

			SCREAMING_SNAKE_CASE__:List[Any]			=	"""0.12"""  # assumed parallelism: 8






@require_flax
@is_staging_test
class 				snake_case__     (   unittest.TestCase   ):


	@classmethod
	def 		a__ (				cls		):
								__a     				=  TOKEN
								HfFolder.save_token(lowercase_		)


	@classmethod
	def 		a__ (				cls		):
								try:
															delete_repo(token=cls._token  ,						repo_id="test-model-flax"		)
								except HTTPError:
															pass

								try:
															delete_repo(token=cls._token  ,						repo_id="valid_org/test-model-flax-org"		)
								except HTTPError:
															pass


	def 		a__ (				self		):
								__a     				=  BertConfig(
								    vocab_size=99  ,						hidden_size=32  ,						num_hidden_layers=5  ,						num_attention_heads=4  ,						intermediate_size=37		)
								__a     				=  FlaxBertModel(lowercase_		)
								model.push_to_hub("test-model-flax"  ,						use_auth_token=self._token		)

								__a     				=  FlaxBertModel.from_pretrained(F"{USER}/test-model-flax"		)

								__a     				=  flatten_dict(unfreeze(model.params		)		)
								__a     				=  flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															__a     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(lowercase_  ,						1E-3  ,						msg=F"{key} not identical"		)

								# Reset repo
								delete_repo(token=self._token  ,						repo_id="test-model-flax"		)

								# Push to hub via save_pretrained
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(lowercase_  ,						repo_id="test-model-flax"  ,						push_to_hub=lowercase_  ,						use_auth_token=self._token		)

								__a     				=  FlaxBertModel.from_pretrained(F"{USER}/test-model-flax"		)

								__a     				=  flatten_dict(unfreeze(model.params		)		)
								__a     				=  flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															__a     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(lowercase_  ,						1E-3  ,						msg=F"{key} not identical"		)






	def 		a__ (				self		):
								__a     				=  BertConfig(
								    vocab_size=99  ,						hidden_size=32  ,						num_hidden_layers=5  ,						num_attention_heads=4  ,						intermediate_size=37		)
								__a     				=  FlaxBertModel(lowercase_		)
								model.push_to_hub("valid_org/test-model-flax-org"  ,						use_auth_token=self._token		)

								__a     				=  FlaxBertModel.from_pretrained("valid_org/test-model-flax-org"		)

								__a     				=  flatten_dict(unfreeze(model.params		)		)
								__a     				=  flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															__a     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(lowercase_  ,						1E-3  ,						msg=F"{key} not identical"		)

								# Reset repo
								delete_repo(token=self._token  ,						repo_id="valid_org/test-model-flax-org"		)

								# Push to hub via save_pretrained
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(
															    lowercase_  ,						repo_id="valid_org/test-model-flax-org"  ,						push_to_hub=lowercase_  ,						use_auth_token=self._token		)

								__a     				=  FlaxBertModel.from_pretrained("valid_org/test-model-flax-org"		)

								__a     				=  flatten_dict(unfreeze(model.params		)		)
								__a     				=  flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															__a     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(lowercase_  ,						1E-3  ,						msg=F"{key} not identical"		)



def 					_lowerCamelCase(					a ,  a							):
							__a     				=  True
							__a     				=  flatten_dict(modela.params							)
							__a     				=  flatten_dict(modela.params							)
							for key in flat_params_a.keys():
														if np.sum(np.abs(flat_params_a[key] - flat_params_a[key]							)							) > 1E-4:
																					__a     				=  False

							return models_are_equal






@require_flax
class 				snake_case__     (   unittest.TestCase   ):


	def 		a__ (				self		):
								__a     				=  BertConfig.from_pretrained("hf-internal-testing/tiny-bert-flax-only"		)
								__a     				=  FlaxBertModel(lowercase_		)

								__a     				=  "bert"
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(lowercase_  ,						lowercase_		)		)

															with self.assertRaises(lowercase_		):
																						__a     				=  FlaxBertModel.from_pretrained(lowercase_		)

															__a     				=  FlaxBertModel.from_pretrained(lowercase_  ,						subfolder=lowercase_		)

								self.assertTrue(check_models_equal(lowercase_  ,						lowercase_		)		)


	def 		a__ (				self		):
								__a     				=  BertConfig.from_pretrained("hf-internal-testing/tiny-bert-flax-only"		)
								__a     				=  FlaxBertModel(lowercase_		)

								__a     				=  "bert"
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(lowercase_  ,						lowercase_		)  ,						max_shard_size="10KB"		)

															with self.assertRaises(lowercase_		):
																						__a     				=  FlaxBertModel.from_pretrained(lowercase_		)

															__a     				=  FlaxBertModel.from_pretrained(lowercase_  ,						subfolder=lowercase_		)

								self.assertTrue(check_models_equal(lowercase_  ,						lowercase_		)		)


	def 		a__ (				self		):
								__a     				=  "bert"
								__a     				=  "hf-internal-testing/tiny-random-bert-subfolder"

								with self.assertRaises(lowercase_		):
															__a     				=  FlaxBertModel.from_pretrained(lowercase_		)

								__a     				=  FlaxBertModel.from_pretrained(lowercase_  ,						subfolder=lowercase_		)

								self.assertIsNotNone(lowercase_		)






	def 		a__ (				self		):
								__a     				=  "bert"
								__a     				=  "hf-internal-testing/tiny-random-bert-sharded-subfolder"
								with self.assertRaises(lowercase_		):
															__a     				=  FlaxBertModel.from_pretrained(lowercase_		)

								__a     				=  FlaxBertModel.from_pretrained(lowercase_  ,						subfolder=lowercase_		)

								self.assertIsNotNone(lowercase_		)

"""simple docstring"""



from ...processing_utils import ProcessorMixin






class 				snake_case__     (   snake_case_   ):
 _snake_case   :			List[str]           =		"""WhisperFeatureExtractor"""
 _snake_case   :			Any           =		"""WhisperTokenizer"""


 def __init__(				self  ,						lowerCamelCase  ,						lowerCamelCase		):
        super().__init__(lowerCamelCase  ,						lowerCamelCase		)
        __a     				=  self.feature_extractor
        __a     				=  False


 def 		a__ (				self  ,						lowerCamelCase=None  ,						lowerCamelCase=None  ,						lowerCamelCase=True		):
        return self.tokenizer.get_decoder_prompt_ids(task=lowerCamelCase  ,						language=lowerCamelCase  ,						no_timestamps=lowerCamelCase		)


 def __call__(				self  ,						*lowerCamelCase  ,						**lowerCamelCase		):
        # For backward compatibility
        if self._in_target_context_manager:
               return self.current_processor(*lowerCamelCase  ,						**lowerCamelCase		)

        __a     				=  kwargs.pop("audio"  ,						lowerCamelCase		)
        __a     				=  kwargs.pop("sampling_rate"  ,						lowerCamelCase		)
        __a     				=  kwargs.pop("text"  ,						lowerCamelCase		)
        if len(lowerCamelCase		) > 0:
               __a     				=  args[0]
               __a     				=  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:
               __a     				=  self.feature_extractor(lowerCamelCase  ,						*lowerCamelCase  ,						sampling_rate=lowerCamelCase  ,						**lowerCamelCase		)
        if text is not None:
               __a     				=  self.tokenizer(lowerCamelCase  ,						**lowerCamelCase		)

        if text is None:
               return inputs

        elif audio is None:
               return encodings
        else:
               __a     				=  encodings["input_ids"]
               return inputs


 def 		a__ (				self  ,						*lowerCamelCase  ,						**lowerCamelCase		):
        return self.tokenizer.batch_decode(*lowerCamelCase  ,						**lowerCamelCase		)


 def 		a__ (				self  ,						*lowerCamelCase  ,						**lowerCamelCase		):
        return self.tokenizer.decode(*lowerCamelCase  ,						**lowerCamelCase		)






 def 		a__ (				self  ,						lowerCamelCase  ,						lowerCamelCase="np"		):
        return self.tokenizer.get_prompt_ids(lowerCamelCase  ,						return_tensors=lowerCamelCase		)
