
"""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 		a_    (			unittest.TestCase						):







    def __init__(      self      :							int  ,       __UpperCamelCase      :							Any  ,       __UpperCamelCase      :							Union[str, Any]=13  ,       __UpperCamelCase      :							Union[str, Any]=7  ,       __UpperCamelCase      :							Optional[Any]=True  ,       __UpperCamelCase      :							List[Any]=True  ,       __UpperCamelCase      :							Any=True  ,       __UpperCamelCase      :							Tuple=True  ,       __UpperCamelCase      :							Tuple=99  ,       __UpperCamelCase      :							str=32  ,       __UpperCamelCase      :							Optional[int]=5  ,       __UpperCamelCase      :							List[Any]=4  ,       __UpperCamelCase      :							List[Any]=37  ,       __UpperCamelCase      :							str="gelu"  ,       __UpperCamelCase      :							Optional[Any]=0.1  ,       __UpperCamelCase      :							Union[str, Any]=0.1  ,       __UpperCamelCase      :							List[Any]=5_12  ,       __UpperCamelCase      :							Dict=16  ,       __UpperCamelCase      :							int=2  ,       __UpperCamelCase      :							Union[str, Any]=0.0_2  ,       __UpperCamelCase      :							Optional[Any]=4  ,       )		->Tuple:






        '''simple docstring'''






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







    def   _snake_case       (      self      :							List[str]  )		->str:






        '''simple docstring'''






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

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

        _UpperCAmelCase        = 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   _snake_case       (      self      :							List[Any]  )		->Union[str, Any]:






        '''simple docstring'''






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







@require_flax
class 		a_    (			__A  ,    unittest.TestCase						):
    a   :					Dict			     =					(
        (
            FlaxDistilBertModel,
            FlaxDistilBertForMaskedLM,
            FlaxDistilBertForMultipleChoice,
            FlaxDistilBertForQuestionAnswering,
            FlaxDistilBertForSequenceClassification,
            FlaxDistilBertForTokenClassification,
            FlaxDistilBertForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )







    def   _snake_case       (      self      :							List[str]  )		->Optional[Any]:






        '''simple docstring'''






        _UpperCAmelCase        = FlaxDistilBertModelTester(self  )







    @slow
    def   _snake_case       (      self      :							Optional[Any]  )		->List[str]:






        '''simple docstring'''






        for model_class_name in self.all_model_classes:
            _UpperCAmelCase        = model_class_name.from_pretrained("""distilbert-base-uncased"""  )
            _UpperCAmelCase        = model(np.ones((1, 1)  )  )
            self.assertIsNotNone(lowerCAmelCase_  )







@require_flax
class 		a_    (			unittest.TestCase						):







    @slow
    def   _snake_case       (      self      :							str  )		->Optional[int]:






        '''simple docstring'''






        _UpperCAmelCase        = FlaxDistilBertModel.from_pretrained("""distilbert-base-uncased"""  )
        _UpperCAmelCase        = np.array([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]  )
        _UpperCAmelCase        = np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
        _UpperCAmelCase        = model(lowerCAmelCase_  ,       attention_mask=lowerCAmelCase_  )[0]
        _UpperCAmelCase        = (1, 11, 7_68)
        self.assertEqual(output.shape  ,       lowerCAmelCase_  )
        _UpperCAmelCase        = np.array([[[-0.1_6_3_9, 0.3_2_9_9, 0.1_6_4_8], [-0.1_7_4_6, 0.3_2_8_9, 0.1_7_1_0], [-0.1_8_8_4, 0.3_3_5_7, 0.1_8_1_0]]]  )

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







"""simple docstring"""
import os
from typing import Dict, List, Tuple, TypeVar, Union


lowerCamelCase_							      = TypeVar('''T''')

lowerCamelCase_							      = Union[List[T], Tuple[T, ...]]
lowerCamelCase_							      = Union[T, List[T], Dict[str, T]]
lowerCamelCase_							      = Union[str, bytes, os.PathLike]



