"""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

	lowerCAmelCase_:					Optional[Any]     =					"0.12"  # assumed parallelism: 8


if is_torch_available():
	import torch







def 					__a   (		A    ,				A    ,				A=None      ):






					'''simple docstring'''



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

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

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

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

					return output







def 					__a   (		A    ,				A=None      ):






					'''simple docstring'''



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






@require_flax
class    a__    :
			snake_case_							=      None
			snake_case_							=      ()


			def 	snake_case__	(     self  ):
								'''simple docstring'''



								lowercase__ ,						lowercase__				 =							self.model_tester.prepare_config_and_inputs_for_common()

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

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

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


			@is_pt_flax_cross_test
			def 	snake_case__	(     self  ):
								'''simple docstring'''



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

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

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

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

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

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


			def 	snake_case__	(     self  ):
								'''simple docstring'''



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

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

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

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

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


			def 	snake_case__	(     self  ):
								'''simple docstring'''



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

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

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

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

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


			def 	snake_case__	(     self  ):
								'''simple docstring'''



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

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

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

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

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


			def 	snake_case__	(     self  ):
								'''simple docstring'''



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

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

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


			def 	snake_case__	(     self  ):
								'''simple docstring'''



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

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

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

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

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


			def 	snake_case__	(     self  ):
								'''simple docstring'''



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

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

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

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

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


			def 	snake_case__	(     self  ):
								'''simple docstring'''



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

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

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

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

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


			def 	snake_case__	(     self  ):
								'''simple docstring'''



								lowercase__ ,						lowercase__ ,						lowercase__ ,						lowercase__				 =							self._get_input_ids_and_config()

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

								lowercase__				 =							False
								lowercase__				 =							max_length

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

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

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

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


			def 	snake_case__	(     self  ):
								'''simple docstring'''



								lowercase__ ,						lowercase__ ,						lowercase__ ,						lowercase__				 =							self._get_input_ids_and_config()

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

								lowercase__				 =							True
								lowercase__				 =							max_length

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

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

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

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


			def 	snake_case__	(     self  ):
								'''simple docstring'''



								lowercase__ ,						lowercase__ ,						lowercase__ ,						lowercase__				 =							self._get_input_ids_and_config()

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

								lowercase__				 =							2
								lowercase__				 =							max_length

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

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

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

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






@require_flax
class    a__    (   unittest.TestCase	):


			def 	snake_case__	(     self  ):
								'''simple docstring'''



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

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

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

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





"""simple docstring"""


from typing import TYPE_CHECKING

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


lowerCAmelCase_:					Union[str, Any]     =					{
    "configuration_distilbert": [
        "DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "DistilBertConfig",
        "DistilBertOnnxConfig",
    ],
    "tokenization_distilbert": ["DistilBertTokenizer"],
}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase_:					Union[str, Any]     =					["DistilBertTokenizerFast"]

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase_:					Any     =					[
     "DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
     "DistilBertForMaskedLM",
     "DistilBertForMultipleChoice",
     "DistilBertForQuestionAnswering",
     "DistilBertForSequenceClassification",
     "DistilBertForTokenClassification",
     "DistilBertModel",
     "DistilBertPreTrainedModel",
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase_:					Tuple     =					[
     "TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
     "TFDistilBertForMaskedLM",
     "TFDistilBertForMultipleChoice",
     "TFDistilBertForQuestionAnswering",
     "TFDistilBertForSequenceClassification",
     "TFDistilBertForTokenClassification",
     "TFDistilBertMainLayer",
     "TFDistilBertModel",
     "TFDistilBertPreTrainedModel",
 ]

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase_:					Optional[Any]     =					[
     "FlaxDistilBertForMaskedLM",
     "FlaxDistilBertForMultipleChoice",
     "FlaxDistilBertForQuestionAnswering",
     "FlaxDistilBertForSequenceClassification",
     "FlaxDistilBertForTokenClassification",
     "FlaxDistilBertModel",
     "FlaxDistilBertPreTrainedModel",
 ]


if TYPE_CHECKING:
 from .configuration_distilbert import (
     DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
     DistilBertConfig,
     DistilBertOnnxConfig,
 )
 from .tokenization_distilbert import DistilBertTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_distilbert_fast import DistilBertTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_distilbert import (
      DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      DistilBertForMaskedLM,
      DistilBertForMultipleChoice,
      DistilBertForQuestionAnswering,
      DistilBertForSequenceClassification,
      DistilBertForTokenClassification,
      DistilBertModel,
      DistilBertPreTrainedModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_distilbert import (
      TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFDistilBertForMaskedLM,
      TFDistilBertForMultipleChoice,
      TFDistilBertForQuestionAnswering,
      TFDistilBertForSequenceClassification,
      TFDistilBertForTokenClassification,
      TFDistilBertMainLayer,
      TFDistilBertModel,
      TFDistilBertPreTrainedModel,
  )

 try:
  if not is_flax_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_flax_distilbert import (
      FlaxDistilBertForMaskedLM,
      FlaxDistilBertForMultipleChoice,
      FlaxDistilBertForQuestionAnswering,
      FlaxDistilBertForSequenceClassification,
      FlaxDistilBertForTokenClassification,
      FlaxDistilBertModel,
      FlaxDistilBertPreTrainedModel,
  )

else:
 import sys

 lowerCAmelCase_:					Optional[int]     =					_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




