



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__       =	'0.12'  # assumed parallelism: 8


if is_torch_available():
							import torch





def 							A     (					_UpperCAmelCase    :       Tuple							,       _UpperCAmelCase    :       str							,       _UpperCAmelCase    :       Any=None   )    ->			str:







			'''simple docstring'''
			if rng is None:
						_UpperCAmelCase			=       random.Random()

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

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

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

			return output





def 							A     (					_UpperCAmelCase    :       str							,       _UpperCAmelCase    :       Tuple=None   )    ->			Optional[Any]:







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







@require_flax
class       __lowerCAmelCase   :
			UpperCamelCase        =					None
			UpperCamelCase        =					()



			def    _lowerCamelCase	(						self       :  int)       ->   Tuple:



						"""simple docstring"""






						_UpperCAmelCase			=       self.model_tester.prepare_config_and_inputs_for_common()

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

						_UpperCAmelCase			=       jnp.ones_like(_SCREAMING_SNAKE_CASE)
						_UpperCAmelCase			=       attention_mask[:max_batch_size, :sequence_length]

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



			@is_pt_flax_cross_test
			def    _lowerCamelCase	(						self       :  int)       ->   str:



						"""simple docstring"""






						_UpperCAmelCase			=       self._get_input_ids_and_config()
						_UpperCAmelCase			=       False
						_UpperCAmelCase			=       max_length
						_UpperCAmelCase			=       0

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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

									_UpperCAmelCase			=       flax_model.generate(_SCREAMING_SNAKE_CASE).sequences
									_UpperCAmelCase			=       pt_model.generate(torch.tensor(_SCREAMING_SNAKE_CASE			,		dtype=torch.long))

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

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



			def    _lowerCamelCase	(						self       :  List[str])       ->   List[str]:



						"""simple docstring"""






						_UpperCAmelCase			=       self._get_input_ids_and_config()
						_UpperCAmelCase			=       False
						_UpperCAmelCase			=       max_length

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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

									_UpperCAmelCase			=       jit(model.generate)
									_UpperCAmelCase			=       jit_generate(_SCREAMING_SNAKE_CASE).sequences

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



			def    _lowerCamelCase	(						self       :  Dict)       ->   Union[str, Any]:



						"""simple docstring"""






						_UpperCAmelCase			=       self._get_input_ids_and_config()
						_UpperCAmelCase			=       True
						_UpperCAmelCase			=       max_length

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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

									_UpperCAmelCase			=       jit(model.generate)
									_UpperCAmelCase			=       jit_generate(_SCREAMING_SNAKE_CASE).sequences

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



			def    _lowerCamelCase	(						self       :  Optional[int])       ->   Dict:



						"""simple docstring"""






						_UpperCAmelCase			=       self._get_input_ids_and_config()
						_UpperCAmelCase			=       False
						_UpperCAmelCase			=       max_length
						_UpperCAmelCase			=       2

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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

									_UpperCAmelCase			=       jit(model.generate)
									_UpperCAmelCase			=       jit_generate(_SCREAMING_SNAKE_CASE).sequences

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



			def    _lowerCamelCase	(						self       :  str)       ->   str:



						"""simple docstring"""






						_UpperCAmelCase			=       self._get_input_ids_and_config()
						_UpperCAmelCase			=       False
						_UpperCAmelCase			=       max_length
						_UpperCAmelCase			=       2
						_UpperCAmelCase			=       2

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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



			def    _lowerCamelCase	(						self       :  Any)       ->   Optional[Any]:



						"""simple docstring"""






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

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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

									_UpperCAmelCase			=       jit(model.generate)
									_UpperCAmelCase			=       jit_generate(_SCREAMING_SNAKE_CASE).sequences

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



			def    _lowerCamelCase	(						self       :  Tuple)       ->   List[Any]:



						"""simple docstring"""






						_UpperCAmelCase			=       self._get_input_ids_and_config()
						_UpperCAmelCase			=       max_length
						_UpperCAmelCase			=       1
						_UpperCAmelCase			=       8
						_UpperCAmelCase			=       9

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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

									_UpperCAmelCase			=       jit(model.generate)
									_UpperCAmelCase			=       jit_generate(_SCREAMING_SNAKE_CASE).sequences

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



			def    _lowerCamelCase	(						self       :  List[str])       ->   Dict:



						"""simple docstring"""






						_UpperCAmelCase			=       self._get_input_ids_and_config()
						_UpperCAmelCase			=       max_length
						_UpperCAmelCase			=       2
						_UpperCAmelCase			=       1
						_UpperCAmelCase			=       8
						_UpperCAmelCase			=       9

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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

									_UpperCAmelCase			=       jit(model.generate)
									_UpperCAmelCase			=       jit_generate(_SCREAMING_SNAKE_CASE).sequences

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



			def    _lowerCamelCase	(						self       :  str)       ->   Any:



						"""simple docstring"""






						_UpperCAmelCase			=       self._get_input_ids_and_config()

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

						_UpperCAmelCase			=       False
						_UpperCAmelCase			=       max_length

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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

									_UpperCAmelCase			=       jit(model.generate)
									_UpperCAmelCase			=       jit_generate(_SCREAMING_SNAKE_CASE			,		attention_mask=_SCREAMING_SNAKE_CASE).sequences

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



			def    _lowerCamelCase	(						self       :  Tuple)       ->   str:



						"""simple docstring"""






						_UpperCAmelCase			=       self._get_input_ids_and_config()

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

						_UpperCAmelCase			=       True
						_UpperCAmelCase			=       max_length

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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

									_UpperCAmelCase			=       jit(model.generate)
									_UpperCAmelCase			=       jit_generate(_SCREAMING_SNAKE_CASE			,		attention_mask=_SCREAMING_SNAKE_CASE).sequences

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



			def    _lowerCamelCase	(						self       :  Union[str, Any])       ->   str:



						"""simple docstring"""






						_UpperCAmelCase			=       self._get_input_ids_and_config()

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

						_UpperCAmelCase			=       2
						_UpperCAmelCase			=       max_length

						for model_class in self.all_generative_model_classes:
									_UpperCAmelCase			=       model_class(_SCREAMING_SNAKE_CASE)

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

									_UpperCAmelCase			=       jit(model.generate)
									_UpperCAmelCase			=       jit_generate(_SCREAMING_SNAKE_CASE			,		attention_mask=_SCREAMING_SNAKE_CASE).sequences

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







