





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


if is_torch_available():
      import torch
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase=None  ):





      """simple docstring"""


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

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

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

      _lowerCAmelCase												=							np.array(lowerCAmelCase_							,     dtype=jnp.intaa  ).reshape(lowerCAmelCase_  )

      return output
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase=None  ):





      """simple docstring"""


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




@require_flax
class 			UpperCAmelCase					:
      _lowercase:      Any											=       None
      _lowercase:      List[Any]											=       ()




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				Dict:
            _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(_UpperCAmelCase  )
            _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     lowercase__   (      self      :			Optional[int]  )    ->				Dict:
            _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(_UpperCAmelCase  )

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

                  _lowerCAmelCase												=							flax_model.generate(_UpperCAmelCase  ).sequences
                  _lowerCAmelCase												=							pt_model.generate(torch.tensor(_UpperCAmelCase      ,      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     lowercase__   (      self      :			Any  )    ->				Optional[int]:
            _lowerCAmelCase												=							self._get_input_ids_and_config()
            _lowerCAmelCase												=							False
            _lowerCAmelCase												=							max_length

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

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

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

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




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				str:
            _lowerCAmelCase												=							self._get_input_ids_and_config()
            _lowerCAmelCase												=							True
            _lowerCAmelCase												=							max_length

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

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

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

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




      def     lowercase__   (      self      :			int  )    ->				Union[str, Any]:
            _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(_UpperCAmelCase  )

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

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

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




      def     lowercase__   (      self      :			Dict  )    ->				Optional[Any]:
            _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(_UpperCAmelCase  )

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




      def     lowercase__   (      self      :			Optional[int]  )    ->				List[Any]:
            _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(_UpperCAmelCase  )

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

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

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




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				Tuple:
            _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(_UpperCAmelCase  )

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

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

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




      def     lowercase__   (      self      :			Tuple  )    ->				Dict:
            _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(_UpperCAmelCase  )

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

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

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




      def     lowercase__   (      self      :			Dict  )    ->				str:
            _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(_UpperCAmelCase  )

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

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

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




      def     lowercase__   (      self      :			Tuple  )    ->				List[Any]:
            _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(_UpperCAmelCase  )

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

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

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




      def     lowercase__   (      self      :			Optional[int]  )    ->				Any:
            _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(_UpperCAmelCase  )

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

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

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




@require_flax
class 			UpperCAmelCase					(      unittest.TestCase       ):




      def     lowercase__   (      self      :			int  )    ->				Optional[int]:
            _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(_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"""  ):
                  _lowerCAmelCase												=							{'foo': 'bar'}
                  model.generate(_UpperCAmelCase      ,      **_UpperCAmelCase  )













'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available


A__    :       int              ={'''tokenization_herbert''': ['''HerbertTokenizer''']}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A__    :       Dict              =['''HerbertTokenizerFast''']


if TYPE_CHECKING:
      from .tokenization_herbert import HerbertTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_herbert_fast import HerbertTokenizerFast

else:
      import sys

      A__    :       List[Any]              =_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






