



"""simple docstring"""

import tempfile
import unittest

from transformers import TaConfig, is_torch_available
from transformers.testing_utils import (
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
    torch_device,
)

from ...generation.test_utils import GenerationTesterMixin
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import AutoTokenizer, UMTaForConditionalGeneration, UMTaForQuestionAnswering, UMTaModel




class        _UpperCAmelCase       :

 def __init__(     self		:       Tuple  ,     lowercase_		:       Union[str, Any]  ,     lowercase_		:       List[str]=99  ,     lowercase_		:       List[str]=13  ,     lowercase_		:       Optional[int]=7  ,     lowercase_		:       Tuple=9  ,     lowercase_		:       Any=True  ,     lowercase_		:       str=True  ,     lowercase_		:       Optional[Any]=False  ,     lowercase_		:       str=32  ,     lowercase_		:       List[str]=5  ,     lowercase_		:       Union[str, Any]=4  ,     lowercase_		:       Tuple=37  ,     lowercase_		:       Optional[int]=8  ,     lowercase_		:       str=0.1  ,     lowercase_		:       Tuple=0.0_02  ,     lowercase_		:       Optional[Any]=1  ,     lowercase_		:       Optional[int]=0  ,     lowercase_		:       Tuple=0  ,     lowercase_		:       str=None  ,     lowercase_		:       Optional[Any]=None  ,     ):
     snake_case_ :     int				=	parent
     snake_case_ :     Tuple				=	batch_size
     snake_case_ :     List[Any]				=	encoder_seq_length
     snake_case_ :     List[str]				=	decoder_seq_length
     # For common tests
     snake_case_ :     str				=	self.decoder_seq_length
     snake_case_ :     Optional[int]				=	is_training
     snake_case_ :     List[str]				=	use_attention_mask
     snake_case_ :     Union[str, Any]				=	use_labels
     snake_case_ :     List[str]				=	vocab_size
     snake_case_ :     Union[str, Any]				=	hidden_size
     snake_case_ :     int				=	num_hidden_layers
     snake_case_ :     Dict				=	num_attention_heads
     snake_case_ :     List[str]				=	d_ff
     snake_case_ :     Tuple				=	relative_attention_num_buckets
     snake_case_ :     List[Any]				=	dropout_rate
     snake_case_ :     List[str]				=	initializer_factor
     snake_case_ :     str				=	eos_token_id
     snake_case_ :     Union[str, Any]				=	pad_token_id
     snake_case_ :     Any				=	decoder_start_token_id
     snake_case_ :     Union[str, Any]				=	None
     snake_case_ :     int				=	decoder_layers

 def 						_snake_case					(     self		:       int				):
     return TaConfig.from_pretrained('''google/umt5-base'''				)

 def 						_snake_case					(     self		:       List[Any]  ,     lowercase_		:       Any  ,     lowercase_		:       Any  ,     lowercase_		:       Tuple  ,     lowercase_		:       Optional[int]=None  ,     lowercase_		:       Optional[Any]=None  ,     lowercase_		:       Any=None  ,     lowercase_		:       int=None  ,     lowercase_		:       List[str]=None  ,     ):
     if attention_mask is None:
         snake_case_ :     str				=	input_ids.ne(config.pad_token_id				)
     if decoder_attention_mask is None:
         snake_case_ :     str				=	decoder_input_ids.ne(config.pad_token_id				)
     if head_mask is None:
         snake_case_ :     Optional[int]				=	torch.ones(config.num_hidden_layers  ,     config.num_attention_heads  ,     device=lowercase_				)
     if decoder_head_mask is None:
         snake_case_ :     int				=	torch.ones(config.num_decoder_layers  ,     config.num_attention_heads  ,     device=lowercase_				)
     if cross_attn_head_mask is None:
         snake_case_ :     Tuple				=	torch.ones(
             config.num_decoder_layers  ,     config.num_attention_heads  ,     device=lowercase_				)
     return {
         "input_ids": input_ids,
         "decoder_input_ids": decoder_input_ids,
         "attention_mask": attention_mask,
         "decoder_attention_mask": decoder_attention_mask,
         "head_mask": head_mask,
         "decoder_head_mask": decoder_head_mask,
         "cross_attn_head_mask": cross_attn_head_mask,
     }

 def 						_snake_case					(     self		:       Union[str, Any]				):
     snake_case_ :     List[str]				=	ids_tensor([self.batch_size, self.encoder_seq_length]  ,     self.vocab_size				)
     snake_case_ :     Optional[int]				=	ids_tensor([self.batch_size, self.decoder_seq_length]  ,     self.vocab_size				)

     # we need to clamp the input ids here to avoid having pad token in between
     # this is because for NllbMoe the position_ids are prepared such that
     # all pad tokens have pos id = 2 and rest are between 2..seq_length
     # and the seq_length here is seq_length - num_pad_tokens
     # but when using past, there is no way of knowing if the past input ids had
     # pad tokens in them, which results in incorrect seq_lenth and which in turn results in
     # position_ids being off by num_pad_tokens in past input
     snake_case_ :     Optional[int]				=	input_ids.clamp(self.pad_token_id + 1				)
     snake_case_ :     List[str]				=	decoder_input_ids.clamp(self.pad_token_id + 1				)

     snake_case_ :     List[Any]				=	self.get_config()
     snake_case_ :     int				=	config.num_attention_heads
     snake_case_ :     Union[str, Any]				=	self.prepare_inputs_dict(lowercase_  ,     lowercase_  ,     lowercase_				)
     return config, input_dict

 def 						_snake_case					(     self		:       int				):
     snake_case_ :     Dict				=	self.prepare_config_and_inputs()
     return config, inputs_dict

 def 						_snake_case					(     self		:       List[Any]				):
     return TaConfig(
         vocab_size=166  ,     d_model=self.hidden_size  ,     d_ff=self.d_ff  ,     d_kv=self.hidden_size // self.num_attention_heads  ,     num_layers=self.num_hidden_layers  ,     num_decoder_layers=self.decoder_layers  ,     num_heads=self.num_attention_heads  ,     relative_attention_num_buckets=self.relative_attention_num_buckets  ,     dropout_rate=self.dropout_rate  ,     initializer_factor=self.initializer_factor  ,     eos_token_id=self.eos_token_id  ,     bos_token_id=self.pad_token_id  ,     pad_token_id=self.pad_token_id  ,     decoder_start_token_id=self.decoder_start_token_id  ,     )

 def 						_snake_case					(     self		:       Dict				):
     return TaConfig(
         vocab_size=self.vocab_size  ,     d_model=self.hidden_size  ,     d_ff=self.d_ff  ,     d_kv=self.hidden_size // self.num_attention_heads  ,     num_layers=self.num_hidden_layers  ,     num_decoder_layers=self.decoder_layers  ,     num_heads=self.num_attention_heads  ,     relative_attention_num_buckets=self.relative_attention_num_buckets  ,     dropout_rate=self.dropout_rate  ,     initializer_factor=self.initializer_factor  ,     eos_token_id=self.eos_token_id  ,     bos_token_id=self.pad_token_id  ,     pad_token_id=self.pad_token_id  ,     decoder_start_token_id=self.decoder_start_token_id  ,     )

 def 						_snake_case					(     self		:       Optional[int]  ,     lowercase_		:       Optional[int]  ,     lowercase_		:       Any  ,     lowercase_		:       Tuple  ,     lowercase_		:       Dict  ,     lowercase_		:       List[str]  ,     lowercase_		:       Tuple  ,     ):
     snake_case_ :     Optional[Any]				=	UMTaModel(config=lowercase_				)
     model.to(lowercase_				)
     model.eval()
     snake_case_ :     Optional[int]				=	model(
         input_ids=lowercase_  ,     decoder_input_ids=lowercase_  ,     attention_mask=lowercase_  ,     decoder_attention_mask=lowercase_  ,     )
     snake_case_ :     Optional[int]				=	model(input_ids=lowercase_  ,     decoder_input_ids=lowercase_				)
     snake_case_ :     Any				=	result.last_hidden_state
     snake_case_ :     List[Any]				=	result.past_key_values
     snake_case_ :     List[str]				=	result.encoder_last_hidden_state

     self.parent.assertEqual(encoder_output.size()  ,     (self.batch_size, self.encoder_seq_length, self.hidden_size)				)
     self.parent.assertEqual(decoder_output.size()  ,     (self.batch_size, self.decoder_seq_length, self.hidden_size)				)
     # There should be `num_layers` key value embeddings stored in decoder_past
     self.parent.assertEqual(len(lowercase_				)  ,     config.num_layers				)
     # There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past tuple
     self.parent.assertEqual(len(decoder_past[0]				)  ,     4				)

 def 						_snake_case					(     self		:       Dict  ,     lowercase_		:       Union[str, Any]  ,     lowercase_		:       Tuple  ,     lowercase_		:       Tuple  ,     lowercase_		:       List[str]  ,     lowercase_		:       List[Any]  ,     lowercase_		:       List[str]  ,     ):
     snake_case_ :     Any				=	UMTaModel(config=lowercase_				).get_decoder().to(lowercase_				).eval()
     # first forward pass
     snake_case_ :     Tuple				=	model(lowercase_  ,     use_cache=lowercase_				)
     snake_case_ :     List[Any]				=	model(lowercase_				)
     snake_case_ :     str				=	model(lowercase_  ,     use_cache=lowercase_				)

     self.parent.assertTrue(len(lowercase_				) == len(lowercase_				)				)
     self.parent.assertTrue(len(lowercase_				) == len(lowercase_				) + 1				)

     snake_case_ :     List[Any]				=	outputs.to_tuple()

     # create hypothetical next token and extent to next_input_ids
     snake_case_ :     str				=	ids_tensor((self.batch_size, 1)  ,     config.vocab_size				)

     # append to next input_ids and
     snake_case_ :     Tuple				=	torch.cat([input_ids, next_tokens]  ,     dim=-1				)

     snake_case_ :     Optional[Any]				=	model(lowercase_				)["""last_hidden_state"""]
     snake_case_ :     Optional[int]				=	model(lowercase_  ,     past_key_values=lowercase_				)["""last_hidden_state"""]

     # select random slice
     snake_case_ :     Union[str, Any]				=	ids_tensor((1,)  ,     output_from_past.shape[-1]				).item()
     snake_case_ :     Optional[int]				=	output_from_no_past[:, -1, random_slice_idx].detach()
     snake_case_ :     str				=	output_from_past[:, 0, random_slice_idx].detach()

     # test that outputs are equal for slice
     self.parent.assertTrue(torch.allclose(lowercase_  ,     lowercase_  ,     atol=1E-3				)				)






 def 						_snake_case					(     self		:       Dict  ,     lowercase_		:       Dict  ,     lowercase_		:       str  ,     ):
     snake_case_ :     Dict				=	UMTaModel(config=lowercase_				).to(lowercase_				).half().eval()
     snake_case_ :     int				=	model(**lowercase_				)["""last_hidden_state"""]
     self.parent.assertFalse(torch.isnan(lowercase_				).any().item()				)





@require_torch
class        _UpperCAmelCase       (			UpperCAmelCase__    ,     UpperCAmelCase__    ,     UpperCAmelCase__    ,     unittest.TestCase):
 _lowerCAmelCase   :			Any							=		(
     (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
 )
 _lowerCAmelCase   :			Optional[int]							=		(UMTaForConditionalGeneration,) if is_torch_available() else ()
 _lowerCAmelCase   :			Union[str, Any]							=		(
     {
         "conversational": UMTaForConditionalGeneration,
         "feature-extraction": UMTaModel,
         "summarization": UMTaForConditionalGeneration,
         "text2text-generation": UMTaForConditionalGeneration,
         "translation": UMTaForConditionalGeneration,
         "question-answering": UMTaForQuestionAnswering,
     }
     if is_torch_available()
     else {}
 )
 _lowerCAmelCase   :			Any							=		True
 _lowerCAmelCase   :			Union[str, Any]							=		False
 _lowerCAmelCase   :			int							=		False
 _lowerCAmelCase   :			int							=		True
 _lowerCAmelCase   :			str							=		True
 # The small UMT5 model needs higher percentages for CPU/MP tests
 _lowerCAmelCase   :			Optional[Any]							=		[0.8, 0.9]

 def 						_snake_case					(     self		:       Optional[Any]				):
     snake_case_ :     Tuple				=	UMTaModelTester(self				)

 @unittest.skip('''Test has a segmentation fault on torch 1.8.0'''				)
 def 						_snake_case					(     self		:       Dict				):
     snake_case_ :     List[str]				=	self.model_tester.prepare_config_and_inputs()
     snake_case_ :     Union[str, Any]				=	UMTaModel(config_and_inputs[0]				).to(lowercase_				)
     with tempfile.TemporaryDirectory() as tmpdirname:
         torch.onnx.export(
             lowercase_  ,     (config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])  ,     f"{tmpdirname}/t5_test.onnx"  ,     export_params=lowercase_  ,     opset_version=9  ,     input_names=['''input_ids''', '''decoder_input_ids''']  ,     )

 @unittest.skipIf(torch_device == '''cpu'''  ,     '''Cant do half precision'''				)
 def 						_snake_case					(     self		:       Optional[int]				):
     snake_case_ :     Optional[int]				=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model_fpaa_forward(*lowercase_				)

 def 						_snake_case					(     self		:       List[Any]				):
     snake_case_ :     Any				=	["""encoder_attentions""", """decoder_attentions""", """cross_attentions"""]
     snake_case_ :     Any				=	self.model_tester.prepare_config_and_inputs()
     snake_case_ :     int				=	config_and_inputs[0]
     snake_case_ :     Tuple				=	UMTaForConditionalGeneration(lowercase_				).eval()
     model.to(lowercase_				)

     snake_case_ :     List[Any]				=	{
         """head_mask""": torch.zeros(config.num_layers  ,     config.num_heads  ,     device=lowercase_				),
         """decoder_head_mask""": torch.zeros(config.num_decoder_layers  ,     config.num_heads  ,     device=lowercase_				),
         """cross_attn_head_mask""": torch.zeros(config.num_decoder_layers  ,     config.num_heads  ,     device=lowercase_				),
     }

     for attn_name, (name, mask) in zip(lowercase_  ,     head_masking.items()				):
         snake_case_ :     str				=	{name: mask}
         # Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
         if name == "head_mask":
             snake_case_ :     Optional[Any]				=	torch.ones(
                 config.num_decoder_layers  ,     config.num_heads  ,     device=lowercase_				)

         snake_case_ :     Union[str, Any]				=	model.generate(
             config_and_inputs[1]['''input_ids''']  ,     num_beams=1  ,     max_length=3  ,     output_attentions=lowercase_  ,     return_dict_in_generate=lowercase_  ,     **lowercase_  ,     )
         # We check the state of decoder_attentions and cross_attentions just from the last step
         snake_case_ :     int				=	out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
         self.assertEqual(sum([w.sum().item() for w in attn_weights]				)  ,     0.0				)






 @unittest.skip('''Does not work on the tiny model as we keep hitting edge cases.'''				)
 def 						_snake_case					(     self		:       List[Any]				):
     pass





@require_torch
@require_sentencepiece
@require_tokenizers
class        _UpperCAmelCase       (			unittest.TestCase):

 @slow
 @unittest.skip(
     '''Unless we stop stripping left and right by default for all special tokens, the expected ids obtained here will not match the original ones. Wait for https://github.com/huggingface/transformers/pull/23909 to be merged'''				)
 def 						_snake_case					(     self		:       Tuple				):
     snake_case_ :     Tuple				=	UMTaForConditionalGeneration.from_pretrained('''google/umt5-small'''  ,     return_dict=lowercase_				).to(lowercase_				)
     snake_case_ :     Dict				=	AutoTokenizer.from_pretrained('''google/umt5-small'''  ,     use_fast=lowercase_  ,     legacy=lowercase_				)
     snake_case_ :     Tuple				=	[
         """Bonjour monsieur <extra_id_0> bien <extra_id_1>.""",
         """No se como puedo <extra_id_0>.""",
         """This is the reason why we <extra_id_0> them.""",
         """The <extra_id_0> walks in <extra_id_1>, seats""",
         """A <extra_id_0> walks into a bar and orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.""",
     ]
     snake_case_ :     int				=	tokenizer(lowercase_  ,     return_tensors='''pt'''  ,     padding=lowercase_				).input_ids
     # fmt: off
     snake_case_ :     Optional[int]				=	torch.tensor(
         [
             [ 38530, 210703, 256299, 1410, 256298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
             [   826, 321, 671, 25922, 256299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
             [  1460, 339, 312, 19014, 10620, 758, 256299, 2355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
             [   517, 256299, 14869, 281, 301, 256298, 275, 119983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
             [   320, 256299, 14869, 281, 2234, 289, 2275, 333,61391, 289, 256298, 543, 256297, 168714, 329, 256296,274, 1],
         ]				)
     # fmt: on
     torch.testing.assert_allclose(lowercase_  ,     lowercase_				)

     snake_case_ :     Tuple				=	model.generate(input_ids.to(lowercase_				)				)
     snake_case_ :     Union[str, Any]				=	[
         """<pad><extra_id_0> et<extra_id_1> [eod] <extra_id_2><extra_id_55>.. [eod] 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 <extra_id_56>ajšietosto<extra_id_56>lleux<extra_id_19><extra_id_6>ajšie</s>""",
         """<pad><extra_id_0>.<extra_id_1>.,<0x0A>...spech <0x0A><extra_id_20> <extra_id_21></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>""",
         """<pad><extra_id_0> are not going to be a part of the world. We are not going to be a part of<extra_id_1> and<extra_id_2><0x0A><extra_id_48>.<extra_id_48></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>""",
         """<pad><extra_id_0> door<extra_id_1>, the door<extra_id_2> 피해[/</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>""",
         """<pad><extra_id_0>nyone who<extra_id_1> drink<extra_id_2> a<extra_id_3> alcohol<extra_id_4> A<extra_id_5> A. This<extra_id_6> I<extra_id_7><extra_id_52><extra_id_53></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>""",
     ]
     snake_case_ :     Tuple				=	tokenizer.batch_decode(lowercase_				)
     self.assertEqual(lowercase_  ,     lowercase_				)









'''simple docstring'''







class A__       :






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




             '''simple docstring'''
             _a    : Tuple     					=0
             _a    : Any     					=0
             _a    : int     					={}






      def __UpperCAmelCase       ( self			:Any     ,						SCREAMING_SNAKE_CASE			:List[str]	)       ->					Optional[int]:




             '''simple docstring'''
             if vertex not in self.adjacency:
                    _a    : Dict     					={}
                    self.num_vertices += 1






      def __UpperCAmelCase       ( self			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:str     ,						SCREAMING_SNAKE_CASE			:str     ,						SCREAMING_SNAKE_CASE			:Any	)       ->					List[str]:




             '''simple docstring'''
             self.add_vertex(SCREAMING_SNAKE_CASE	)
             self.add_vertex(SCREAMING_SNAKE_CASE	)

             if head == tail:
                    return

             _a    : Any     					=weight
             _a    : Tuple     					=weight






      def __UpperCAmelCase       ( self			:Dict	)       ->					Optional[int]:




             '''simple docstring'''
             _a    : Union[str, Any]     					=self.get_edges()
             for edge in edges:
                    _a					,    _a					,    _a    : List[str]     					=edge
                    edges.remove((tail, head, weight)	)
             for i in range(len(SCREAMING_SNAKE_CASE	)	):
                    _a    : str     					=list(edges[i]	)

             edges.sort(key=lambda SCREAMING_SNAKE_CASE	: e[2]	)
             for i in range(len(SCREAMING_SNAKE_CASE	) - 1	):
                    if edges[i][2] >= edges[i + 1][2]:
                           _a    : Union[str, Any]     					=edges[i][2] + 1
             for edge in edges:
                    _a					,    _a					,    _a    : Tuple     					=edge
                    _a    : Tuple     					=weight
                    _a    : List[Any]     					=weight






      def __str__( self			:int	)       ->					str:




             '''simple docstring'''
             _a    : int     					=""""""
             for tail in self.adjacency:
                    for head in self.adjacency[tail]:
                           _a    : str     					=self.adjacency[head][tail]
                           string += f"{head} -> {tail} == {weight}\n"
             return string.rstrip("""\n"""	)






      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					Optional[Any]:




             '''simple docstring'''
             _a    : Union[str, Any]     					=[]
             for tail in self.adjacency:
                    for head in self.adjacency[tail]:
                           output.append((tail, head, self.adjacency[head][tail])	)
             return output






      def __UpperCAmelCase       ( self			:List[Any]	)       ->					List[Any]:




             '''simple docstring'''
             return self.adjacency.keys()






      @staticmethod
      def __UpperCAmelCase       ( SCREAMING_SNAKE_CASE			:Dict=None     ,						SCREAMING_SNAKE_CASE			:List[Any]=None	)       ->					Optional[int]:




             '''simple docstring'''
             _a    : str     					=Graph()
             if vertices is None:
                    _a    : Union[str, Any]     					=[]
             if edges is None:
                    _a    : List[Any]     					=[]
             for vertex in vertices:
                    g.add_vertex(SCREAMING_SNAKE_CASE	)
             for edge in edges:
                    g.add_edge(*SCREAMING_SNAKE_CASE	)
             return g







      class A__       :






            def __init__( self			:Optional[int]	)       ->					Union[str, Any]:




                   '''simple docstring'''
                   _a    : Optional[int]     					={}
                   _a    : List[str]     					={}






            def __len__( self			:List[Any]	)       ->					List[Any]:




                   '''simple docstring'''
                   return len(self.parent	)






            def __UpperCAmelCase       ( self			:Tuple     ,						SCREAMING_SNAKE_CASE			:Tuple	)       ->					Dict:




                   '''simple docstring'''
                   if item in self.parent:
                          return self.find(SCREAMING_SNAKE_CASE	)

                   _a    : Optional[Any]     					=item
                   _a    : List[str]     					=0
                   return item






            def __UpperCAmelCase       ( self			:int     ,						SCREAMING_SNAKE_CASE			:Dict	)       ->					List[str]:




                   '''simple docstring'''
                   if item not in self.parent:
                          return self.make_set(SCREAMING_SNAKE_CASE	)
                   if item != self.parent[item]:
                          _a    : str     					=self.find(self.parent[item]	)
                   return self.parent[item]






            def __UpperCAmelCase       ( self			:Optional[int]     ,						SCREAMING_SNAKE_CASE			:Tuple     ,						SCREAMING_SNAKE_CASE			:List[Any]	)       ->					Optional[Any]:




                   '''simple docstring'''
                   _a    : Optional[int]     					=self.find(SCREAMING_SNAKE_CASE	)
                   _a    : Dict     					=self.find(SCREAMING_SNAKE_CASE	)

                   if roota == roota:
                          return roota

                   if self.rank[roota] > self.rank[roota]:
                          _a    : Any     					=roota
                          return roota

                   if self.rank[roota] < self.rank[roota]:
                          _a    : List[str]     					=roota
                          return roota

                   if self.rank[roota] == self.rank[roota]:
                          self.rank[roota] += 1
                          _a    : List[Any]     					=roota
                          return roota
                   return None






      @staticmethod
      def __UpperCAmelCase       ( SCREAMING_SNAKE_CASE			:Dict	)       ->					Union[str, Any]:




             '''simple docstring'''
             _a    : Any     					=graph.num_vertices

             _a    : Union[str, Any]     					=Graph.UnionFind()
             _a    : Optional[int]     					=[]
             while num_components > 1:
                    _a    : str     					={}
                    for vertex in graph.get_vertices():
                           _a    : List[str]     					=-1

                    _a    : Any     					=graph.get_edges()
                    for edge in edges:
                           _a					,    _a					,    _a    : Tuple     					=edge
                           edges.remove((tail, head, weight)	)
                    for edge in edges:
                           _a					,    _a					,    _a    : Any     					=edge
                           _a    : Any     					=union_find.find(SCREAMING_SNAKE_CASE	)
                           _a    : List[Any]     					=union_find.find(SCREAMING_SNAKE_CASE	)
                           if seta != seta:
                                  if cheap_edge[seta] == -1 or cheap_edge[seta][2] > weight:
                                         _a    : Optional[int]     					=[head, tail, weight]

                                  if cheap_edge[seta] == -1 or cheap_edge[seta][2] > weight:
                                         _a    : List[Any]     					=[head, tail, weight]
                    for vertex in cheap_edge:
                           if cheap_edge[vertex] != -1:
                                  _a					,    _a					,    _a    : Optional[Any]     					=cheap_edge[vertex]
                                  if union_find.find(SCREAMING_SNAKE_CASE	) != union_find.find(SCREAMING_SNAKE_CASE	):
                                         union_find.union(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)
                                         mst_edges.append(cheap_edge[vertex]	)
                                         _a    : str     					=num_components - 1
             _a    : str     					=Graph.build(edges=SCREAMING_SNAKE_CASE	)
             return mst