@require_flax
class       __lowerCAmelCase   (      unittest.TestCase						):



			def    _lowerCamelCase	(						self       :  Tuple)       ->   Any:



						"""simple docstring"""






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

						_UpperCAmelCase			=       "Hello world"
						_UpperCAmelCase			=       tokenizer(_SCREAMING_SNAKE_CASE			,		return_tensors='np').input_ids

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

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




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase__       =	{
    "configuration_lxmert": ["LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "LxmertConfig"],
    "tokenization_lxmert": ["LxmertTokenizer"],
}

try:
							if not is_tokenizers_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__       =	["LxmertTokenizerFast"]

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__       =	[
							    "LxmertEncoder",
							    "LxmertForPreTraining",
							    "LxmertForQuestionAnswering",
							    "LxmertModel",
							    "LxmertPreTrainedModel",
							    "LxmertVisualFeatureEncoder",
							    "LxmertXLayer",
							]

try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__       =	[
							    "TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "TFLxmertForPreTraining",
							    "TFLxmertMainLayer",
							    "TFLxmertModel",
							    "TFLxmertPreTrainedModel",
							    "TFLxmertVisualFeatureEncoder",
							]


if TYPE_CHECKING:
							from .configuration_lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig
							from .tokenization_lxmert import LxmertTokenizer

							try:
														if not is_tokenizers_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_lxmert_fast import LxmertTokenizerFast

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_lxmert import (
														    LxmertEncoder,
														    LxmertForPreTraining,
														    LxmertForQuestionAnswering,
														    LxmertModel,
														    LxmertPreTrainedModel,
														    LxmertVisualFeatureEncoder,
														    LxmertXLayer,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_lxmert import (
														    TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST,
														    TFLxmertForPreTraining,
														    TFLxmertMainLayer,
														    TFLxmertModel,
														    TFLxmertPreTrainedModel,
														    TFLxmertVisualFeatureEncoder,
														)

else:
							import sys

							UpperCAmelCase__       =	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
