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






    '''simple docstring'''
    SCREAMING_SNAKE_CASE:				Dict          =    MODEL_FOR_MASKED_LM_MAPPING
    SCREAMING_SNAKE_CASE:				Optional[int]          =    TF_MODEL_FOR_MASKED_LM_MAPPING




    def      _a    (    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      _a    (    self ):
       lowerCAmelCase_:							Union[str, Any]       =	pipeline(task="fill-mask"       ,						model="sshleifer/tiny-distilroberta-base"       ,						top_k=2       ,						framework="tf" )
       lowerCAmelCase_:							Optional[Any]       =	unmasker("My name is <mask>" )
       self.assertEqual(
           nested_simplify(lowerCamelCase__       ,						decimals=6 )       ,						[
               {"sequence": "My name is grouped", "score": 2.1E-05, "token": 38_015, "token_str": " grouped"},
               {"sequence": "My name is accuser", "score": 2.1E-05, "token": 25_506, "token_str": " accuser"},
           ]       ,						)

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

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




    @require_torch
    def      _a    (    self ):
       lowerCAmelCase_:							int       =	pipeline(task="fill-mask"       ,						model="sshleifer/tiny-distilroberta-base"       ,						top_k=2       ,						framework="pt" )

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

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

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

       lowerCAmelCase_:							Union[str, Any]       =	unmasker("My name is <mask> <mask>"       ,						top_k=2 )

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




    @require_torch_gpu
    def      _a    (    self ):
       lowerCAmelCase_:							Tuple       =	pipeline("fill-mask"       ,						model="hf-internal-testing/tiny-random-distilbert"       ,						device=0       ,						framework="pt" )

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

       lowerCAmelCase_:							List[Any]       =	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(lowerCamelCase__       ,						lowerCamelCase__ )




    @slow
    @require_torch
    def      _a    (    self ):
       lowerCAmelCase_:							Dict       =	pipeline(task="fill-mask"       ,						model="distilroberta-base"       ,						top_k=2       ,						framework="pt" )
       self.run_large_test(lowerCamelCase__ )




    @slow
    @require_tf
    def      _a    (    self ):
       lowerCAmelCase_:							List[Any]       =	pipeline(task="fill-mask"       ,						model="distilroberta-base"       ,						top_k=2       ,						framework="tf" )
       self.run_large_test(lowerCamelCase__ )




    def      _a    (    self       ,						lowerCamelCase__ ):
       lowerCAmelCase_:							Optional[int]       =	unmasker("My name is <mask>" )
       self.assertEqual(
           nested_simplify(lowerCamelCase__ )       ,						[
               {"sequence": "My name is John", "score": 0.0_0_8, "token": 610, "token_str": " John"},
               {"sequence": "My name is Chris", "score": 0.0_0_7, "token": 1_573, "token_str": " Chris"},
           ]       ,						)
       lowerCAmelCase_:							str       =	unmasker("The largest city in France is <mask>" )
       self.assertEqual(
           nested_simplify(lowerCamelCase__ )       ,						[
               {
                   "sequence": "The largest city in France is Paris",
                   "score": 0.2_5_1,
                   "token": 2_201,
                   "token_str": " Paris",
               },
               {
                   "sequence": "The largest city in France is Lyon",
                   "score": 0.2_1_4,
                   "token": 12_790,
                   "token_str": " Lyon",
               },
           ]       ,						)

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




    @require_torch
    def      _a    (    self ):
       lowerCAmelCase_:							List[Any]       =	pipeline(task="fill-mask"       ,						model="sshleifer/tiny-distilroberta-base"       ,						framework="pt" )
       lowerCAmelCase_:							str       =	None
       lowerCAmelCase_:							Dict       =	None
       self.run_pipeline_test(lowerCamelCase__       ,						[] )




    @require_tf
    def      _a    (    self ):
       lowerCAmelCase_:							Tuple       =	pipeline(task="fill-mask"       ,						model="sshleifer/tiny-distilroberta-base"       ,						framework="tf" )
       lowerCAmelCase_:							str       =	None
       lowerCAmelCase_:							Dict       =	None
       self.run_pipeline_test(lowerCamelCase__       ,						[] )




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__ ):
       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_:							Union[str, Any]       =	FillMaskPipeline(model=lowerCamelCase__       ,						tokenizer=lowerCamelCase__ )
       lowerCAmelCase_:							Union[str, Any]       =	[
           F'''This is another {tokenizer.mask_token} test''',
       ]
       return fill_masker, examples




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__ ):
       lowerCAmelCase_:							int       =	fill_masker.tokenizer
       lowerCAmelCase_:							str       =	fill_masker.model

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

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

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

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

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




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

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

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

       # Raises with invalid
       with self.assertRaises(lowerCamelCase__ ):
          lowerCAmelCase_:							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(lowerCamelCase__ ):
             lowerCAmelCase_:							Union[str, Any]       =	fill_masker(F'''This is a {tokenizer.mask_token}'''       ,						targets=[""] )
          with self.assertRaises(lowerCamelCase__ ):
             lowerCAmelCase_:							int       =	fill_masker(F'''This is a {tokenizer.mask_token}'''       ,						targets="" )




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__ ):
       lowerCAmelCase_:							Union[str, Any]       =	FillMaskPipeline(model=lowerCamelCase__       ,						tokenizer=lowerCamelCase__       ,						top_k=2 )
       lowerCAmelCase_:							Any       =	fill_masker(F'''This is a {tokenizer.mask_token}''' )
       self.assertEqual(
           lowerCamelCase__       ,						[
               {"sequence": ANY(lowerCamelCase__ ), "score": ANY(lowerCamelCase__ ), "token": ANY(lowerCamelCase__ ), "token_str": ANY(lowerCamelCase__ )},
               {"sequence": ANY(lowerCamelCase__ ), "score": ANY(lowerCamelCase__ ), "token": ANY(lowerCamelCase__ ), "token_str": ANY(lowerCamelCase__ )},
           ]       ,						)

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




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__ ):
       lowerCAmelCase_:							Dict       =	tokenizer.get_vocab()
       lowerCAmelCase_:							Optional[Any]       =	FillMaskPipeline(model=lowerCamelCase__       ,						tokenizer=lowerCamelCase__ )

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

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




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__ ):
       lowerCAmelCase_:							int       =	FillMaskPipeline(model=lowerCamelCase__       ,						tokenizer=lowerCamelCase__ )
       lowerCAmelCase_:							int       =	tokenizer.get_vocab()
       # String duplicates + id duplicates
       lowerCAmelCase_:							List[Any]       =	sorted(vocab.keys() )[:3]
       lowerCAmelCase_:							str       =	[targets[0], targets[1], targets[0], targets[2], targets[1]]
       lowerCAmelCase_:							Tuple       =	fill_masker(F'''My name is {tokenizer.mask_token}'''       ,						targets=lowerCamelCase__       ,						top_k=10 )

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




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__ ):
       lowerCAmelCase_:							Tuple       =	FillMaskPipeline(model=lowerCamelCase__       ,						tokenizer=lowerCamelCase__ )

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

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotSmallConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
 import os

 # The slow tests are often failing with OOM error on GPU
 # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
 # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
 a  :      Tuple            =    """platform"""

 import jax
 import jax.numpy as jnp

 from transformers.models.blenderbot_small.modeling_flax_blenderbot_small import (
     FlaxBlenderbotSmallForConditionalGeneration,
     FlaxBlenderbotSmallModel,
     shift_tokens_right,
 )
