
import unittest

import numpy as np

from transformers import AlbertConfig, 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.albert.modeling_flax_albert import (
          FlaxAlbertForMaskedLM,
          FlaxAlbertForMultipleChoice,
          FlaxAlbertForPreTraining,
          FlaxAlbertForQuestionAnswering,
          FlaxAlbertForSequenceClassification,
          FlaxAlbertForTokenClassification,
          FlaxAlbertModel,
      )




class      UpperCamelCase (     unittest.TestCase   ):

    def __init__(			self :      Union[str, Any]							,    snake_case__ :      List[str]							,    snake_case__ :      str=1_3							,    snake_case__ :      Dict=7							,    snake_case__ :      Any=True							,    snake_case__ :      List[str]=True							,    snake_case__ :      List[str]=True							,    snake_case__ :      Tuple=True							,    snake_case__ :      str=9_9							,    snake_case__ :      Dict=3_2							,    snake_case__ :      Any=5							,    snake_case__ :      List[Any]=4							,    snake_case__ :      Tuple=3_7							,    snake_case__ :      List[Any]="gelu"							,    snake_case__ :      str=0.1							,    snake_case__ :      Optional[int]=0.1							,    snake_case__ :      int=5_1_2							,    snake_case__ :      int=1_6							,    snake_case__ :      List[str]=2							,    snake_case__ :      str=0.02							,    snake_case__ :      List[str]=4							,    ):

     """simple docstring"""






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

    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






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

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

     SCREAMING_SNAKE_CASE								=			None
     if self.use_token_type_ids:
      SCREAMING_SNAKE_CASE								=			ids_tensor([self.batch_size, self.seq_length]							,    self.type_vocab_size    )

     SCREAMING_SNAKE_CASE								=			AlbertConfig(
         vocab_size=self.vocab_size							,    hidden_size=self.hidden_size							,    num_hidden_layers=self.num_hidden_layers							,    num_attention_heads=self.num_attention_heads							,    intermediate_size=self.intermediate_size							,    hidden_act=self.hidden_act							,    hidden_dropout_prob=self.hidden_dropout_prob							,    attention_probs_dropout_prob=self.attention_probs_dropout_prob							,    max_position_embeddings=self.max_position_embeddings							,    type_vocab_size=self.type_vocab_size							,    is_decoder=snake_case__							,    initializer_range=self.initializer_range							,    )

     return config, input_ids, token_type_ids, attention_mask




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

     """simple docstring"""






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




@require_flax
class      UpperCamelCase (     SCREAMING_SNAKE_CASE	,			unittest.TestCase   ):
    __UpperCamelCase         =(
        (
            FlaxAlbertModel,
            FlaxAlbertForPreTraining,
            FlaxAlbertForMaskedLM,
            FlaxAlbertForMultipleChoice,
            FlaxAlbertForQuestionAnswering,
            FlaxAlbertForSequenceClassification,
            FlaxAlbertForTokenClassification,
            FlaxAlbertForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			FlaxAlbertModelTester(self    )




    @slow
    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     for model_class_name in self.all_model_classes:
      SCREAMING_SNAKE_CASE								=			model_class_name.from_pretrained('albert-base-v2'    )
      SCREAMING_SNAKE_CASE								=			model(np.ones((1, 1)    )    )
      self.assertIsNotNone(snake_case__    )




@require_flax
class      UpperCamelCase (     unittest.TestCase   ):

    @slow
    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			FlaxAlbertModel.from_pretrained('albert-base-v2'    )
     SCREAMING_SNAKE_CASE								=			np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]    )
     SCREAMING_SNAKE_CASE								=			np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__    )[0]
     SCREAMING_SNAKE_CASE								=			(1, 1_1, 7_6_8)
     self.assertEqual(output.shape							,    snake_case__    )
     SCREAMING_SNAKE_CASE								=			np.array(
         [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]]    )

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




import random


def        __lowerCAmelCase    (       _UpperCamelCase		:  int		,							_UpperCamelCase		:  float		,							_UpperCamelCase		:  bool = False	)		->							dict:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE								=			{i: [] for i in range(_UpperCamelCase	)}

 # if probability is greater or equal than 1, then generate a complete graph
 if probability >= 1:
  return complete_graph(_UpperCamelCase	)
 # if probability is lower or equal than 0, then return a graph without edges
 if probability <= 0:
  return graph

 # for each couple of nodes, add an edge from u to v
 # if the number randomly generated is greater than probability probability
 for i in range(_UpperCamelCase	):
  for j in range(i + 1		,							_UpperCamelCase	):
   if random.random() < probability:
    graph[i].append(_UpperCamelCase	)
    if not directed:
     # if the graph is undirected, add an edge in from j to i, either
     graph[j].append(_UpperCamelCase	)
 return graph


def        __lowerCAmelCase    (       _UpperCamelCase		:  int	)		->							dict:
 '''simple docstring'''


 return {
     i: [j for j in range(_UpperCamelCase	) if i != j] for i in range(_UpperCamelCase	)
 }


if __name__ == "__main__":
      import doctest

      doctest.testmod()

