





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_)},
           ],
       ]       , )




'''simple docstring'''



import unittest

import numpy as np
import torch

from diffusers import PNDMPipeline, PNDMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()







class 						lowerCAmelCase__   (      unittest.TestCase					):



 """simple docstring"""



 @property
 def       UpperCAmelCase__		(		self					:       List[str]			)  -> Optional[int]:


      """simple docstring"""




      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      UNetaDModel(
          block_out_channels=(32, 64)					,							layers_per_block=2					,							sample_size=32					,							in_channels=3					,							out_channels=3					,							down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")					,							up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")					,							)
      return model


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


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.dummy_uncond_unet
      __SCREAMING_SNAKE_CASE     =      PNDMScheduler()

      __SCREAMING_SNAKE_CASE     =      PNDMPipeline(unet=__SCREAMING_SNAKE_CASE					,							scheduler=__SCREAMING_SNAKE_CASE			)
      pndm.to(__SCREAMING_SNAKE_CASE			)
      pndm.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pndm(generator=__SCREAMING_SNAKE_CASE					,							num_inference_steps=20					,							output_type="""numpy"""			).images

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pndm(generator=__SCREAMING_SNAKE_CASE					,							num_inference_steps=20					,							output_type="""numpy"""					,							return_dict=__SCREAMING_SNAKE_CASE			)[0]

      __SCREAMING_SNAKE_CASE     =      image[0, -3:, -3:, -1]
      __SCREAMING_SNAKE_CASE     =      image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      __SCREAMING_SNAKE_CASE     =      np.array([1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0]			)

      assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2
      assert np.abs(image_from_tuple_slice.flatten() - expected_slice			).max() < 1E-2







@slow
@require_torch
class 						lowerCAmelCase__   (      unittest.TestCase					):



 """simple docstring"""



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


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      """google/ddpm-cifar10-32"""

      __SCREAMING_SNAKE_CASE     =      UNetaDModel.from_pretrained(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      PNDMScheduler()

      __SCREAMING_SNAKE_CASE     =      PNDMPipeline(unet=__SCREAMING_SNAKE_CASE					,							scheduler=__SCREAMING_SNAKE_CASE			)
      pndm.to(__SCREAMING_SNAKE_CASE			)
      pndm.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pndm(generator=__SCREAMING_SNAKE_CASE					,							output_type="""numpy"""			).images

      __SCREAMING_SNAKE_CASE     =      image[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      __SCREAMING_SNAKE_CASE     =      np.array([0.1564, 0.14645, 0.1406, 0.14715, 0.12425, 0.14045, 0.13115, 0.12175, 0.125]			)

      assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2






