




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


if is_torch_available():
     import torch



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase,			lowerCamelCase=None   ):
  if rng is None:
    lowercase    :Dict      	=						random.Random()

  lowercase    :Optional[Any]      	=						1
  for dim in shape:
    total_dims *= dim

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

  lowercase    :List[str]      	=						np.array(lowerCamelCase,			dtype=jnp.intaa   ).reshape(lowerCamelCase   )

  return output



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase=None   ):
  lowercase    :str      	=						ids_tensor(lowerCamelCase,			vocab_size=2,			rng=lowerCamelCase   )
  # make sure that at least one token is attended to for each batch
  lowercase    :Dict      	=						1
  return attn_mask




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






       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         lowercase    :Union[str, Any]      	=						self.model_tester.prepare_config_and_inputs_for_common()

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

         lowercase    :List[str]      	=						jnp.ones_like(_lowerCAmelCase				)
         lowercase    :Optional[Any]      	=						attention_mask[:max_batch_size, :sequence_length]

         # generate max 5 tokens
         lowercase    :Optional[int]      	=						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    :Tuple      	=						config.eos_token_id
         return config, input_ids, attention_mask, max_length






       @is_pt_flax_cross_test
       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         lowercase    :Optional[int]      	=						self._get_input_ids_and_config()
         lowercase    :Tuple      	=						False
         lowercase    :List[str]      	=						max_length
         lowercase    :int      	=						0

         for model_class in self.all_generative_model_classes:
           lowercase    :Optional[int]      	=						model_class(_lowerCAmelCase				)

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

           lowercase    :Optional[Any]      	=						flax_model.generate(_lowerCAmelCase				).sequences
           lowercase    :List[str]      	=						pt_model.generate(torch.tensor(_lowerCAmelCase       ,					dtype=torch.long				)				)

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

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






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Tuple      	=						self._get_input_ids_and_config()
         lowercase    :int      	=						False
         lowercase    :str      	=						max_length

         for model_class in self.all_generative_model_classes:
           lowercase    :List[str]      	=						model_class(_lowerCAmelCase				)

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

           lowercase    :Optional[int]      	=						jit(model.generate				)
           lowercase    :List[Any]      	=						jit_generate(_lowerCAmelCase				).sequences

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






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Optional[Any]      	=						self._get_input_ids_and_config()
         lowercase    :Any      	=						True
         lowercase    :Union[str, Any]      	=						max_length

         for model_class in self.all_generative_model_classes:
           lowercase    :Any      	=						model_class(_lowerCAmelCase				)

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

           lowercase    :Optional[Any]      	=						jit(model.generate				)
           lowercase    :Any      	=						jit_generate(_lowerCAmelCase				).sequences

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






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :int      	=						self._get_input_ids_and_config()
         lowercase    :Optional[Any]      	=						False
         lowercase    :Optional[int]      	=						max_length
         lowercase    :Optional[Any]      	=						2

         for model_class in self.all_generative_model_classes:
           lowercase    :int      	=						model_class(_lowerCAmelCase				)

           lowercase    :Optional[int]      	=						model.generate(_lowerCAmelCase				).sequences
           self.assertEqual(generation_outputs.shape[-1]       ,					_lowerCAmelCase				)

           lowercase    :Dict      	=						jit(model.generate				)
           lowercase    :Tuple      	=						jit_generate(_lowerCAmelCase				).sequences

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






       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         lowercase    :List[str]      	=						self._get_input_ids_and_config()
         lowercase    :Union[str, Any]      	=						False
         lowercase    :int      	=						max_length
         lowercase    :List[str]      	=						2
         lowercase    :Dict      	=						2

         for model_class in self.all_generative_model_classes:
           lowercase    :List[str]      	=						model_class(_lowerCAmelCase				)

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






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :Union[str, Any]      	=						self._get_input_ids_and_config()
         lowercase    :Union[str, Any]      	=						True
         lowercase    :int      	=						max_length
         lowercase    :Optional[int]      	=						0.8
         lowercase    :List[str]      	=						10
         lowercase    :Optional[Any]      	=						0.3
         lowercase    :int      	=						1
         lowercase    :int      	=						8
         lowercase    :List[str]      	=						9

         for model_class in self.all_generative_model_classes:
           lowercase    :List[Any]      	=						model_class(_lowerCAmelCase				)

           lowercase    :List[Any]      	=						model.generate(_lowerCAmelCase				).sequences
           self.assertEqual(generation_outputs.shape[-1]       ,					_lowerCAmelCase				)

           lowercase    :Tuple      	=						jit(model.generate				)
           lowercase    :int      	=						jit_generate(_lowerCAmelCase				).sequences

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






       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         lowercase    :Any      	=						self._get_input_ids_and_config()
         lowercase    :Any      	=						max_length
         lowercase    :Optional[int]      	=						1
         lowercase    :Dict      	=						8
         lowercase    :List[str]      	=						9

         for model_class in self.all_generative_model_classes:
           lowercase    :Dict      	=						model_class(_lowerCAmelCase				)

           lowercase    :Union[str, Any]      	=						model.generate(_lowerCAmelCase				).sequences
           self.assertEqual(generation_outputs.shape[-1]       ,					_lowerCAmelCase				)

           lowercase    :Optional[int]      	=						jit(model.generate				)
           lowercase    :List[Any]      	=						jit_generate(_lowerCAmelCase				).sequences

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






       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
         lowercase    :Dict      	=						self._get_input_ids_and_config()
         lowercase    :List[Any]      	=						max_length
         lowercase    :Optional[Any]      	=						2
         lowercase    :List[str]      	=						1
         lowercase    :Union[str, Any]      	=						8
         lowercase    :Any      	=						9

         for model_class in self.all_generative_model_classes:
           lowercase    :Tuple      	=						model_class(_lowerCAmelCase				)

           lowercase    :Optional[Any]      	=						model.generate(_lowerCAmelCase				).sequences
           self.assertEqual(generation_outputs.shape[-1]       ,					_lowerCAmelCase				)

           lowercase    :Union[str, Any]      	=						jit(model.generate				)
           lowercase    :Optional[Any]      	=						jit_generate(_lowerCAmelCase				).sequences

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






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         lowercase    :Dict      	=						self._get_input_ids_and_config()

         # pad attention mask on the left
         lowercase    :List[Any]      	=						attention_mask.at[(0, 0)].set(0				)

         lowercase    :List[Any]      	=						False
         lowercase    :Any      	=						max_length

         for model_class in self.all_generative_model_classes:
           lowercase    :Tuple      	=						model_class(_lowerCAmelCase				)

           lowercase    :Union[str, Any]      	=						model.generate(_lowerCAmelCase       ,					attention_mask=_lowerCAmelCase				).sequences
           self.assertEqual(generation_outputs.shape[-1]       ,					_lowerCAmelCase				)

           lowercase    :Any      	=						jit(model.generate				)
           lowercase    :int      	=						jit_generate(_lowerCAmelCase       ,					attention_mask=_lowerCAmelCase				).sequences

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






       def 					SCREAMING_SNAKE_CASE    (   self:					Dict				):
         lowercase    :List[Any]      	=						self._get_input_ids_and_config()

         # pad attention mask on the left
         lowercase    :List[str]      	=						attention_mask.at[(0, 0)].set(0				)

         lowercase    :Tuple      	=						True
         lowercase    :Tuple      	=						max_length

         for model_class in self.all_generative_model_classes:
           lowercase    :int      	=						model_class(_lowerCAmelCase				)

           lowercase    :Optional[int]      	=						model.generate(_lowerCAmelCase       ,					attention_mask=_lowerCAmelCase				).sequences
           self.assertEqual(generation_outputs.shape[-1]       ,					_lowerCAmelCase				)

           lowercase    :List[str]      	=						jit(model.generate				)
           lowercase    :Tuple      	=						jit_generate(_lowerCAmelCase       ,					attention_mask=_lowerCAmelCase				).sequences

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






       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         lowercase    :int      	=						self._get_input_ids_and_config()

         # pad attention mask on the left
         lowercase    :List[Any]      	=						attention_mask.at[(0, 0)].set(0				)

         lowercase    :List[str]      	=						2
         lowercase    :Union[str, Any]      	=						max_length

         for model_class in self.all_generative_model_classes:
           lowercase    :List[Any]      	=						model_class(_lowerCAmelCase				)

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

           lowercase    :Optional[int]      	=						jit(model.generate				)
           lowercase    :int      	=						jit_generate(_lowerCAmelCase       ,					attention_mask=_lowerCAmelCase				).sequences

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







@require_flax
class    __lowerCAmelCase       (							unittest.TestCase):






       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         lowercase    :Any      	=						AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"				)
         lowercase    :List[Any]      	=						FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"				)

         lowercase    :str      	=						"Hello world"
         lowercase    :Any      	=						tokenizer(_lowerCAmelCase       ,					return_tensors="np"				).input_ids

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

         # arbitrary arguments that will not be used anywhere are also not accepted
         with self.assertRaisesRegex(_lowerCAmelCase       ,					"foo"				):
           lowercase    :Union[str, Any]      	=						{"foo": "bar"}
           model.generate(_lowerCAmelCase       ,					**_lowerCAmelCase				)









_UpperCAmelCase	:						Union[str, Any]			=     "0.21.0"

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
     from .utils import rich