def 		snake_case__	(		lowercase							,				lowercase							,				lowercase=None							,				lowercase=None							,				lowercase=None							,				lowercase=None							,				lowercase=None							,				lowercase=None							,				):
   if attention_mask is None:
      lowerCAmelCase_:							Optional[int]       =	np.where(input_ids != config.pad_token_id							,				1							,				0					)
   if decoder_attention_mask is None:
      lowerCAmelCase_:							Union[str, Any]       =	np.where(decoder_input_ids != config.pad_token_id							,				1							,				0					)
   if head_mask is None:
      lowerCAmelCase_:							Optional[int]       =	np.ones((config.encoder_layers, config.encoder_attention_heads)					)
   if decoder_head_mask is None:
      lowerCAmelCase_:							str       =	np.ones((config.decoder_layers, config.decoder_attention_heads)					)
   if cross_attn_head_mask is None:
      lowerCAmelCase_:							Any       =	np.ones((config.decoder_layers, config.decoder_attention_heads)					)
   return {
       "input_ids": input_ids,
       "decoder_input_ids": decoder_input_ids,
       "attention_mask": attention_mask,
       "decoder_attention_mask": attention_mask,
   }

class 					_lowercase    :






    '''simple docstring'''




    def __init__(    self       ,						lowerCamelCase__       ,						lowerCamelCase__=13       ,						lowerCamelCase__=7       ,						lowerCamelCase__=True       ,						lowerCamelCase__=False       ,						lowerCamelCase__=99       ,						lowerCamelCase__=16       ,						lowerCamelCase__=2       ,						lowerCamelCase__=4       ,						lowerCamelCase__=4       ,						lowerCamelCase__="gelu"       ,						lowerCamelCase__=0.1       ,						lowerCamelCase__=0.1       ,						lowerCamelCase__=32       ,						lowerCamelCase__=2       ,						lowerCamelCase__=1       ,						lowerCamelCase__=0       ,						lowerCamelCase__=0.0_2       ,						):
       lowerCAmelCase_:							Union[str, Any]       =	parent
       lowerCAmelCase_:							Tuple       =	batch_size
       lowerCAmelCase_:							Any       =	seq_length
       lowerCAmelCase_:							Tuple       =	is_training
       lowerCAmelCase_:							Optional[int]       =	use_labels
       lowerCAmelCase_:							List[Any]       =	vocab_size
       lowerCAmelCase_:							str       =	hidden_size
       lowerCAmelCase_:							Union[str, Any]       =	num_hidden_layers
       lowerCAmelCase_:							List[str]       =	num_attention_heads
       lowerCAmelCase_:							Dict       =	intermediate_size
       lowerCAmelCase_:							int       =	hidden_act
       lowerCAmelCase_:							Any       =	hidden_dropout_prob
       lowerCAmelCase_:							str       =	attention_probs_dropout_prob
       lowerCAmelCase_:							Union[str, Any]       =	max_position_embeddings
       lowerCAmelCase_:							Any       =	eos_token_id
       lowerCAmelCase_:							Union[str, Any]       =	pad_token_id
       lowerCAmelCase_:							Union[str, Any]       =	bos_token_id
       lowerCAmelCase_:							Optional[int]       =	initializer_range




    def      _a    (    self ):
       lowerCAmelCase_:							List[str]       =	np.clip(ids_tensor([self.batch_size, self.seq_length - 1]       ,						self.vocab_size )       ,						3       ,						self.vocab_size )
       lowerCAmelCase_:							int       =	np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)       ,						dtype=np.intaa ))       ,						-1 )

       lowerCAmelCase_:							Optional[Any]       =	shift_tokens_right(lowerCamelCase__       ,						1       ,						2 )

       lowerCAmelCase_:							Optional[int]       =	BlenderbotSmallConfig(
           vocab_size=self.vocab_size       ,						d_model=self.hidden_size       ,						encoder_layers=self.num_hidden_layers       ,						decoder_layers=self.num_hidden_layers       ,						encoder_attention_heads=self.num_attention_heads       ,						decoder_attention_heads=self.num_attention_heads       ,						encoder_ffn_dim=self.intermediate_size       ,						decoder_ffn_dim=self.intermediate_size       ,						dropout=self.hidden_dropout_prob       ,						attention_dropout=self.attention_probs_dropout_prob       ,						max_position_embeddings=self.max_position_embeddings       ,						eos_token_id=self.eos_token_id       ,						bos_token_id=self.bos_token_id       ,						pad_token_id=self.pad_token_id       ,						initializer_range=self.initializer_range       ,						use_cache=lowerCamelCase__       ,						)
       lowerCAmelCase_:							Optional[Any]       =	prepare_blenderbot_inputs_dict(lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__ )
       return config, inputs_dict




    def      _a    (    self ):
       lowerCAmelCase_  ,	lowerCAmelCase_:							Dict       =	self.prepare_config_and_inputs()
       return config, inputs_dict




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__ ):
       lowerCAmelCase_:							Optional[Any]       =	20
       lowerCAmelCase_:							int       =	model_class_name(lowerCamelCase__ )

       lowerCAmelCase_:							Any       =	model.encode(inputs_dict["input_ids"] )

       lowerCAmelCase_  ,	lowerCAmelCase_:							Optional[int]       =	(
           inputs_dict["decoder_input_ids"],
           inputs_dict["decoder_attention_mask"],
       )

       lowerCAmelCase_:							List[str]       =	model.init_cache(decoder_input_ids.shape[0]       ,						lowerCamelCase__       ,						lowerCamelCase__ )
       lowerCAmelCase_:							Any       =	jnp.ones((decoder_input_ids.shape[0], max_decoder_length)       ,						dtype="i4" )

       lowerCAmelCase_:							Optional[Any]       =	jnp.broadcast_to(
           jnp.arange(decoder_input_ids.shape[-1] - 1 )[None, :]       ,						(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)       ,						)
       lowerCAmelCase_:							Optional[Any]       =	model.decode(
           decoder_input_ids[:, :-1]       ,						lowerCamelCase__       ,						decoder_attention_mask=lowerCamelCase__       ,						past_key_values=lowerCamelCase__       ,						decoder_position_ids=lowerCamelCase__       ,						)

       lowerCAmelCase_:							Optional[Any]       =	jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]       ,						dtype="i4" )
       lowerCAmelCase_:							Optional[int]       =	model.decode(
           decoder_input_ids[:, -1:]       ,						lowerCamelCase__       ,						decoder_attention_mask=lowerCamelCase__       ,						past_key_values=outputs_cache.past_key_values       ,						decoder_position_ids=lowerCamelCase__       ,						)

       lowerCAmelCase_:							Any       =	model.decode(lowerCamelCase__       ,						lowerCamelCase__ )

       lowerCAmelCase_:							str       =	np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
       self.parent.assertTrue(diff < 1E-3       ,						msg=F'''Max diff is {diff}''' )




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__ ):
       lowerCAmelCase_:							str       =	20
       lowerCAmelCase_:							Union[str, Any]       =	model_class_name(lowerCamelCase__ )

       lowerCAmelCase_:							Tuple       =	model.encode(inputs_dict["input_ids"] )

       lowerCAmelCase_  ,	lowerCAmelCase_:							List[str]       =	(
           inputs_dict["decoder_input_ids"],
           inputs_dict["decoder_attention_mask"],
       )

       lowerCAmelCase_:							int       =	jnp.concatenate(
           [
               decoder_attention_mask,
               jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1]) ),
           ]       ,						axis=-1       ,						)

       lowerCAmelCase_:							Any       =	model.init_cache(decoder_input_ids.shape[0]       ,						lowerCamelCase__       ,						lowerCamelCase__ )
       lowerCAmelCase_:							List[Any]       =	jnp.broadcast_to(
           jnp.arange(decoder_input_ids.shape[-1] - 1 )[None, :]       ,						(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)       ,						)

       lowerCAmelCase_:							Tuple       =	model.decode(
           decoder_input_ids[:, :-1]       ,						lowerCamelCase__       ,						decoder_attention_mask=lowerCamelCase__       ,						past_key_values=lowerCamelCase__       ,						decoder_position_ids=lowerCamelCase__       ,						)
       lowerCAmelCase_:							str       =	jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]       ,						dtype="i4" )
       lowerCAmelCase_:							Tuple       =	model.decode(
           decoder_input_ids[:, -1:]       ,						lowerCamelCase__       ,						past_key_values=outputs_cache.past_key_values       ,						decoder_attention_mask=lowerCamelCase__       ,						decoder_position_ids=lowerCamelCase__       ,						)

       lowerCAmelCase_:							Union[str, Any]       =	model.decode(lowerCamelCase__       ,						lowerCamelCase__       ,						decoder_attention_mask=lowerCamelCase__ )

       lowerCAmelCase_:							Optional[int]       =	np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
       self.parent.assertTrue(diff < 1E-3       ,						msg=F'''Max diff is {diff}''' )





