
'''simple docstring'''

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

		UpperCAmelCase__ :			str											=		"0.12"  # assumed parallelism: 8


if is_torch_available():
		import torch







def 	A  (   UpperCamelCase_       :  Tuple    ,						UpperCamelCase_       :  Any    ,						UpperCamelCase_       :  Optional[int]=None   )       ->							Optional[int]:







			'''simple docstring'''



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

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

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

			lowerCAmelCase__										=				np.array(UpperCamelCase_    ,						dtype=jnp.intaa   ).reshape(UpperCamelCase_   )

			return output







def 	A  (   UpperCamelCase_       :  Tuple    ,						UpperCamelCase_       :  List[str]=None   )       ->							int:







			'''simple docstring'''



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






@require_flax
class  A  :
	snake_case__   :Dict   =  None
	snake_case__   :Optional[Any]   =  ()





	def 				__SCREAMING_SNAKE_CASE	(       self      :							Optional[Any]				):


				"""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(__magic_name__				)
				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 				__SCREAMING_SNAKE_CASE	(       self      :							Optional[Any]				):


				"""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(__magic_name__				)

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

							lowerCAmelCase__										=				flax_model.generate(__magic_name__				).sequences
							lowerCAmelCase__										=				pt_model.generate(torch.tensor(__magic_name__   ,				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 				__SCREAMING_SNAKE_CASE	(       self      :							Optional[int]				):


				"""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(__magic_name__				)

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

							lowerCAmelCase__										=				jit(model.generate				)
							lowerCAmelCase__										=				jit_generate(__magic_name__				).sequences

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





	def 				__SCREAMING_SNAKE_CASE	(       self      :							List[str]				):


				"""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(__magic_name__				)

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

							lowerCAmelCase__										=				jit(model.generate				)
							lowerCAmelCase__										=				jit_generate(__magic_name__				).sequences

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





	def 				__SCREAMING_SNAKE_CASE	(       self      :							Tuple				):


				"""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(__magic_name__				)

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

							lowerCAmelCase__										=				jit(model.generate				)
							lowerCAmelCase__										=				jit_generate(__magic_name__				).sequences

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





	def 				__SCREAMING_SNAKE_CASE	(       self      :							Optional[int]				):


				"""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(__magic_name__				)

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





	def 				__SCREAMING_SNAKE_CASE	(       self      :							str				):


				"""simple docstring"""

				lowerCAmelCase__						,lowerCAmelCase__						,lowerCAmelCase__						,lowerCAmelCase__										=				self._get_input_ids_and_config()
				lowerCAmelCase__										=				True
				lowerCAmelCase__										=				max_length
				lowerCAmelCase__										=				0.8
				lowerCAmelCase__										=				10
				lowerCAmelCase__										=				0.3
				lowerCAmelCase__										=				1
				lowerCAmelCase__										=				8
				lowerCAmelCase__										=				9

				for model_class in self.all_generative_model_classes:
							lowerCAmelCase__										=				model_class(__magic_name__				)

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

							lowerCAmelCase__										=				jit(model.generate				)
							lowerCAmelCase__										=				jit_generate(__magic_name__				).sequences

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





	def 				__SCREAMING_SNAKE_CASE	(       self      :							List[str]				):


				"""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(__magic_name__				)

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

							lowerCAmelCase__										=				jit(model.generate				)
							lowerCAmelCase__										=				jit_generate(__magic_name__				).sequences

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





	def 				__SCREAMING_SNAKE_CASE	(       self      :							int				):


				"""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(__magic_name__				)

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

							lowerCAmelCase__										=				jit(model.generate				)
							lowerCAmelCase__										=				jit_generate(__magic_name__				).sequences

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





	def 				__SCREAMING_SNAKE_CASE	(       self      :							Union[str, Any]				):


				"""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(__magic_name__				)

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

							lowerCAmelCase__										=				jit(model.generate				)
							lowerCAmelCase__										=				jit_generate(__magic_name__   ,				attention_mask=__magic_name__				).sequences

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





	def 				__SCREAMING_SNAKE_CASE	(       self      :							Optional[Any]				):


				"""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(__magic_name__				)

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

							lowerCAmelCase__										=				jit(model.generate				)
							lowerCAmelCase__										=				jit_generate(__magic_name__   ,				attention_mask=__magic_name__				).sequences

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







	def 				__SCREAMING_SNAKE_CASE	(       self      :							str				):


				"""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(__magic_name__				)

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

							lowerCAmelCase__										=				jit(model.generate				)
							lowerCAmelCase__										=				jit_generate(__magic_name__   ,				attention_mask=__magic_name__				).sequences

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






@require_flax
class  A  (						unittest.TestCase     ):





	def 				__SCREAMING_SNAKE_CASE	(       self      :							Optional[Any]				):


				"""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(__magic_name__   ,				return_tensors="np"				).input_ids

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

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






"""simple docstring"""



import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin
class 						__UpperCAmelCase    (unittest.TestCase							,   __A     ):







			'''simple docstring'''





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[Any]								=			load_tool("""text-classification"""							)
										self.tool.setup()
										A__       :     Any								=			load_tool("""text-classification"""   ,      remote=snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     int								=			self.tool("""That's quite cool"""   ,      ["""positive""", """negative"""]							)
										self.assertEqual(snake_case_   ,      """positive"""							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Optional[Any]								=			self.remote_tool("""That's quite cool"""   ,      ["""positive""", """negative"""]							)
										self.assertEqual(snake_case_   ,      """positive"""							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     str								=			self.tool(text="""That's quite cool"""   ,      labels=["""positive""", """negative"""]							)
										self.assertEqual(snake_case_   ,      """positive"""							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     int								=			self.remote_tool(text="""That's quite cool"""   ,      labels=["""positive""", """negative"""]							)
										self.assertEqual(snake_case_   ,      """positive"""							)





