





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 							_UpperCamelCase					(	unittest.TestCase     ):
       _UpperCamelCase				:					Any         =   MODEL_FOR_MASKED_LM_MAPPING
       _UpperCamelCase				:					int         =   TF_MODEL_FOR_MASKED_LM_MAPPING

       def lowercase (		self: int		)       -> List[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 lowercase (		self: Optional[Any]		)       -> Tuple:




              """simple docstring"""


              UpperCamelCase_          =  pipeline(task="fill-mask"					,   model="sshleifer/tiny-distilroberta-base"					,   top_k=2					,   framework="tf"		)
              UpperCamelCase_          =  unmasker("My name is <mask>"		)
              self.assertEqual(
                  nested_simplify(_SCREAMING_SNAKE_CASE					,   decimals=6		)					,   [
                      {"sequence": "My name is grouped", "score": 2.1e-05, "token": 38015, "token_str": " grouped"},
                      {"sequence": "My name is accuser", "score": 2.1e-05, "token": 25506, "token_str": " accuser"},
                  ]					,   )

              UpperCamelCase_          =  unmasker("The largest city in France is <mask>"		)
              self.assertEqual(
                  nested_simplify(_SCREAMING_SNAKE_CASE					,   decimals=6		)					,   [
                      {
                          "sequence": "The largest city in France is grouped",
                          "score": 2.1e-05,
                          "token": 38015,
                          "token_str": " grouped",
                      },
                      {
                          "sequence": "The largest city in France is accuser",
                          "score": 2.1e-05,
                          "token": 25506,
                          "token_str": " accuser",
                      },
                  ]					,   )

              UpperCamelCase_          =  unmasker("My name is <mask>"					,   targets=[" Patrick", " Clara", " Teven"]					,   top_k=3		)
              self.assertEqual(
                  nested_simplify(_SCREAMING_SNAKE_CASE					,   decimals=6		)					,   [
                      {"sequence": "My name is Clara", "score": 2e-05, "token": 13606, "token_str": " Clara"},
                      {"sequence": "My name is Patrick", "score": 2e-05, "token": 3499, "token_str": " Patrick"},
                      {"sequence": "My name is Te", "score": 1.9e-05, "token": 2941, "token_str": " Te"},
                  ]					,   )

       @require_torch
       def lowercase (		self: Tuple		)       -> List[str]:




              """simple docstring"""


              UpperCamelCase_          =  pipeline(task="fill-mask"					,   model="sshleifer/tiny-distilroberta-base"					,   top_k=2					,   framework="pt"		)

              UpperCamelCase_          =  unmasker("My name is <mask>"		)
              self.assertEqual(
                  nested_simplify(_SCREAMING_SNAKE_CASE					,   decimals=6		)					,   [
                      {"sequence": "My name is Maul", "score": 2.2e-05, "token": 35676, "token_str": " Maul"},
                      {"sequence": "My name isELS", "score": 2.2e-05, "token": 16416, "token_str": "ELS"},
                  ]					,   )

              UpperCamelCase_          =  unmasker("The largest city in France is <mask>"		)
              self.assertEqual(
                  nested_simplify(_SCREAMING_SNAKE_CASE					,   decimals=6		)					,   [
                      {
                          "sequence": "The largest city in France is Maul",
                          "score": 2.2e-05,
                          "token": 35676,
                          "token_str": " Maul",
                      },
                      {"sequence": "The largest city in France isELS", "score": 2.2e-05, "token": 16416, "token_str": "ELS"},
                  ]					,   )

              UpperCamelCase_          =  unmasker("My name is <mask>"					,   targets=[" Patrick", " Clara", " Teven"]					,   top_k=3		)
              self.assertEqual(
                  nested_simplify(_SCREAMING_SNAKE_CASE					,   decimals=6		)					,   [
                      {"sequence": "My name is Patrick", "score": 2.1e-05, "token": 3499, "token_str": " Patrick"},
                      {"sequence": "My name is Te", "score": 2e-05, "token": 2941, "token_str": " Te"},
                      {"sequence": "My name is Clara", "score": 2e-05, "token": 13606, "token_str": " Clara"},
                  ]					,   )

              UpperCamelCase_          =  unmasker("My name is <mask> <mask>"					,   top_k=2		)

              self.assertEqual(
                  nested_simplify(_SCREAMING_SNAKE_CASE					,   decimals=6		)					,   [
                      [
                          {
                              "score": 2.2e-05,
                              "token": 35676,
                              "token_str": " Maul",
                              "sequence": "<s>My name is Maul<mask></s>",
                          },
                          {"score": 2.2e-05, "token": 16416, "token_str": "ELS", "sequence": "<s>My name isELS<mask></s>"},
                      ],
                      [
                          {
                              "score": 2.2e-05,
                              "token": 35676,
                              "token_str": " Maul",
                              "sequence": "<s>My name is<mask> Maul</s>",
                          },
                          {"score": 2.2e-05, "token": 16416, "token_str": "ELS", "sequence": "<s>My name is<mask>ELS</s>"},
                      ],
                  ]					,   )

       @require_torch_gpu
       def lowercase (		self: Tuple		)       -> Any:




              """simple docstring"""


              UpperCamelCase_          =  pipeline("fill-mask"					,   model="hf-internal-testing/tiny-random-distilbert"					,   device=0					,   framework="pt"		)

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

              UpperCamelCase_          =  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(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

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




              """simple docstring"""


              UpperCamelCase_          =  pipeline(task="fill-mask"					,   model="distilroberta-base"					,   top_k=2					,   framework="pt"		)
              self.run_large_test(_SCREAMING_SNAKE_CASE		)

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




              """simple docstring"""


              UpperCamelCase_          =  pipeline(task="fill-mask"					,   model="distilroberta-base"					,   top_k=2					,   framework="tf"		)
              self.run_large_test(_SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Any					,   _SCREAMING_SNAKE_CASE: Optional[Any]		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  unmasker("My name is <mask>"		)
              self.assertEqual(
                  nested_simplify(_SCREAMING_SNAKE_CASE		)					,   [
                      {"sequence": "My name is John", "score": 0.0_08, "token": 610, "token_str": " John"},
                      {"sequence": "My name is Chris", "score": 0.0_07, "token": 1573, "token_str": " Chris"},
                  ]					,   )
              UpperCamelCase_          =  unmasker("The largest city in France is <mask>"		)
              self.assertEqual(
                  nested_simplify(_SCREAMING_SNAKE_CASE		)					,   [
                      {
                          "sequence": "The largest city in France is Paris",
                          "score": 0.2_51,
                          "token": 2201,
                          "token_str": " Paris",
                      },
                      {
                          "sequence": "The largest city in France is Lyon",
                          "score": 0.2_14,
                          "token": 12790,
                          "token_str": " Lyon",
                      },
                  ]					,   )

              UpperCamelCase_          =  unmasker("My name is <mask>"					,   targets=[" Patrick", " Clara", " Teven"]					,   top_k=3		)
              self.assertEqual(
                  nested_simplify(_SCREAMING_SNAKE_CASE		)					,   [
                      {"sequence": "My name is Patrick", "score": 0.0_05, "token": 3499, "token_str": " Patrick"},
                      {"sequence": "My name is Clara", "score": 0.0_00, "token": 13606, "token_str": " Clara"},
                      {"sequence": "My name is Te", "score": 0.0_00, "token": 2941, "token_str": " Te"},
                  ]					,   )

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




              """simple docstring"""


              UpperCamelCase_          =  pipeline(task="fill-mask"					,   model="sshleifer/tiny-distilroberta-base"					,   framework="pt"		)
              UpperCamelCase_          =  None
              UpperCamelCase_          =  None
              self.run_pipeline_test(_SCREAMING_SNAKE_CASE					,   []		)

       @require_tf
       def lowercase (		self: Union[str, Any]		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  pipeline(task="fill-mask"					,   model="sshleifer/tiny-distilroberta-base"					,   framework="tf"		)
              UpperCamelCase_          =  None
              UpperCamelCase_          =  None
              self.run_pipeline_test(_SCREAMING_SNAKE_CASE					,   []		)

       def lowercase (		self: Dict					,   _SCREAMING_SNAKE_CASE: Dict					,   _SCREAMING_SNAKE_CASE: Dict					,   _SCREAMING_SNAKE_CASE: Union[str, Any]		)       -> Any:




              """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)"		)

              UpperCamelCase_          =  FillMaskPipeline(model=_SCREAMING_SNAKE_CASE					,   tokenizer=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  [
                  f'''This is another {tokenizer.mask_token} test''',
              ]
              return fill_masker, examples

       def lowercase (		self: str					,   _SCREAMING_SNAKE_CASE: Any					,   _SCREAMING_SNAKE_CASE: str		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  fill_masker.tokenizer
              UpperCamelCase_          =  fill_masker.model

              UpperCamelCase_          =  fill_masker(
                  f'''This is a {tokenizer.mask_token}'''					,   )
              self.assertEqual(
                  _SCREAMING_SNAKE_CASE					,   [
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                  ]					,   )

              UpperCamelCase_          =  fill_masker([f'''This is a {tokenizer.mask_token}''']		)
              self.assertEqual(
                  _SCREAMING_SNAKE_CASE					,   [
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                  ]					,   )

              UpperCamelCase_          =  fill_masker([f'''This is a {tokenizer.mask_token}''', f'''Another {tokenizer.mask_token} great test.''']		)
              self.assertEqual(
                  _SCREAMING_SNAKE_CASE					,   [
                      [
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      ],
                      [
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      ],
                  ]					,   )

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

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

       def lowercase (		self: Union[str, Any]					,   _SCREAMING_SNAKE_CASE: Optional[Any]					,   _SCREAMING_SNAKE_CASE: Tuple		)       -> List[str]:




              """simple docstring"""


              UpperCamelCase_          =  tokenizer.get_vocab()
              UpperCamelCase_          =  sorted(vocab.keys()		)[:2]
              # Pipeline argument
              UpperCamelCase_          =  FillMaskPipeline(model=_SCREAMING_SNAKE_CASE					,   tokenizer=_SCREAMING_SNAKE_CASE					,   targets=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  fill_masker(f'''This is a {tokenizer.mask_token}'''		)
              self.assertEqual(
                  _SCREAMING_SNAKE_CASE					,   [
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                  ]					,   )
              UpperCamelCase_          =  {vocab[el] for el in targets}
              self.assertEqual({el["token"] for el in outputs}					,   _SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  [tokenizer.decode([x]		) for x in target_ids]
              self.assertEqual({el["token_str"] for el in outputs}					,   set(_SCREAMING_SNAKE_CASE		)		)

              # Call argument
              UpperCamelCase_          =  FillMaskPipeline(model=_SCREAMING_SNAKE_CASE					,   tokenizer=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  fill_masker(f'''This is a {tokenizer.mask_token}'''					,   targets=_SCREAMING_SNAKE_CASE		)
              self.assertEqual(
                  _SCREAMING_SNAKE_CASE					,   [
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                  ]					,   )
              UpperCamelCase_          =  {vocab[el] for el in targets}
              self.assertEqual({el["token"] for el in outputs}					,   _SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  [tokenizer.decode([x]		) for x in target_ids]
              self.assertEqual({el["token_str"] for el in outputs}					,   set(_SCREAMING_SNAKE_CASE		)		)

              # Score equivalence
              UpperCamelCase_          =  fill_masker(f'''This is a {tokenizer.mask_token}'''					,   targets=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  [top_mask["token_str"] for top_mask in outputs]
              UpperCamelCase_          =  [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(_SCREAMING_SNAKE_CASE		) == set(_SCREAMING_SNAKE_CASE		):
                     UpperCamelCase_          =  fill_masker(f'''This is a {tokenizer.mask_token}'''					,   targets=_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  [top_mask["score"] for top_mask in unmasked_targets]
                     self.assertEqual(nested_simplify(_SCREAMING_SNAKE_CASE		)					,   nested_simplify(_SCREAMING_SNAKE_CASE		)		)

              # Raises with invalid
              with self.assertRaises(_SCREAMING_SNAKE_CASE		):
                     UpperCamelCase_          =  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(_SCREAMING_SNAKE_CASE		):
                            UpperCamelCase_          =  fill_masker(f'''This is a {tokenizer.mask_token}'''					,   targets=[""]		)
                     with self.assertRaises(_SCREAMING_SNAKE_CASE		):
                            UpperCamelCase_          =  fill_masker(f'''This is a {tokenizer.mask_token}'''					,   targets=""		)

       def lowercase (		self: str					,   _SCREAMING_SNAKE_CASE: List[str]					,   _SCREAMING_SNAKE_CASE: Union[str, Any]		)       -> Tuple:




              """simple docstring"""


              UpperCamelCase_          =  FillMaskPipeline(model=_SCREAMING_SNAKE_CASE					,   tokenizer=_SCREAMING_SNAKE_CASE					,   top_k=2		)
              UpperCamelCase_          =  fill_masker(f'''This is a {tokenizer.mask_token}'''		)
              self.assertEqual(
                  _SCREAMING_SNAKE_CASE					,   [
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                  ]					,   )

              UpperCamelCase_          =  FillMaskPipeline(model=_SCREAMING_SNAKE_CASE					,   tokenizer=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  fill_masker(f'''This is a {tokenizer.mask_token}'''					,   top_k=2		)
              self.assertEqual(
                  _SCREAMING_SNAKE_CASE					,   [
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                  ]					,   )
              self.assertEqual(nested_simplify(_SCREAMING_SNAKE_CASE		)					,   nested_simplify(_SCREAMING_SNAKE_CASE		)		)

       def lowercase (		self: Optional[int]					,   _SCREAMING_SNAKE_CASE: Tuple					,   _SCREAMING_SNAKE_CASE: Dict		)       -> Optional[int]:




              """simple docstring"""


              UpperCamelCase_          =  tokenizer.get_vocab()
              UpperCamelCase_          =  FillMaskPipeline(model=_SCREAMING_SNAKE_CASE					,   tokenizer=_SCREAMING_SNAKE_CASE		)

              # top_k=2, ntargets=3
              UpperCamelCase_          =  sorted(vocab.keys()		)[:3]
              UpperCamelCase_          =  fill_masker(f'''This is a {tokenizer.mask_token}'''					,   top_k=2					,   targets=_SCREAMING_SNAKE_CASE		)

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

       def lowercase (		self: Any					,   _SCREAMING_SNAKE_CASE: str					,   _SCREAMING_SNAKE_CASE: int		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  FillMaskPipeline(model=_SCREAMING_SNAKE_CASE					,   tokenizer=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  tokenizer.get_vocab()
              # String duplicates + id duplicates
              UpperCamelCase_          =  sorted(vocab.keys()		)[:3]
              UpperCamelCase_          =  [targets[0], targets[1], targets[0], targets[2], targets[1]]
              UpperCamelCase_          =  fill_masker(f'''My name is {tokenizer.mask_token}'''					,   targets=_SCREAMING_SNAKE_CASE					,   top_k=10		)

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



       def lowercase (		self: str					,   _SCREAMING_SNAKE_CASE: Optional[Any]					,   _SCREAMING_SNAKE_CASE: Optional[Any]		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  FillMaskPipeline(model=_SCREAMING_SNAKE_CASE					,   tokenizer=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  fill_masker(
                  f'''This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}'''					,   top_k=2		)
              self.assertEqual(
                  _SCREAMING_SNAKE_CASE					,   [
                      [
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      ],
                      [
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      ],
                      [
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                          {"sequence": ANY(_SCREAMING_SNAKE_CASE		), "score": ANY(_SCREAMING_SNAKE_CASE		), "token": ANY(_SCREAMING_SNAKE_CASE		), "token_str": ANY(_SCREAMING_SNAKE_CASE		)},
                      ],
                  ]					,   )








from typing import Optional

import pyspark

from .. import Features, NamedSplit
from ..download import DownloadMode
from ..packaged_modules.spark.spark import Spark
from .abc import AbstractDatasetReader



class 							_UpperCamelCase					(	lowerCAmelCase_     ):

       def __init__(		self: List[Any]					,   _SCREAMING_SNAKE_CASE: pyspark.sql.DataFrame					,   _SCREAMING_SNAKE_CASE: Optional[NamedSplit] = None					,   _SCREAMING_SNAKE_CASE: Optional[Features] = None					,   _SCREAMING_SNAKE_CASE: bool = True					,   _SCREAMING_SNAKE_CASE: str = None					,   _SCREAMING_SNAKE_CASE: bool = False					,   _SCREAMING_SNAKE_CASE: str = None					,   _SCREAMING_SNAKE_CASE: bool = True					,   _SCREAMING_SNAKE_CASE: str = "arrow"					,   **_SCREAMING_SNAKE_CASE: Union[str, Any]					,   )       -> Any:




              """simple docstring"""


              super().__init__(
                  split=_SCREAMING_SNAKE_CASE					,   features=_SCREAMING_SNAKE_CASE					,   cache_dir=_SCREAMING_SNAKE_CASE					,   keep_in_memory=_SCREAMING_SNAKE_CASE					,   streaming=_SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE					,   )
              UpperCamelCase_          =  load_from_cache_file
              UpperCamelCase_          =  file_format
              UpperCamelCase_          =  Spark(
                  df=_SCREAMING_SNAKE_CASE					,   features=_SCREAMING_SNAKE_CASE					,   cache_dir=_SCREAMING_SNAKE_CASE					,   working_dir=_SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE					,   )



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




              """simple docstring"""


              if self.streaming:
                     return self.builder.as_streaming_dataset(split=self.split		)
              UpperCamelCase_          =  None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
              self.builder.download_and_prepare(
                  download_mode=_SCREAMING_SNAKE_CASE					,   file_format=self._file_format					,   )
              return self.builder.as_dataset(split=self.split		)

