


"""simple docstring"""



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 					__magic_name__		(			unittest.TestCase     ):
       __A			:      Optional[Any]        =   MODEL_FOR_MASKED_LM_MAPPING
       __A			:      List[str]        =   TF_MODEL_FOR_MASKED_LM_MAPPING






       def  __snake_case						(   self  :    Optional[int]				):

              '''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  __snake_case						(   self  :    List[str]				):

              '''simple docstring'''






              lowercase		:Tuple       							=				pipeline(task='''fill-mask'''  ,     model='''sshleifer/tiny-distilroberta-base'''  ,     top_k=2  ,     framework='''tf'''				)
              lowercase		:int       							=				unmasker('''My name is <mask>'''				)
              self.assertEqual(
                  nested_simplify(snake_case__  ,     decimals=6				)  ,     [
                      {'''sequence''': '''My name is grouped''', '''score''': 2.1e-0_5, '''token''': 3_8_0_1_5, '''token_str''': ''' grouped'''},
                      {'''sequence''': '''My name is accuser''', '''score''': 2.1e-0_5, '''token''': 2_5_5_0_6, '''token_str''': ''' accuser'''},
                  ]  ,     )

              lowercase		:Dict       							=				unmasker('''The largest city in France is <mask>'''				)
              self.assertEqual(
                  nested_simplify(snake_case__  ,     decimals=6				)  ,     [
                      {
                          '''sequence''': '''The largest city in France is grouped''',
                          '''score''': 2.1e-0_5,
                          '''token''': 3_8_0_1_5,
                          '''token_str''': ''' grouped''',
                      },
                      {
                          '''sequence''': '''The largest city in France is accuser''',
                          '''score''': 2.1e-0_5,
                          '''token''': 2_5_5_0_6,
                          '''token_str''': ''' accuser''',
                      },
                  ]  ,     )

              lowercase		:List[Any]       							=				unmasker('''My name is <mask>'''  ,     targets=[''' Patrick''', ''' Clara''', ''' Teven''']  ,     top_k=3				)
              self.assertEqual(
                  nested_simplify(snake_case__  ,     decimals=6				)  ,     [
                      {'''sequence''': '''My name is Clara''', '''score''': 2e-0_5, '''token''': 1_3_6_0_6, '''token_str''': ''' Clara'''},
                      {'''sequence''': '''My name is Patrick''', '''score''': 2e-0_5, '''token''': 3_4_9_9, '''token_str''': ''' Patrick'''},
                      {'''sequence''': '''My name is Te''', '''score''': 1.9e-0_5, '''token''': 2_9_4_1, '''token_str''': ''' Te'''},
                  ]  ,     )






       @require_torch
       def  __snake_case						(   self  :    int				):

              '''simple docstring'''






              lowercase		:Any       							=				pipeline(task='''fill-mask'''  ,     model='''sshleifer/tiny-distilroberta-base'''  ,     top_k=2  ,     framework='''pt'''				)

              lowercase		:Optional[Any]       							=				unmasker('''My name is <mask>'''				)
              self.assertEqual(
                  nested_simplify(snake_case__  ,     decimals=6				)  ,     [
                      {'''sequence''': '''My name is Maul''', '''score''': 2.2e-0_5, '''token''': 3_5_6_7_6, '''token_str''': ''' Maul'''},
                      {'''sequence''': '''My name isELS''', '''score''': 2.2e-0_5, '''token''': 1_6_4_1_6, '''token_str''': '''ELS'''},
                  ]  ,     )

              lowercase		:List[Any]       							=				unmasker('''The largest city in France is <mask>'''				)
              self.assertEqual(
                  nested_simplify(snake_case__  ,     decimals=6				)  ,     [
                      {
                          '''sequence''': '''The largest city in France is Maul''',
                          '''score''': 2.2e-0_5,
                          '''token''': 3_5_6_7_6,
                          '''token_str''': ''' Maul''',
                      },
                      {'''sequence''': '''The largest city in France isELS''', '''score''': 2.2e-0_5, '''token''': 1_6_4_1_6, '''token_str''': '''ELS'''},
                  ]  ,     )

              lowercase		:Any       							=				unmasker('''My name is <mask>'''  ,     targets=[''' Patrick''', ''' Clara''', ''' Teven''']  ,     top_k=3				)
              self.assertEqual(
                  nested_simplify(snake_case__  ,     decimals=6				)  ,     [
                      {'''sequence''': '''My name is Patrick''', '''score''': 2.1e-0_5, '''token''': 3_4_9_9, '''token_str''': ''' Patrick'''},
                      {'''sequence''': '''My name is Te''', '''score''': 2e-0_5, '''token''': 2_9_4_1, '''token_str''': ''' Te'''},
                      {'''sequence''': '''My name is Clara''', '''score''': 2e-0_5, '''token''': 1_3_6_0_6, '''token_str''': ''' Clara'''},
                  ]  ,     )

              lowercase		:str       							=				unmasker('''My name is <mask> <mask>'''  ,     top_k=2				)

              self.assertEqual(
                  nested_simplify(snake_case__  ,     decimals=6				)  ,     [
                      [
                          {
                              '''score''': 2.2e-0_5,
                              '''token''': 3_5_6_7_6,
                              '''token_str''': ''' Maul''',
                              '''sequence''': '''<s>My name is Maul<mask></s>''',
                          },
                          {'''score''': 2.2e-0_5, '''token''': 1_6_4_1_6, '''token_str''': '''ELS''', '''sequence''': '''<s>My name isELS<mask></s>'''},
                      ],
                      [
                          {
                              '''score''': 2.2e-0_5,
                              '''token''': 3_5_6_7_6,
                              '''token_str''': ''' Maul''',
                              '''sequence''': '''<s>My name is<mask> Maul</s>''',
                          },
                          {'''score''': 2.2e-0_5, '''token''': 1_6_4_1_6, '''token_str''': '''ELS''', '''sequence''': '''<s>My name is<mask>ELS</s>'''},
                      ],
                  ]  ,     )






       @require_torch_gpu
       def  __snake_case						(   self  :    List[str]				):

              '''simple docstring'''






              lowercase		:List[str]       							=				pipeline('''fill-mask'''  ,     model='''hf-internal-testing/tiny-random-distilbert'''  ,     device=0  ,     framework='''pt'''				)

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

              lowercase		:List[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(snake_case__  ,     snake_case__				)






       @slow
       @require_torch
       def  __snake_case						(   self  :    Any				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				pipeline(task='''fill-mask'''  ,     model='''distilroberta-base'''  ,     top_k=2  ,     framework='''pt'''				)
              self.run_large_test(snake_case__				)






       @slow
       @require_tf
       def  __snake_case						(   self  :    Tuple				):

              '''simple docstring'''






              lowercase		:Tuple       							=				pipeline(task='''fill-mask'''  ,     model='''distilroberta-base'''  ,     top_k=2  ,     framework='''tf'''				)
              self.run_large_test(snake_case__				)






       def  __snake_case						(   self  :    List[str]  ,     snake_case__  :    int				):

              '''simple docstring'''






              lowercase		:int       							=				unmasker('''My name is <mask>'''				)
              self.assertEqual(
                  nested_simplify(snake_case__				)  ,     [
                      {'''sequence''': '''My name is John''', '''score''': 0.0_08, '''token''': 6_1_0, '''token_str''': ''' John'''},
                      {'''sequence''': '''My name is Chris''', '''score''': 0.0_07, '''token''': 1_5_7_3, '''token_str''': ''' Chris'''},
                  ]  ,     )
              lowercase		:Optional[Any]       							=				unmasker('''The largest city in France is <mask>'''				)
              self.assertEqual(
                  nested_simplify(snake_case__				)  ,     [
                      {
                          '''sequence''': '''The largest city in France is Paris''',
                          '''score''': 0.2_51,
                          '''token''': 2_2_0_1,
                          '''token_str''': ''' Paris''',
                      },
                      {
                          '''sequence''': '''The largest city in France is Lyon''',
                          '''score''': 0.2_14,
                          '''token''': 1_2_7_9_0,
                          '''token_str''': ''' Lyon''',
                      },
                  ]  ,     )

              lowercase		:Any       							=				unmasker('''My name is <mask>'''  ,     targets=[''' Patrick''', ''' Clara''', ''' Teven''']  ,     top_k=3				)
              self.assertEqual(
                  nested_simplify(snake_case__				)  ,     [
                      {'''sequence''': '''My name is Patrick''', '''score''': 0.0_05, '''token''': 3_4_9_9, '''token_str''': ''' Patrick'''},
                      {'''sequence''': '''My name is Clara''', '''score''': 0.0_00, '''token''': 1_3_6_0_6, '''token_str''': ''' Clara'''},
                      {'''sequence''': '''My name is Te''', '''score''': 0.0_00, '''token''': 2_9_4_1, '''token_str''': ''' Te'''},
                  ]  ,     )






       @require_torch
       def  __snake_case						(   self  :    Dict				):

              '''simple docstring'''






              lowercase		:int       							=				pipeline(task='''fill-mask'''  ,     model='''sshleifer/tiny-distilroberta-base'''  ,     framework='''pt'''				)
              lowercase		:int       							=				None
              lowercase		:List[str]       							=				None
              self.run_pipeline_test(snake_case__  ,     []				)






       @require_tf
       def  __snake_case						(   self  :    Tuple				):

              '''simple docstring'''






              lowercase		:int       							=				pipeline(task='''fill-mask'''  ,     model='''sshleifer/tiny-distilroberta-base'''  ,     framework='''tf'''				)
              lowercase		:int       							=				None
              lowercase		:Tuple       							=				None
              self.run_pipeline_test(snake_case__  ,     []				)






       def  __snake_case						(   self  :    Optional[int]  ,     snake_case__  :    Union[str, Any]  ,     snake_case__  :    Tuple  ,     snake_case__  :    str				):

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

              lowercase		:List[str]       							=				FillMaskPipeline(model=snake_case__  ,     tokenizer=snake_case__				)
              lowercase		:int       							=				[
                  f"""This is another {tokenizer.mask_token} test""",
              ]
              return fill_masker, examples






       def  __snake_case						(   self  :    Any  ,     snake_case__  :    Any  ,     snake_case__  :    Any				):

              '''simple docstring'''






              lowercase		:Tuple       							=				fill_masker.tokenizer
              lowercase		:Dict       							=				fill_masker.model

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

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

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

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

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






       def  __snake_case						(   self  :    List[str]  ,     snake_case__  :    Dict  ,     snake_case__  :    Tuple				):

              '''simple docstring'''






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

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

              # Score equivalence
              lowercase		:Optional[Any]       							=				fill_masker(f"""This is a {tokenizer.mask_token}"""  ,     targets=snake_case__				)
              lowercase		:Union[str, Any]       							=				[top_mask['''token_str'''] for top_mask in outputs]
              lowercase		:Optional[int]       							=				[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(snake_case__				) == set(snake_case__				):
                     lowercase		:Tuple       							=				fill_masker(f"""This is a {tokenizer.mask_token}"""  ,     targets=snake_case__				)
                     lowercase		:List[Any]       							=				[top_mask['''score'''] for top_mask in unmasked_targets]
                     self.assertEqual(nested_simplify(snake_case__				)  ,     nested_simplify(snake_case__				)				)

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






       def  __snake_case						(   self  :    str  ,     snake_case__  :    List[str]  ,     snake_case__  :    List[str]				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				FillMaskPipeline(model=snake_case__  ,     tokenizer=snake_case__  ,     top_k=2				)
              lowercase		:List[str]       							=				fill_masker(f"""This is a {tokenizer.mask_token}"""				)
              self.assertEqual(
                  snake_case__  ,     [
                      {'''sequence''': ANY(snake_case__				), '''score''': ANY(snake_case__				), '''token''': ANY(snake_case__				), '''token_str''': ANY(snake_case__				)},
                      {'''sequence''': ANY(snake_case__				), '''score''': ANY(snake_case__				), '''token''': ANY(snake_case__				), '''token_str''': ANY(snake_case__				)},
                  ]  ,     )

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






       def  __snake_case						(   self  :    List[Any]  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    Union[str, Any]				):

              '''simple docstring'''






              lowercase		:List[str]       							=				tokenizer.get_vocab()
              lowercase		:List[Any]       							=				FillMaskPipeline(model=snake_case__  ,     tokenizer=snake_case__				)

              # top_k=2, ntargets=3
              lowercase		:Optional[Any]       							=				sorted(vocab.keys()				)[:3]
              lowercase		:List[Any]       							=				fill_masker(f"""This is a {tokenizer.mask_token}"""  ,     top_k=2  ,     targets=snake_case__				)

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






       def  __snake_case						(   self  :    str  ,     snake_case__  :    str  ,     snake_case__  :    Any				):

              '''simple docstring'''






              lowercase		:List[Any]       							=				FillMaskPipeline(model=snake_case__  ,     tokenizer=snake_case__				)
              lowercase		:List[str]       							=				tokenizer.get_vocab()
              # String duplicates + id duplicates
              lowercase		:List[Any]       							=				sorted(vocab.keys()				)[:3]
              lowercase		:List[Any]       							=				[targets[0], targets[1], targets[0], targets[2], targets[1]]
              lowercase		:Union[str, Any]       							=				fill_masker(f"""My name is {tokenizer.mask_token}"""  ,     targets=snake_case__  ,     top_k=1_0				)

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







       def  __snake_case						(   self  :    Dict  ,     snake_case__  :    List[str]  ,     snake_case__  :    List[str]				):

              '''simple docstring'''






              lowercase		:List[str]       							=				FillMaskPipeline(model=snake_case__  ,     tokenizer=snake_case__				)

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





"""simple docstring"""



import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
 import jax
@require_flax
class 					__magic_name__		(			__UpperCAmelCase			,						unittest.TestCase     ):
       __A			:      str        =   FlaxAutoencoderKL






       @property
       def  __snake_case						(   self  :    List[str]				):

              '''simple docstring'''






              lowercase		:Dict       							=				4
              lowercase		:Optional[int]       							=				3
              lowercase		:Any       							=				(3_2, 3_2)

              lowercase		:Optional[int]       							=				jax.random.PRNGKey(0				)
              lowercase		:Any       							=				jax.random.uniform(snake_case__  ,     ((batch_size, num_channels) + sizes)				)

              return {"sample": image, "prng_key": prng_key}







       def  __snake_case						(   self  :    Union[str, Any]				):

              '''simple docstring'''






              lowercase		:List[Any]       							=				{
                  '''block_out_channels''': [3_2, 6_4],
                  '''in_channels''': 3,
                  '''out_channels''': 3,
                  '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
                  '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
                  '''latent_channels''': 4,
              }
              lowercase		:Optional[int]       							=				self.dummy_input
              return init_dict, inputs_dict

