
import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_PRETRAINING_MAPPING,
      AlbertForMaskedLM,
      AlbertForMultipleChoice,
      AlbertForPreTraining,
      AlbertForQuestionAnswering,
      AlbertForSequenceClassification,
      AlbertForTokenClassification,
      AlbertModel,
  )
  from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST

class __snake_case		:
    def __init__(  self		,snake_case		,snake_case=13		,snake_case=7		,snake_case=True		,snake_case=True		,snake_case=True		,snake_case=True		,snake_case=99		,snake_case=16		,snake_case=36		,snake_case=6		,snake_case=6		,snake_case=6		,snake_case=37		,snake_case="gelu"		,snake_case=0.1		,snake_case=0.1		,snake_case=512		,snake_case=16		,snake_case=2		,snake_case=0.02		,snake_case=3		,snake_case=4		,snake_case=None		,):



           '''simple docstring'''







           lowercase							:     Union[str, Any]       	=      parent
           lowercase							:     Dict       	=      batch_size
           lowercase							:     Optional[int]       	=      seq_length
           lowercase							:     Union[str, Any]       	=      is_training
           lowercase							:     Dict       	=      use_input_mask
           lowercase							:     Dict       	=      use_token_type_ids
           lowercase							:     str       	=      use_labels
           lowercase							:     Union[str, Any]       	=      vocab_size
           lowercase							:     int       	=      embedding_size
           lowercase							:     List[str]       	=      hidden_size
           lowercase							:     Dict       	=      num_hidden_layers
           lowercase							:     Optional[Any]       	=      num_hidden_groups
           lowercase							:     List[Any]       	=      num_attention_heads
           lowercase							:     Union[str, Any]       	=      intermediate_size
           lowercase							:     Any       	=      hidden_act
           lowercase							:     Tuple       	=      hidden_dropout_prob
           lowercase							:     Dict       	=      attention_probs_dropout_prob
           lowercase							:     Any       	=      max_position_embeddings
           lowercase							:     List[str]       	=      type_vocab_size
           lowercase							:     int       	=      type_sequence_label_size
           lowercase							:     Optional[Any]       	=      initializer_range
           lowercase							:     int       	=      num_labels
           lowercase							:     Optional[Any]       	=      num_choices
           lowercase							:     Union[str, Any]       	=      scope
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     str       	=      ids_tensor([self.batch_size, self.seq_length]		,self.vocab_size      )

           lowercase							:     Union[str, Any]       	=      None
           if self.use_input_mask:
                  lowercase							:     Optional[Any]       	=      random_attention_mask([self.batch_size, self.seq_length]      )

           lowercase							:     Optional[int]       	=      None
           if self.use_token_type_ids:
                  lowercase							:     Any       	=      ids_tensor([self.batch_size, self.seq_length]		,self.type_vocab_size      )

           lowercase							:     Tuple       	=      None
           lowercase							:     Any       	=      None
           lowercase							:     Any       	=      None
           if self.use_labels:
                  lowercase							:     List[Any]       	=      ids_tensor([self.batch_size]		,self.type_sequence_label_size      )
                  lowercase							:     Dict       	=      ids_tensor([self.batch_size, self.seq_length]		,self.num_labels      )
                  lowercase							:     List[str]       	=      ids_tensor([self.batch_size]		,self.num_choices      )

           lowercase							:     Tuple       	=      self.get_config()

           return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           return 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		,initializer_range=self.initializer_range		,num_hidden_groups=self.num_hidden_groups		,)
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case      ):



           '''simple docstring'''







           lowercase							:     int       	=      AlbertModel(config=snake_case      )
           model.to(snake_case      )
           model.eval()
           lowercase							:     List[str]       	=      model(snake_case		,attention_mask=snake_case		,token_type_ids=snake_case      )
           lowercase							:     str       	=      model(snake_case		,token_type_ids=snake_case      )
           lowercase							:     Optional[Any]       	=      model(snake_case      )
           self.parent.assertEqual(result.last_hidden_state.shape		,(self.batch_size, self.seq_length, self.hidden_size)      )
           self.parent.assertEqual(result.pooler_output.shape		,(self.batch_size, self.hidden_size)      )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case      ):



           '''simple docstring'''







           lowercase							:     Tuple       	=      AlbertForPreTraining(config=snake_case      )
           model.to(snake_case      )
           model.eval()
           lowercase							:     Union[str, Any]       	=      model(
               snake_case		,attention_mask=snake_case		,token_type_ids=snake_case		,labels=snake_case		,sentence_order_label=snake_case		,)
           self.parent.assertEqual(result.prediction_logits.shape		,(self.batch_size, self.seq_length, self.vocab_size)      )
           self.parent.assertEqual(result.sop_logits.shape		,(self.batch_size, config.num_labels)      )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case      ):



           '''simple docstring'''







           lowercase							:     int       	=      AlbertForMaskedLM(config=snake_case      )
           model.to(snake_case      )
           model.eval()
           lowercase							:     int       	=      model(snake_case		,attention_mask=snake_case		,token_type_ids=snake_case		,labels=snake_case      )
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.seq_length, self.vocab_size)      )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case      ):



           '''simple docstring'''







           lowercase							:     str       	=      AlbertForQuestionAnswering(config=snake_case      )
           model.to(snake_case      )
           model.eval()
           lowercase							:     List[str]       	=      model(
               snake_case		,attention_mask=snake_case		,token_type_ids=snake_case		,start_positions=snake_case		,end_positions=snake_case		,)
           self.parent.assertEqual(result.start_logits.shape		,(self.batch_size, self.seq_length)      )
           self.parent.assertEqual(result.end_logits.shape		,(self.batch_size, self.seq_length)      )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case      ):



           '''simple docstring'''







           lowercase							:     str       	=      self.num_labels
           lowercase							:     Tuple       	=      AlbertForSequenceClassification(snake_case      )
           model.to(snake_case      )
           model.eval()
           lowercase							:     Tuple       	=      model(snake_case		,attention_mask=snake_case		,token_type_ids=snake_case		,labels=snake_case      )
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.num_labels)      )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case      ):



           '''simple docstring'''







           lowercase							:     Tuple       	=      self.num_labels
           lowercase							:     Optional[int]       	=      AlbertForTokenClassification(config=snake_case      )
           model.to(snake_case      )
           model.eval()
           lowercase							:     int       	=      model(snake_case		,attention_mask=snake_case		,token_type_ids=snake_case		,labels=snake_case      )
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.seq_length, self.num_labels)      )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case      ):



           '''simple docstring'''







           lowercase							:     Any       	=      self.num_choices
           lowercase							:     Any       	=      AlbertForMultipleChoice(config=snake_case      )
           model.to(snake_case      )
           model.eval()
           lowercase							:     List[Any]       	=      input_ids.unsqueeze(1      ).expand(-1		,self.num_choices		,-1      ).contiguous()
           lowercase							:     Tuple       	=      token_type_ids.unsqueeze(1      ).expand(-1		,self.num_choices		,-1      ).contiguous()
           lowercase							:     Tuple       	=      input_mask.unsqueeze(1      ).expand(-1		,self.num_choices		,-1      ).contiguous()
           lowercase							:     Dict       	=      model(
               snake_case		,attention_mask=snake_case		,token_type_ids=snake_case		,labels=snake_case		,)
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.num_choices)      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Tuple       	=      self.prepare_config_and_inputs()
           (
               (
               lowercase
           )	,			(
               lowercase
           )	,			(
               lowercase
           )	,			(
               lowercase
           )	,			(
               lowercase
           )	,			(
               lowercase
           )	,			(
               lowercase
           )	,			
           )							:     str       	=      config_and_inputs
           lowercase							:     int       	=      {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
           return config, inputs_dict






@require_torch
class __snake_case		(      lowerCAmelCase					,	lowerCAmelCase					,	unittest.TestCase     ):
    _a		:						Tuple=				(
        (
            AlbertModel,
            AlbertForPreTraining,
            AlbertForMaskedLM,
            AlbertForMultipleChoice,
            AlbertForSequenceClassification,
            AlbertForTokenClassification,
            AlbertForQuestionAnswering,
        )
        if is_torch_available()
        else ()
    )
    _a		:						List[str]=				(
        {
            "feature-extraction": AlbertModel,
            "fill-mask": AlbertForMaskedLM,
            "question-answering": AlbertForQuestionAnswering,
            "text-classification": AlbertForSequenceClassification,
            "token-classification": AlbertForTokenClassification,
            "zero-shot": AlbertForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    _a		:						Dict=				True
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case=False      ):



           '''simple docstring'''







           lowercase							:     List[str]       	=      super()._prepare_for_class(snake_case		,snake_case		,return_labels=snake_case      )

           if return_labels:
                  if model_class in get_values(snake_case      ):
                         lowercase							:     List[Any]       	=      torch.zeros(
                             (self.model_tester.batch_size, self.model_tester.seq_length)		,dtype=torch.long		,device=snake_case      )
                         lowercase							:     str       	=      torch.zeros(
                             self.model_tester.batch_size		,dtype=torch.long		,device=snake_case      )
           return inputs_dict
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     str       	=      AlbertModelTester(self      )
           lowercase							:     Optional[int]       	=      ConfigTester(self		,config_class=snake_case		,hidden_size=37      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           self.config_tester.run_common_tests()
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Dict       	=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*snake_case      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Union[str, Any]       	=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_pretraining(*snake_case      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Dict       	=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*snake_case      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     List[Any]       	=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*snake_case      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Optional[int]       	=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*snake_case      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Tuple       	=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*snake_case      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     str       	=      self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
                  lowercase							:     Optional[Any]       	=      type
                  self.model_tester.create_and_check_model(*snake_case      )
    @slow
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowercase							:     List[Any]       	=      AlbertModel.from_pretrained(snake_case      )
                  self.assertIsNotNone(snake_case      )






@require_torch
class __snake_case		(      unittest.TestCase     ):
    @slow
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Optional[int]       	=      AlbertModel.from_pretrained("""albert-base-v2"""      )
           lowercase							:     int       	=      torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]      )
           lowercase							:     Optional[int]       	=      torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]      )
           with torch.no_grad():
                  lowercase							:     List[str]       	=      model(snake_case		,attention_mask=snake_case      )[0]
           lowercase							:     Any       	=      torch.Size((1, 11, 768)      )
           self.assertEqual(output.shape		,snake_case      )
           lowercase							:     int       	=      torch.tensor(
               [[[-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(torch.allclose(output[:, 1:4, 1:4]		,snake_case		,atol=1e-4      )      )

"""simple docstring"""




from __future__ import annotations

from typing import Generic, TypeVar

__lowerCamelCase							 =     TypeVar("T")
class 			UpperCamelCase__(	Generic[T]       ):





						def __init__(				self						,__UpperCAmelCase   )    -> None:
													A__		   =							data
													A__		   =							self
													A__		   =							0






class 			UpperCamelCase__(	Generic[T]       ):





						def __init__(				self   )    -> None:
													# map from node name to the node object
													A__		   =							{}





						def  snake_case__		(				self						,__UpperCAmelCase   )    -> None:
													# create a new set with x as its member
													A__		   =							DisjointSetTreeNode(__UpperCAmelCase   )





						def  snake_case__		(				self						,__UpperCAmelCase   )    -> DisjointSetTreeNode[T]:
													# find the set x belongs to (with path-compression)
													A__		   =							self.map[data]
													if elem_ref != elem_ref.parent:
																				A__		   =							self.find_set(elem_ref.parent.data   )
													return elem_ref.parent





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase   )    -> None:
													# helper function for union operation
													if nodea.rank > nodea.rank:
																				A__		   =							nodea
													else:
																				A__		   =							nodea
																				if nodea.rank == nodea.rank:
																											nodea.rank += 1





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase   )    -> None:
													# merge 2 disjoint sets
													self.link(self.find_set(__UpperCAmelCase   )						,self.find_set(__UpperCAmelCase   )   )






class 			UpperCamelCase__(	Generic[T]       ):





						def __init__(				self   )    -> None:
													# connections: map from the node to the neighbouring nodes (with weights)
													A__		   =							{}





						def  snake_case__		(				self						,__UpperCAmelCase   )    -> None:
													# add a node ONLY if its not present in the graph
													if node not in self.connections:
																				A__		   =							{}





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )    -> None:
													# add an edge with the given weight
													self.add_node(__UpperCAmelCase   )
													self.add_node(__UpperCAmelCase   )
													A__		   =							weight
													A__		   =							weight





						def  snake_case__		(				self   )    -> GraphUndirectedWeighted[T]:
													A__		   =							[]
													A__		   =							set()
													for start in self.connections:
																				for end in self.connections[start]:
																											if (start, end) not in seen:
																																		seen.add((end, start)   )
																																		edges.append((start, end, self.connections[start][end])   )
													edges.sort(key=lambda __UpperCAmelCase   : x[2]   )

													# creating the disjoint set
													A__		   =							DisjointSetTree[T]()
													for node in self.connections:
																				disjoint_set.make_set(__UpperCAmelCase   )

													# MST generation
													A__		   =							0
													A__		   =							0
													A__		   =							GraphUndirectedWeighted[T]()
													while num_edges < len(self.connections   ) - 1:
																				A__   ,     A__   ,     A__		   =							edges[index]
																				index += 1
																				A__		   =							disjoint_set.find_set(__UpperCAmelCase   )
																				A__		   =							disjoint_set.find_set(__UpperCAmelCase   )
																				if parent_u != parent_v:
																											num_edges += 1
																											graph.add_edge(__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )
																											disjoint_set.union(__UpperCAmelCase						,__UpperCAmelCase   )
													return graph