@require_flax
class 					_lowercase    (     unittest.TestCase			):






    '''simple docstring'''
    SCREAMING_SNAKE_CASE:				Optional[Any]          =    99




    def      _a    (    self ):
       lowerCAmelCase_:							Optional[int]       =	np.array(
           [
               [71, 82, 18, 33, 46, 91, 2],
               [68, 34, 26, 58, 30, 82, 2],
               [5, 97, 17, 39, 94, 40, 2],
               [76, 83, 94, 25, 70, 78, 2],
               [87, 59, 41, 35, 48, 66, 2],
               [55, 13, 16, 58, 5, 2, 1],  # note padding
               [64, 27, 31, 51, 12, 75, 2],
               [52, 64, 86, 17, 83, 39, 2],
               [48, 61, 9, 24, 71, 82, 2],
               [26, 1, 60, 48, 22, 13, 2],
               [21, 5, 62, 28, 14, 76, 2],
               [45, 98, 37, 86, 59, 48, 2],
               [70, 70, 50, 9, 28, 0, 2],
           ]       ,						dtype=np.intaa       ,						)

       lowerCAmelCase_:							Optional[int]       =	input_ids.shape[0]
       lowerCAmelCase_:							str       =	BlenderbotSmallConfig(
           vocab_size=self.vocab_size       ,						d_model=24       ,						encoder_layers=2       ,						decoder_layers=2       ,						encoder_attention_heads=2       ,						decoder_attention_heads=2       ,						encoder_ffn_dim=32       ,						decoder_ffn_dim=32       ,						max_position_embeddings=48       ,						eos_token_id=2       ,						pad_token_id=1       ,						bos_token_id=0       ,						)
       return config, input_ids, batch_size




    def      _a    (    self ):
       lowerCAmelCase_  ,	lowerCAmelCase_  ,	lowerCAmelCase_:							List[Any]       =	self._get_config_and_data()
       lowerCAmelCase_:							Union[str, Any]       =	FlaxBlenderbotSmallForConditionalGeneration(lowerCamelCase__ )
       lowerCAmelCase_:							Any       =	lm_model(input_ids=lowerCamelCase__ )
       lowerCAmelCase_:							List[str]       =	(batch_size, input_ids.shape[1], config.vocab_size)
       self.assertEqual(outputs["logits"].shape       ,						lowerCamelCase__ )




    def      _a    (    self ):
       lowerCAmelCase_:							Optional[int]       =	BlenderbotSmallConfig(
           vocab_size=self.vocab_size       ,						d_model=14       ,						encoder_layers=2       ,						decoder_layers=2       ,						encoder_attention_heads=2       ,						decoder_attention_heads=2       ,						encoder_ffn_dim=8       ,						decoder_ffn_dim=8       ,						max_position_embeddings=48       ,						)
       lowerCAmelCase_:							Union[str, Any]       =	FlaxBlenderbotSmallForConditionalGeneration(lowerCamelCase__ )
       lowerCAmelCase_:							str       =	np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]       ,						dtype=np.intaa )
       lowerCAmelCase_:							str       =	np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]       ,						dtype=np.intaa )
       lowerCAmelCase_:							Any       =	lm_model(input_ids=lowerCamelCase__       ,						decoder_input_ids=lowerCamelCase__ )
       lowerCAmelCase_:							List[str]       =	(*summary.shape, config.vocab_size)
       self.assertEqual(outputs["logits"].shape       ,						lowerCamelCase__ )




    def      _a    (    self ):
       lowerCAmelCase_:							Tuple       =	np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]       ,						dtype=np.intaa )
       lowerCAmelCase_:							Optional[Any]       =	shift_tokens_right(lowerCamelCase__       ,						1       ,						2 )
       lowerCAmelCase_:							List[Any]       =	np.equal(lowerCamelCase__       ,						1 ).astype(np.floataa ).sum()
       lowerCAmelCase_:							Any       =	np.equal(lowerCamelCase__       ,						1 ).astype(np.floataa ).sum()
       self.assertEqual(shifted.shape       ,						input_ids.shape )
       self.assertEqual(lowerCamelCase__       ,						n_pad_before - 1 )
       self.assertTrue(np.equal(shifted[:, 0]       ,						2 ).all() )





