





import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY



@is_pipeline_test
class    _SCREAMING_SNAKE_CASE						(	unittest.TestCase    ):







 '''simple docstring'''
 lowercase_		=			MODEL_FOR_MASKED_LM_MAPPING
 lowercase_		=			TF_MODEL_FOR_MASKED_LM_MAPPING

 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->str:



   '''simple docstring'''



   super().tearDown()
   # clean-up as much as possible GPU memory occupied by PyTorch
   gc.collect()
   if is_torch_available():
     import torch

     torch.cuda.empty_cache()

 @require_tf
 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Tuple:



   '''simple docstring'''



   lowerCamelCase__:      Any			  =pipeline(task="fill-mask"       , model="sshleifer/tiny-distilroberta-base"       , top_k=2       , framework="tf")
   lowerCamelCase__:      List[Any]			  =unmasker("My name is <mask>")
   self.assertEqual(
       nested_simplify(UpperCAmelCase_       , decimals=6)       , [
           {"sequence": "My name is grouped", "score": 2.1E-0_5, "token": 38_015, "token_str": " grouped"},
           {"sequence": "My name is accuser", "score": 2.1E-0_5, "token": 25_506, "token_str": " accuser"},
       ]       , )

   lowerCamelCase__:      List[str]			  =unmasker("The largest city in France is <mask>")
   self.assertEqual(
       nested_simplify(UpperCAmelCase_       , decimals=6)       , [
           {
               "sequence": "The largest city in France is grouped",
               "score": 2.1E-0_5,
               "token": 38_015,
               "token_str": " grouped",
           },
           {
               "sequence": "The largest city in France is accuser",
               "score": 2.1E-0_5,
               "token": 25_506,
               "token_str": " accuser",
           },
       ]       , )

   lowerCamelCase__:      List[str]			  =unmasker("My name is <mask>"       , targets=[" Patrick", " Clara", " Teven"]       , top_k=3)
   self.assertEqual(
       nested_simplify(UpperCAmelCase_       , decimals=6)       , [
           {"sequence": "My name is Clara", "score": 2E-0_5, "token": 13_606, "token_str": " Clara"},
           {"sequence": "My name is Patrick", "score": 2E-0_5, "token": 3_499, "token_str": " Patrick"},
           {"sequence": "My name is Te", "score": 1.9E-0_5, "token": 2_941, "token_str": " Te"},
       ]       , )

 @require_torch
 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:      List[str]			  =pipeline(task="fill-mask"       , model="sshleifer/tiny-distilroberta-base"       , top_k=2       , framework="pt")

   lowerCamelCase__:      Tuple			  =unmasker("My name is <mask>")
   self.assertEqual(
       nested_simplify(UpperCAmelCase_       , decimals=6)       , [
           {"sequence": "My name is Maul", "score": 2.2E-0_5, "token": 35_676, "token_str": " Maul"},
           {"sequence": "My name isELS", "score": 2.2E-0_5, "token": 16_416, "token_str": "ELS"},
       ]       , )

   lowerCamelCase__:      str			  =unmasker("The largest city in France is <mask>")
   self.assertEqual(
       nested_simplify(UpperCAmelCase_       , decimals=6)       , [
           {
               "sequence": "The largest city in France is Maul",
               "score": 2.2E-0_5,
               "token": 35_676,
               "token_str": " Maul",
           },
           {"sequence": "The largest city in France isELS", "score": 2.2E-0_5, "token": 16_416, "token_str": "ELS"},
       ]       , )

   lowerCamelCase__:      str			  =unmasker("My name is <mask>"       , targets=[" Patrick", " Clara", " Teven"]       , top_k=3)
   self.assertEqual(
       nested_simplify(UpperCAmelCase_       , decimals=6)       , [
           {"sequence": "My name is Patrick", "score": 2.1E-0_5, "token": 3_499, "token_str": " Patrick"},
           {"sequence": "My name is Te", "score": 2E-0_5, "token": 2_941, "token_str": " Te"},
           {"sequence": "My name is Clara", "score": 2E-0_5, "token": 13_606, "token_str": " Clara"},
       ]       , )

   lowerCamelCase__:      Dict			  =unmasker("My name is <mask> <mask>"       , top_k=2)

   self.assertEqual(
       nested_simplify(UpperCAmelCase_       , decimals=6)       , [
           [
               {
                   "score": 2.2E-0_5,
                   "token": 35_676,
                   "token_str": " Maul",
                   "sequence": "<s>My name is Maul<mask></s>",
               },
               {"score": 2.2E-0_5, "token": 16_416, "token_str": "ELS", "sequence": "<s>My name isELS<mask></s>"},
           ],
           [
               {
                   "score": 2.2E-0_5,
                   "token": 35_676,
                   "token_str": " Maul",
                   "sequence": "<s>My name is<mask> Maul</s>",
               },
               {"score": 2.2E-0_5, "token": 16_416, "token_str": "ELS", "sequence": "<s>My name is<mask>ELS</s>"},
           ],
       ]       , )

 @require_torch_gpu
 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->str:



   '''simple docstring'''



   lowerCamelCase__:      Union[str, Any]			  =pipeline("fill-mask"       , model="hf-internal-testing/tiny-random-distilbert"       , device=0       , framework="pt")

   # convert model to fp16
   pipe.model.half()

   lowerCamelCase__:      str			  =pipe("Paris is the [MASK] of France.")
   # We actually don't care about the result, we just want to make sure
   # it works, meaning the float16 tensor got casted back to float32
   # for postprocessing.
   self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

 @slow
 @require_torch
 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      Any			  =pipeline(task="fill-mask"       , model="distilroberta-base"       , top_k=2       , framework="pt")
   self.run_large_test(UpperCAmelCase_)

 @slow
 @require_tf
 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int])       ->Dict:



   '''simple docstring'''



   lowerCamelCase__:      List[Any]			  =pipeline(task="fill-mask"       , model="distilroberta-base"       , top_k=2       , framework="tf")
   self.run_large_test(UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str]       , UpperCAmelCase_				:   str)       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[int]			  =unmasker("My name is <mask>")
   self.assertEqual(
       nested_simplify(UpperCAmelCase_)       , [
           {"sequence": "My name is John", "score": 0.008, "token": 610, "token_str": " John"},
           {"sequence": "My name is Chris", "score": 0.007, "token": 1_573, "token_str": " Chris"},
       ]       , )
   lowerCamelCase__:      Tuple			  =unmasker("The largest city in France is <mask>")
   self.assertEqual(
       nested_simplify(UpperCAmelCase_)       , [
           {
               "sequence": "The largest city in France is Paris",
               "score": 0.251,
               "token": 2_201,
               "token_str": " Paris",
           },
           {
               "sequence": "The largest city in France is Lyon",
               "score": 0.214,
               "token": 12_790,
               "token_str": " Lyon",
           },
       ]       , )

   lowerCamelCase__:      int			  =unmasker("My name is <mask>"       , targets=[" Patrick", " Clara", " Teven"]       , top_k=3)
   self.assertEqual(
       nested_simplify(UpperCAmelCase_)       , [
           {"sequence": "My name is Patrick", "score": 0.005, "token": 3_499, "token_str": " Patrick"},
           {"sequence": "My name is Clara", "score": 0.000, "token": 13_606, "token_str": " Clara"},
           {"sequence": "My name is Te", "score": 0.000, "token": 2_941, "token_str": " Te"},
       ]       , )

 @require_torch
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =pipeline(task="fill-mask"       , model="sshleifer/tiny-distilroberta-base"       , framework="pt")
   lowerCamelCase__:      Dict			  =None
   lowerCamelCase__:      Dict			  =None
   self.run_pipeline_test(UpperCAmelCase_       , [])

 @require_tf
 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->str:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =pipeline(task="fill-mask"       , model="sshleifer/tiny-distilroberta-base"       , framework="tf")
   lowerCamelCase__:      Any			  =None
   lowerCamelCase__:      Optional[int]			  =None
   self.run_pipeline_test(UpperCAmelCase_       , [])

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any       , UpperCAmelCase_				:   Optional[Any]       , UpperCAmelCase_				:   List[Any]       , UpperCAmelCase_				:   Optional[int])       ->Optional[int]:



   '''simple docstring'''



   if tokenizer is None or tokenizer.mask_token_id is None:
     self.skipTest("The provided tokenizer has no mask token, (probably reformer or wav2vec2)")

   lowerCamelCase__:      List[Any]			  =FillMaskPipeline(model=UpperCAmelCase_       , tokenizer=UpperCAmelCase_)
   lowerCamelCase__:      Optional[int]			  =[
       F"""This is another {tokenizer.mask_token} test""",
   ]
   return fill_masker, examples

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



   '''simple docstring'''



   lowerCamelCase__:      int			  =fill_masker.tokenizer
   lowerCamelCase__:      str			  =fill_masker.model

   lowerCamelCase__:      Optional[int]			  =fill_masker(
       F"""This is a {tokenizer.mask_token}"""       , )
   self.assertEqual(
       UpperCAmelCase_       , [
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
       ]       , )

   lowerCamelCase__:      List[Any]			  =fill_masker([F"""This is a {tokenizer.mask_token}"""])
   self.assertEqual(
       UpperCAmelCase_       , [
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
       ]       , )

   lowerCamelCase__:      Optional[Any]			  =fill_masker([F"""This is a {tokenizer.mask_token}""", F"""Another {tokenizer.mask_token} great test."""])
   self.assertEqual(
       UpperCAmelCase_       , [
           [
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           ],
           [
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           ],
       ]       , )

   with self.assertRaises(UpperCAmelCase_):
     fill_masker([None])
   # No mask_token is not supported
   with self.assertRaises(UpperCAmelCase_):
     fill_masker("This is")

   self.run_test_top_k(UpperCAmelCase_       , UpperCAmelCase_)
   self.run_test_targets(UpperCAmelCase_       , UpperCAmelCase_)
   self.run_test_top_k_targets(UpperCAmelCase_       , UpperCAmelCase_)
   self.fill_mask_with_duplicate_targets_and_top_k(UpperCAmelCase_       , UpperCAmelCase_)
   self.fill_mask_with_multiple_masks(UpperCAmelCase_       , UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   str       , UpperCAmelCase_				:   Dict       , UpperCAmelCase_				:   Optional[int])       ->Tuple:



   '''simple docstring'''



   lowerCamelCase__:      Optional[int]			  =tokenizer.get_vocab()
   lowerCamelCase__:      Optional[int]			  =sorted(vocab.keys())[:2]
   # Pipeline argument
   lowerCamelCase__:      Tuple			  =FillMaskPipeline(model=UpperCAmelCase_       , tokenizer=UpperCAmelCase_       , targets=UpperCAmelCase_)
   lowerCamelCase__:      Optional[Any]			  =fill_masker(F"""This is a {tokenizer.mask_token}""")
   self.assertEqual(
       UpperCAmelCase_       , [
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
       ]       , )
   lowerCamelCase__:      Tuple			  ={vocab[el] for el in targets}
   self.assertEqual({el["token"] for el in outputs}       , UpperCAmelCase_)
   lowerCamelCase__:      List[str]			  =[tokenizer.decode([x]) for x in target_ids]
   self.assertEqual({el["token_str"] for el in outputs}       , set(UpperCAmelCase_))

   # Call argument
   lowerCamelCase__:      List[Any]			  =FillMaskPipeline(model=UpperCAmelCase_       , tokenizer=UpperCAmelCase_)
   lowerCamelCase__:      Optional[Any]			  =fill_masker(F"""This is a {tokenizer.mask_token}"""       , targets=UpperCAmelCase_)
   self.assertEqual(
       UpperCAmelCase_       , [
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
       ]       , )
   lowerCamelCase__:      Optional[int]			  ={vocab[el] for el in targets}
   self.assertEqual({el["token"] for el in outputs}       , UpperCAmelCase_)
   lowerCamelCase__:      str			  =[tokenizer.decode([x]) for x in target_ids]
   self.assertEqual({el["token_str"] for el in outputs}       , set(UpperCAmelCase_))

   # Score equivalence
   lowerCamelCase__:      List[str]			  =fill_masker(F"""This is a {tokenizer.mask_token}"""       , targets=UpperCAmelCase_)
   lowerCamelCase__:      str			  =[top_mask["token_str"] for top_mask in outputs]
   lowerCamelCase__:      Dict			  =[top_mask["score"] for top_mask in outputs]

   # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
   if set(UpperCAmelCase_) == set(UpperCAmelCase_):
     lowerCamelCase__:      Union[str, Any]			  =fill_masker(F"""This is a {tokenizer.mask_token}"""       , targets=UpperCAmelCase_)
     lowerCamelCase__:      Any			  =[top_mask["score"] for top_mask in unmasked_targets]
     self.assertEqual(nested_simplify(UpperCAmelCase_)       , nested_simplify(UpperCAmelCase_))

   # Raises with invalid
   with self.assertRaises(UpperCAmelCase_):
     lowerCamelCase__:      Optional[int]			  =fill_masker(F"""This is a {tokenizer.mask_token}"""       , targets=[])
   # For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
   if "" not in tokenizer.get_vocab():
     with self.assertRaises(UpperCAmelCase_):
       lowerCamelCase__:      Dict			  =fill_masker(F"""This is a {tokenizer.mask_token}"""       , targets=[""])
     with self.assertRaises(UpperCAmelCase_):
       lowerCamelCase__:      Dict			  =fill_masker(F"""This is a {tokenizer.mask_token}"""       , targets="")

 def 							SCREAMING_SNAKE_CASE_					(self				:   Tuple       , UpperCAmelCase_				:   Optional[Any]       , UpperCAmelCase_				:   List[str])       ->str:



   '''simple docstring'''



   lowerCamelCase__:      Optional[int]			  =FillMaskPipeline(model=UpperCAmelCase_       , tokenizer=UpperCAmelCase_       , top_k=2)
   lowerCamelCase__:      List[str]			  =fill_masker(F"""This is a {tokenizer.mask_token}""")
   self.assertEqual(
       UpperCAmelCase_       , [
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
       ]       , )

   lowerCamelCase__:      Dict			  =FillMaskPipeline(model=UpperCAmelCase_       , tokenizer=UpperCAmelCase_)
   lowerCamelCase__:      Tuple			  =fill_masker(F"""This is a {tokenizer.mask_token}"""       , top_k=2)
   self.assertEqual(
       UpperCAmelCase_       , [
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
       ]       , )
   self.assertEqual(nested_simplify(UpperCAmelCase_)       , nested_simplify(UpperCAmelCase_))

 def 							SCREAMING_SNAKE_CASE_					(self				:   Tuple       , UpperCAmelCase_				:   Tuple       , UpperCAmelCase_				:   Any)       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =tokenizer.get_vocab()
   lowerCamelCase__:      Any			  =FillMaskPipeline(model=UpperCAmelCase_       , tokenizer=UpperCAmelCase_)

   # top_k=2, ntargets=3
   lowerCamelCase__:      List[str]			  =sorted(vocab.keys())[:3]
   lowerCamelCase__:      List[Any]			  =fill_masker(F"""This is a {tokenizer.mask_token}"""       , top_k=2       , targets=UpperCAmelCase_)

   # If we use the most probably targets, and filter differently, we should still
   # have the same results
   lowerCamelCase__:      Union[str, Any]			  =[el["token_str"] for el in sorted(UpperCAmelCase_       , key=lambda UpperCAmelCase_: x["score"]       , reverse=UpperCAmelCase_)]
   # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
   if set(UpperCAmelCase_).issubset(UpperCAmelCase_):
     lowerCamelCase__:      int			  =fill_masker(F"""This is a {tokenizer.mask_token}"""       , top_k=3       , targets=UpperCAmelCase_)
     # They should yield exactly the same result
     self.assertEqual(nested_simplify(UpperCAmelCase_)       , nested_simplify(UpperCAmelCase_))

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any]       , UpperCAmelCase_				:   Optional[int]       , UpperCAmelCase_				:   Optional[Any])       ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:      Any			  =FillMaskPipeline(model=UpperCAmelCase_       , tokenizer=UpperCAmelCase_)
   lowerCamelCase__:      Tuple			  =tokenizer.get_vocab()
   # String duplicates + id duplicates
   lowerCamelCase__:      Optional[Any]			  =sorted(vocab.keys())[:3]
   lowerCamelCase__:      Optional[Any]			  =[targets[0], targets[1], targets[0], targets[2], targets[1]]
   lowerCamelCase__:      Union[str, Any]			  =fill_masker(F"""My name is {tokenizer.mask_token}"""       , targets=UpperCAmelCase_       , top_k=10)

   # The target list contains duplicates, so we can't output more
   # than them
   self.assertEqual(len(UpperCAmelCase_)       , 3)







 def 							SCREAMING_SNAKE_CASE_					(self				:   int       , UpperCAmelCase_				:   int       , UpperCAmelCase_				:   str)       ->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =FillMaskPipeline(model=UpperCAmelCase_       , tokenizer=UpperCAmelCase_)

   lowerCamelCase__:      Tuple			  =fill_masker(
       F"""This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}"""       , top_k=2)
   self.assertEqual(
       UpperCAmelCase_       , [
           [
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           ],
           [
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           ],
           [
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
               {"sequence": ANY(UpperCAmelCase_), "score": ANY(UpperCAmelCase_), "token": ANY(UpperCAmelCase_), "token_str": ANY(UpperCAmelCase_)},
           ],
       ]       , )








