
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      __A(    unittest.TestCase			):

  """simple docstring"""
  SCREAMING_SNAKE_CASE__         =      MODEL_FOR_MASKED_LM_MAPPING
  SCREAMING_SNAKE_CASE__         =      TF_MODEL_FOR_MASKED_LM_MAPPING







  def 	UpperCAmelCase_ (self						):
       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 	UpperCAmelCase_ (self						):
       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""": 3_80_15, """token_str""": """ grouped"""},
               {"""sequence""": """My name is accuser""", """score""": 2.1E-05, """token""": 2_55_06, """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""": 3_80_15,
                   """token_str""": """ grouped""",
               },
               {
                   """sequence""": """The largest city in France is accuser""",
                   """score""": 2.1E-05,
                   """token""": 2_55_06,
                   """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""": 1_36_06, """token_str""": """ Clara"""},
               {"""sequence""": """My name is Patrick""", """score""": 2E-05, """token""": 34_99, """token_str""": """ Patrick"""},
               {"""sequence""": """My name is Te""", """score""": 1.9E-05, """token""": 29_41, """token_str""": """ Te"""},
           ]	,      )







  @require_torch
  def 	UpperCAmelCase_ (self						):
       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""": 3_56_76, """token_str""": """ Maul"""},
               {"""sequence""": """My name isELS""", """score""": 2.2E-05, """token""": 1_64_16, """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""": 3_56_76,
                   """token_str""": """ Maul""",
               },
               {"""sequence""": """The largest city in France isELS""", """score""": 2.2E-05, """token""": 1_64_16, """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""": 34_99, """token_str""": """ Patrick"""},
               {"""sequence""": """My name is Te""", """score""": 2E-05, """token""": 29_41, """token_str""": """ Te"""},
               {"""sequence""": """My name is Clara""", """score""": 2E-05, """token""": 1_36_06, """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""": 3_56_76,
                       """token_str""": """ Maul""",
                       """sequence""": """<s>My name is Maul<mask></s>""",
                   },
                   {"""score""": 2.2E-05, """token""": 1_64_16, """token_str""": """ELS""", """sequence""": """<s>My name isELS<mask></s>"""},
               ],
               [
                   {
                       """score""": 2.2E-05,
                       """token""": 3_56_76,
                       """token_str""": """ Maul""",
                       """sequence""": """<s>My name is<mask> Maul</s>""",
                   },
                   {"""score""": 2.2E-05, """token""": 1_64_16, """token_str""": """ELS""", """sequence""": """<s>My name is<mask>ELS</s>"""},
               ],
           ]	,      )







  @require_torch_gpu
  def 	UpperCAmelCase_ (self						):
       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 	UpperCAmelCase_ (self						):
       UpperCamelCase__				=    pipeline(task="""fill-mask"""	,      model="""distilroberta-base"""	,      top_k=2	,      framework="""pt"""						)
       self.run_large_test(SCREAMING_SNAKE_CASE_						)







  @slow
  @require_tf
  def 	UpperCAmelCase_ (self						):
       UpperCamelCase__				=    pipeline(task="""fill-mask"""	,      model="""distilroberta-base"""	,      top_k=2	,      framework="""tf"""						)
       self.run_large_test(SCREAMING_SNAKE_CASE_						)







  def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_						):
       UpperCamelCase__				=    unmasker("""My name is <mask>"""						)
       self.assertEqual(
           nested_simplify(SCREAMING_SNAKE_CASE_						)	,      [
               {"""sequence""": """My name is John""", """score""": 0.008, """token""": 6_10, """token_str""": """ John"""},
               {"""sequence""": """My name is Chris""", """score""": 0.007, """token""": 15_73, """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.251,
                   """token""": 22_01,
                   """token_str""": """ Paris""",
               },
               {
                   """sequence""": """The largest city in France is Lyon""",
                   """score""": 0.214,
                   """token""": 1_27_90,
                   """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.005, """token""": 34_99, """token_str""": """ Patrick"""},
               {"""sequence""": """My name is Clara""", """score""": 0.000, """token""": 1_36_06, """token_str""": """ Clara"""},
               {"""sequence""": """My name is Te""", """score""": 0.000, """token""": 29_41, """token_str""": """ Te"""},
           ]	,      )







  @require_torch
  def 	UpperCAmelCase_ (self						):
       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 	UpperCAmelCase_ (self						):
       UpperCamelCase__				=    pipeline(task="""fill-mask"""	,      model="""sshleifer/tiny-distilroberta-base"""	,      framework="""tf"""						)
       UpperCamelCase__				=    None
       UpperCamelCase__				=    None
       self.run_pipeline_test(SCREAMING_SNAKE_CASE_	,      []						)







  def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						):
       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 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						):
       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 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						):
       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 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						):
       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 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						):
       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 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						):
       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 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						):
       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_						)},
               ],
           ]	,      )