@require_flax
class 					_lowercase    (     UpperCAmelCase__ ,  unittest.TestCase ,  UpperCAmelCase__			):






    '''simple docstring'''
    SCREAMING_SNAKE_CASE:				List[Any]          =    True
    SCREAMING_SNAKE_CASE:				Optional[Any]          =    (
        (
            FlaxBlenderbotSmallModel,
            FlaxBlenderbotSmallForConditionalGeneration,
        )
        if is_flax_available()
        else ()
    )
    SCREAMING_SNAKE_CASE:				Optional[Any]          =    (FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()




    def      _a    (    self ):
       lowerCAmelCase_:							Optional[int]       =	FlaxBlenderbotSmallModelTester(self )




    def      _a    (    self ):
       lowerCAmelCase_  ,	lowerCAmelCase_:							Union[str, Any]       =	self.model_tester.prepare_config_and_inputs()
       for model_class in self.all_model_classes:
          self.model_tester.check_use_cache_forward(lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__ )




    def      _a    (    self ):
       lowerCAmelCase_  ,	lowerCAmelCase_:							Optional[Any]       =	self.model_tester.prepare_config_and_inputs()
       for model_class in self.all_model_classes:
          self.model_tester.check_use_cache_forward_with_attn_mask(lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__ )




    def      _a    (    self ):
       lowerCAmelCase_  ,	lowerCAmelCase_:							Optional[Any]       =	self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
          with self.subTest(model_class.__name__ ):
             lowerCAmelCase_:							Dict       =	self._prepare_for_class(lowerCamelCase__       ,						lowerCamelCase__ )
             lowerCAmelCase_:							Union[str, Any]       =	model_class(lowerCamelCase__ )

             @jax.jit
             def encode_jitted(lowerCamelCase__       ,						lowerCamelCase__=None       ,						**lowerCamelCase__ ):
                return model.encode(input_ids=lowerCamelCase__       ,						attention_mask=lowerCamelCase__ )

             with self.subTest("JIT Enabled" ):
                lowerCAmelCase_:							Optional[int]       =	encode_jitted(**lowerCamelCase__ ).to_tuple()

             with self.subTest("JIT Disabled" ):
                with jax.disable_jit():
                   lowerCAmelCase_:							int       =	encode_jitted(**lowerCamelCase__ ).to_tuple()

             self.assertEqual(len(lowerCamelCase__ )       ,						len(lowerCamelCase__ ) )
             for jitted_output, output in zip(lowerCamelCase__       ,						lowerCamelCase__ ):
                self.assertEqual(jitted_output.shape       ,						output.shape )




    def      _a    (    self ):
       lowerCAmelCase_  ,	lowerCAmelCase_:							Optional[Any]       =	self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
          with self.subTest(model_class.__name__ ):
             lowerCAmelCase_:							Dict       =	model_class(lowerCamelCase__ )
             lowerCAmelCase_:							Optional[int]       =	model.encode(inputs_dict["input_ids"]       ,						inputs_dict["attention_mask"] )

             lowerCAmelCase_:							Any       =	{
                 "decoder_input_ids": inputs_dict["decoder_input_ids"],
                 "decoder_attention_mask": inputs_dict["decoder_attention_mask"],
                 "encoder_outputs": encoder_outputs,
             }

             @jax.jit
             def decode_jitted(lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__ ):
                return model.decode(
                    decoder_input_ids=lowerCamelCase__       ,						decoder_attention_mask=lowerCamelCase__       ,						encoder_outputs=lowerCamelCase__       ,						)

             with self.subTest("JIT Enabled" ):
                lowerCAmelCase_:							Dict       =	decode_jitted(**lowerCamelCase__ ).to_tuple()

             with self.subTest("JIT Disabled" ):
                with jax.disable_jit():
                   lowerCAmelCase_:							Union[str, Any]       =	decode_jitted(**lowerCamelCase__ ).to_tuple()

             self.assertEqual(len(lowerCamelCase__ )       ,						len(lowerCamelCase__ ) )
             for jitted_output, output in zip(lowerCamelCase__       ,						lowerCamelCase__ ):
                self.assertEqual(jitted_output.shape       ,						output.shape )




    @slow
    def      _a    (    self ):
       for model_class_name in self.all_model_classes:
          lowerCAmelCase_:							Any       =	model_class_name.from_pretrained("facebook/blenderbot_small-90M" )
          # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
          lowerCAmelCase_:							str       =	np.ones((1, 1) ) * model.config.eos_token_id
          lowerCAmelCase_:							Optional[Any]       =	model(lowerCamelCase__ )
          self.assertIsNotNone(lowerCamelCase__ )