





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__       (    SCREAMING_SNAKE_CASE_   :      int				,	SCREAMING_SNAKE_CASE_   :      Optional[int]				,	SCREAMING_SNAKE_CASE_   :      Tuple=None						)  ->  Dict:






   """simple docstring"""







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

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

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

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

   return output
def      A__       (    SCREAMING_SNAKE_CASE_   :      Any				,	SCREAMING_SNAKE_CASE_   :      Optional[int]=None						)  ->  int:






   """simple docstring"""







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




@require_flax
class 			__UpperCamelCase    :
     __A	:  Optional[Any]							  =   None
     __A	:  Tuple							  =   ()





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _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(_UpperCamelCase					)
        _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 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _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(_UpperCamelCase					)

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

           _UpperCAmelCase           =   flax_model.generate(_UpperCamelCase					).sequences
           _UpperCAmelCase           =   pt_model.generate(torch.tensor(_UpperCamelCase				,							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 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase           =   self._get_input_ids_and_config()
        _UpperCAmelCase           =   False
        _UpperCAmelCase           =   max_length

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

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

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

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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase           =   self._get_input_ids_and_config()
        _UpperCAmelCase           =   True
        _UpperCAmelCase           =   max_length

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

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

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

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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _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(_UpperCamelCase					)

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

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

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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _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(_UpperCamelCase					)

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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _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(_UpperCamelCase					)

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

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

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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _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(_UpperCamelCase					)

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

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

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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _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(_UpperCamelCase					)

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

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

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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _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(_UpperCamelCase					)

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

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

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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _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(_UpperCamelCase					)

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

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

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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase       ,      _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(_UpperCamelCase					)

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

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

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




@require_flax
class 			__UpperCamelCase    (							unittest.TestCase    ):





     def 		UpperCamelCase(     self					):
        _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(_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'''					):
           _UpperCAmelCase           =   {'''foo''': '''bar'''}
           model.generate(_UpperCamelCase				,							**_UpperCamelCase					)






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


UpperCAmelCase_           =	{
    "configuration_xlm": ["XLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "XLMConfig", "XLMOnnxConfig"],
    "tokenization_xlm": ["XLMTokenizer"],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase_           =	[
       "XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
       "XLMForMultipleChoice",
       "XLMForQuestionAnswering",
       "XLMForQuestionAnsweringSimple",
       "XLMForSequenceClassification",
       "XLMForTokenClassification",
       "XLMModel",
       "XLMPreTrainedModel",
       "XLMWithLMHeadModel",
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase_           =	[
       "TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
       "TFXLMForMultipleChoice",
       "TFXLMForQuestionAnsweringSimple",
       "TFXLMForSequenceClassification",
       "TFXLMForTokenClassification",
       "TFXLMMainLayer",
       "TFXLMModel",
       "TFXLMPreTrainedModel",
       "TFXLMWithLMHeadModel",
   ]


if TYPE_CHECKING:
   from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMOnnxConfig
   from .tokenization_xlm import XLMTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_xlm import (
          XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
          XLMForMultipleChoice,
          XLMForQuestionAnswering,
          XLMForQuestionAnsweringSimple,
          XLMForSequenceClassification,
          XLMForTokenClassification,
          XLMModel,
          XLMPreTrainedModel,
          XLMWithLMHeadModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_xlm import (
          TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFXLMForMultipleChoice,
          TFXLMForQuestionAnsweringSimple,
          TFXLMForSequenceClassification,
          TFXLMForTokenClassification,
          TFXLMMainLayer,
          TFXLMModel,
          TFXLMPreTrainedModel,
          TFXLMWithLMHeadModel,
      )

else:
   import sys

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