from collections import defaultdict

from graphs.minimum_spanning_tree_prims import prisms_algorithm as mst





def      lowerCAmelCase_ (       )     ->		Optional[int]:







  """simple docstring"""





  lowerCamelCase__  ,					lowerCamelCase__:      int			  =9, 14  # noqa: F841
  lowerCamelCase__:      List[Any]			  =[
      [0, 1, 4],
      [0, 7, 8],
      [1, 2, 8],
      [7, 8, 7],
      [7, 6, 1],
      [2, 8, 2],
      [8, 6, 6],
      [2, 3, 7],
      [2, 5, 4],
      [6, 5, 2],
      [3, 5, 14],
      [3, 4, 9],
      [5, 4, 10],
      [1, 7, 11],
  ]

  lowerCamelCase__:      List[str]			  =defaultdict(__a		)
  for nodea, nodea, cost in edges:
    adjancency[nodea].append([nodea, cost]		)
    adjancency[nodea].append([nodea, cost]		)

  lowerCamelCase__:      List[str]			  =mst(__a		)

  lowerCamelCase__:      Union[str, Any]			  =[
      [7, 6, 1],
      [2, 8, 2],
      [6, 5, 2],
      [0, 1, 4],
      [2, 5, 4],
      [2, 3, 7],
      [0, 7, 8],
      [3, 4, 9],
  ]

  for answer in expected:
    lowerCamelCase__:      Optional[int]			  =tuple(answer[:2]		)
    lowerCamelCase__:      List[Any]			  =tuple(edge[::-1]		)
    assert edge in result or reverse in result

