

import unittest

import numpy as np

from transformers import DistilBertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  import jax.numpy as jnp

  from transformers.models.distilbert.modeling_flax_distilbert import (
      FlaxDistilBertForMaskedLM,
      FlaxDistilBertForMultipleChoice,
      FlaxDistilBertForQuestionAnswering,
      FlaxDistilBertForSequenceClassification,
      FlaxDistilBertForTokenClassification,
      FlaxDistilBertModel,
  )





class  __lowerCAmelCase     (					unittest.TestCase							):




  def __init__(				self							,   lowerCAmelCase							,   lowerCAmelCase=13							,   lowerCAmelCase=7							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   lowerCAmelCase=99							,   lowerCAmelCase=32							,   lowerCAmelCase=5							,   lowerCAmelCase=4							,   lowerCAmelCase=37							,   lowerCAmelCase="gelu"							,   lowerCAmelCase=0.1							,   lowerCAmelCase=0.1							,   lowerCAmelCase=512							,   lowerCAmelCase=16							,   lowerCAmelCase=2							,   lowerCAmelCase=0.02							,   lowerCAmelCase=4							,   )  -> Union[str, Any]:
        '''simple docstring'''






        _lowercase	       =parent
        _lowercase	       =batch_size
        _lowercase	       =seq_length
        _lowercase	       =is_training
        _lowercase	       =use_attention_mask
        _lowercase	       =use_token_type_ids
        _lowercase	       =use_labels
        _lowercase	       =vocab_size
        _lowercase	       =hidden_size
        _lowercase	       =num_hidden_layers
        _lowercase	       =num_attention_heads
        _lowercase	       =intermediate_size
        _lowercase	       =hidden_act
        _lowercase	       =hidden_dropout_prob
        _lowercase	       =attention_probs_dropout_prob
        _lowercase	       =max_position_embeddings
        _lowercase	       =type_vocab_size
        _lowercase	       =type_sequence_label_size
        _lowercase	       =initializer_range
        _lowercase	       =num_choices




  def   A__     (				self					)  -> List[Any]:
        '''simple docstring'''






        _lowercase	       =ids_tensor([self.batch_size, self.seq_length]							,   self.vocab_size					)

        _lowercase	       =None
        if self.use_attention_mask:
              _lowercase	       =random_attention_mask([self.batch_size, self.seq_length]					)

        _lowercase	       =DistilBertConfig(
            vocab_size=self.vocab_size							,   dim=self.hidden_size							,   n_layers=self.num_hidden_layers							,   n_heads=self.num_attention_heads							,   hidden_dim=self.intermediate_size							,   hidden_act=self.hidden_act							,   dropout=self.hidden_dropout_prob							,   attention_dropout=self.attention_probs_dropout_prob							,   max_position_embeddings=self.max_position_embeddings							,   initializer_range=self.initializer_range							,   tie_weights_=lowerCAmelCase							,   )

        return config, input_ids, attention_mask




  def   A__     (				self					)  -> Any:
        '''simple docstring'''






        _lowercase	       =self.prepare_config_and_inputs()
        _lowercase       ,       _lowercase       ,       _lowercase	       =config_and_inputs
        _lowercase	       ={'input_ids': input_ids, 'attention_mask': attention_mask}
        return config, inputs_dict





@require_flax
class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE			,			unittest.TestCase							):
  _a									=  (
      (
          FlaxDistilBertModel,
          FlaxDistilBertForMaskedLM,
          FlaxDistilBertForMultipleChoice,
          FlaxDistilBertForQuestionAnswering,
          FlaxDistilBertForSequenceClassification,
          FlaxDistilBertForTokenClassification,
          FlaxDistilBertForQuestionAnswering,
      )
      if is_flax_available()
      else ()
  )




  def   A__     (				self					)  -> Tuple:
        '''simple docstring'''






        _lowercase	       =FlaxDistilBertModelTester(self					)




  @slow
  def   A__     (				self					)  -> Union[str, Any]:
        '''simple docstring'''






        for model_class_name in self.all_model_classes:
              _lowercase	       =model_class_name.from_pretrained('distilbert-base-uncased'					)
              _lowercase	       =model(np.ones((1, 1)					)					)
              self.assertIsNotNone(lowerCAmelCase					)





@require_flax
class  __lowerCAmelCase     (					unittest.TestCase							):




  @slow
  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        _lowercase	       =FlaxDistilBertModel.from_pretrained('distilbert-base-uncased'					)
        _lowercase	       =np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]					)
        _lowercase	       =np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]					)
        _lowercase	       =model(lowerCAmelCase							,   attention_mask=lowerCAmelCase					)[0]
        _lowercase	       =(1, 11, 768)
        self.assertEqual(output.shape							,   lowerCAmelCase					)
        _lowercase	       =np.array([[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]					)

        self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]							,   lowerCAmelCase							,   atol=1e-4					)					)





from .imports import is_tqdm_available


if is_tqdm_available():
  from tqdm.auto import tqdm as _tqdm

from ..state import PartialState
def      a      (			A__	:   bool = True      ,   *A__	:   int      ,   **A__	:   Union[str, Any]     )    ->      List[str]:





      """simple docstring"""



      if not is_tqdm_available():
            raise ImportError('Accelerate\'s `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`.'     )
      _lowercase	       =False
      if main_process_only:
            _lowercase	       =PartialState().local_process_index == 0
      return _tqdm(*A__      ,   **A__      ,   disable=A__     )


