


'''simple docstring'''

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     ):





      '''simple docstring'''






      def __init__(self    :    Tuple   ,						UpperCamelCase    :    int   ,						UpperCamelCase    :    str=13   ,						UpperCamelCase    :    Optional[int]=7   ,						UpperCamelCase    :    Any=True   ,						UpperCamelCase    :    int=True   ,						UpperCamelCase    :    Optional[int]=True   ,						UpperCamelCase    :    Optional[Any]=True   ,						UpperCamelCase    :    Optional[Any]=99   ,						UpperCamelCase    :    Any=32   ,						UpperCamelCase    :    Tuple=5   ,						UpperCamelCase    :    str=4   ,						UpperCamelCase    :    Optional[int]=37   ,						UpperCamelCase    :    Any="gelu"   ,						UpperCamelCase    :    Tuple=0.1   ,						UpperCamelCase    :    Dict=0.1   ,						UpperCamelCase    :    List[Any]=512   ,						UpperCamelCase    :    List[Any]=16   ,						UpperCamelCase    :    List[Any]=2   ,						UpperCamelCase    :    Dict=0.02   ,						UpperCamelCase    :    Dict=4   ,						):

             '''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       UpperCamelCase__		(self    :    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_=UpperCamelCase   ,						)

             return config, input_ids, attention_mask






      def       UpperCamelCase__		(self    :    List[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						(lowercase_  ,     unittest.TestCase     ):





      '''simple docstring'''






      lowerCAmelCase__					:     List[str]  					=							(
          (
              FlaxDistilBertModel,
              FlaxDistilBertForMaskedLM,
              FlaxDistilBertForMultipleChoice,
              FlaxDistilBertForQuestionAnswering,
              FlaxDistilBertForSequenceClassification,
              FlaxDistilBertForTokenClassification,
              FlaxDistilBertForQuestionAnswering,
          )
          if is_flax_available()
          else ()
      )



      def       UpperCamelCase__		(self    :    Union[str, Any] ):

             '''simple docstring'''







             lowercase__      			= FlaxDistilBertModelTester(self )






      @slow
      def       UpperCamelCase__		(self    :    List[str] ):

             '''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(UpperCamelCase )






@require_flax
class   __lowerCAmelCase						(unittest.TestCase     ):





      '''simple docstring'''






      @slow
      def       UpperCamelCase__		(self    :    Optional[Any] ):

             '''simple docstring'''







             lowercase__      			= FlaxDistilBertModel.from_pretrained('''distilbert-base-uncased''' )
             lowercase__      			= np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]] )
             lowercase__      			= np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
             lowercase__      			= model(UpperCamelCase   ,						attention_mask=UpperCamelCase )[0]
             lowercase__      			= (1, 11, 768)
             self.assertEqual(output.shape   ,						UpperCamelCase )
             lowercase__      			= np.array([[[-0.16_39, 0.32_99, 0.16_48], [-0.17_46, 0.32_89, 0.17_10], [-0.18_84, 0.33_57, 0.18_10]]] )

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








'''simple docstring'''

# Lint as: python3
import itertools
import os
import re


lowerCamelCase							:   Any  						=	re.compile(R'([A-Z]+)([A-Z][a-z])')
lowerCamelCase							:   str  						=	re.compile(R'([a-z\d])([A-Z])')

lowerCamelCase							:   Optional[int]  						=	re.compile(R'(?<!_)_(?!_)')
lowerCamelCase							:   List[Any]  						=	re.compile(R'(_{2,})')

lowerCamelCase							:   str  						=	R'^\w+(\.\w+)*$'

lowerCamelCase							:   Dict  						=	R'<>:/\|?*'

def _SCREAMING_SNAKE_CASE      (A       )       ->       Any:






       """simple docstring"""







       lowercase__      			= _uppercase_uppercase_re.sub(R'''\1_\2'''  , A       )
       lowercase__      			= _lowercase_uppercase_re.sub(R'''\1_\2'''  , A       )
       return name.lower()

def _SCREAMING_SNAKE_CASE      (A       )       ->       Tuple:






       """simple docstring"""







       lowercase__      			= _single_underscore_re.split(A       )
       lowercase__      			= [_multiple_underscores_re.split(A       ) for n in name]
       return "".join(n.capitalize() for n in itertools.chain.from_iterable(A       ) if n != ''''''       )

def _SCREAMING_SNAKE_CASE      (A       )       ->       Tuple:






       """simple docstring"""







       if os.path.basename(A       ) != name:
              raise ValueError(f"Should be a dataset name, not a path: {name}"       )
       return camelcase_to_snakecase(A       )

def _SCREAMING_SNAKE_CASE      (A  , A       )       ->       Optional[Any]:






       """simple docstring"""







       if os.path.basename(A       ) != name:
              raise ValueError(f"Should be a dataset name, not a path: {name}"       )
       if not re.match(_split_re  , A       ):
              raise ValueError(f"Split name should match '{_split_re}'' but got '{split}'."       )
       return f"{filename_prefix_for_name(A       )}-{split}"

def _SCREAMING_SNAKE_CASE      (A  , A  , A  , A=None       )       ->       List[str]:






       """simple docstring"""







       lowercase__      			= filename_prefix_for_split(A  , A       )
       if filetype_suffix:
              prefix += f".{filetype_suffix}"
       lowercase__      			= os.path.join(A  , A       )
       return f"{filepath}*"

def _SCREAMING_SNAKE_CASE      (A  , A  , A  , A=None  , A=None       )       ->       Optional[Any]:






       """simple docstring"""







       lowercase__      			= filename_prefix_for_split(A  , A       )
       lowercase__      			= os.path.join(A  , A       )

       if shard_lengths:
              lowercase__      			= len(A       )
              lowercase__      			= [f"{prefix}-{shard_id:05d}-of-{num_shards:05d}" for shard_id in range(A       )]
              if filetype_suffix:
                     lowercase__      			= [filename + f".{filetype_suffix}" for filename in filenames]
              return filenames
       else:
              lowercase__      			= prefix
              if filetype_suffix:
                     filename += f".{filetype_suffix}"
              return [filename]