import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


lowerCamelCase_						     =					logging.get_logger(__name__)






def 						__magic_name__			(   __a      :	Optional[int]						):



     '''simple docstring'''




     UpperCamelCase__				=    R"""\w+[.]\d+"""
     UpperCamelCase__				=    re.findall(__a	,		__a						)
     for pat in pats:
          UpperCamelCase__				=    key.replace(__a	,		"""_""".join(pat.split("""."""						)						)						)
     return key






def 						__magic_name__			(   __a      :	str	,		__a      :	Dict	,		__a      :	int						):



     '''simple docstring'''




     UpperCamelCase__				=    pt_tuple_key[:-1] + ("""scale""",)
     if (
         any("""norm""" in str_ for str_ in pt_tuple_key						)
         and (pt_tuple_key[-1] == "bias")
         and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
         and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
     ):
          UpperCamelCase__				=    pt_tuple_key[:-1] + ("""scale""",)
          return renamed_pt_tuple_key, pt_tensor
     elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
          UpperCamelCase__				=    pt_tuple_key[:-1] + ("""scale""",)
          return renamed_pt_tuple_key, pt_tensor

     # embedding
     if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
          UpperCamelCase__				=    pt_tuple_key[:-1] + ("""embedding""",)
          return renamed_pt_tuple_key, pt_tensor

     # conv layer
     UpperCamelCase__				=    pt_tuple_key[:-1] + ("""kernel""",)
     if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
          UpperCamelCase__				=    pt_tensor.transpose(2	,		3	,		1	,		0						)
          return renamed_pt_tuple_key, pt_tensor

     # linear layer
     UpperCamelCase__				=    pt_tuple_key[:-1] + ("""kernel""",)
     if pt_tuple_key[-1] == "weight":
          UpperCamelCase__				=    pt_tensor.T
          return renamed_pt_tuple_key, pt_tensor

     # old PyTorch layer norm weight
     UpperCamelCase__				=    pt_tuple_key[:-1] + ("""weight""",)
     if pt_tuple_key[-1] == "gamma":
          return renamed_pt_tuple_key, pt_tensor

     # old PyTorch layer norm bias
     UpperCamelCase__				=    pt_tuple_key[:-1] + ("""bias""",)
     if pt_tuple_key[-1] == "beta":
          return renamed_pt_tuple_key, pt_tensor

     return pt_tuple_key, pt_tensor






def 						__magic_name__			(   __a      :	List[Any]	,		__a      :	List[Any]	,		__a      :	Optional[int]=42						):



     '''simple docstring'''




     UpperCamelCase__				=    {k: v.numpy() for k, v in pt_state_dict.items()}

     # Step 2: Since the model is stateless, get random Flax params
     UpperCamelCase__				=    flax_model.init_weights(PRNGKey(__a						)						)

     UpperCamelCase__				=    flatten_dict(__a						)
     UpperCamelCase__				=    {}

     # Need to change some parameters name to match Flax names
     for pt_key, pt_tensor in pt_state_dict.items():
          UpperCamelCase__				=    rename_key(__a						)
          UpperCamelCase__				=    tuple(renamed_pt_key.split("""."""						)						)

          # Correctly rename weight parameters
          UpperCamelCase__		,	UpperCamelCase__				=    rename_key_and_reshape_tensor(__a	,		__a	,		__a						)

          if flax_key in random_flax_state_dict:
               if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                    raise ValueError(
                        f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape "
                        f"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."						)

        # also add unexpected weight so that warning is thrown
          UpperCamelCase__				=    jnp.asarray(__a						)

     return unflatten_dict(__a						)





