

from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


SCREAMING_SNAKE_CASE_          =      logging.get_logger(__name__)  # pylint: disable=invalid-name

SCREAMING_SNAKE_CASE_          =      '\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")\n        >>> pipe_prior.to("cuda")\n        >>> prompt = "red cat, 4k photo"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> zero_image_emb = out.negative_image_embeds\n        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")\n        >>> pipe.to("cuda")\n        >>> image = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=50,\n        ... ).images\n        >>> image[0].save("cat.png")\n        ```\n'

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Tuple  ,      lowerCAmelCase:  Union[str, Any]  ,      lowerCAmelCase:  Optional[Any]=8 )   ->       List[Any]:
				_UpperCAmelCase       :		Optional[Any] 		=    height // scale_factor**2
				if height % scale_factor**2 != 0:
								new_height += 1
				_UpperCAmelCase       :		Optional[int] 		=    width // scale_factor**2
				if width % scale_factor**2 != 0:
								new_width += 1
				return new_height * scale_factor, new_width * scale_factor





class 				a  (	UpperCAmelCase				):


						def __init__(       self						,					A_						,					A_						,					A_						,					):







										'''simple docstring'''


										super().__init__()

										self.register_modules(
										    unet=A_						,					scheduler=A_						,					movq=A_						,					)
										_UpperCAmelCase       :		Optional[Any] 		=    2 ** (len(self.movq.config.block_out_channels	) - 1)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										if latents is None:
														_UpperCAmelCase       :		int 		=    randn_tensor(A_						,					generator=A_						,					device=A_						,					dtype=A_	)
										else:
														if latents.shape != shape:
																		raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {shape}'	)
														_UpperCAmelCase       :		Any 		=    latents.to(A_	)

										_UpperCAmelCase       :		Dict 		=    latents * scheduler.init_noise_sigma
										return latents


						def        _UpperCAmelCase				(       self						,					A_=0	):







										'''simple docstring'''


										if is_accelerate_available():
														from accelerate import cpu_offload
										else:
														raise ImportError("Please install accelerate via `pip install accelerate`"	)

										_UpperCAmelCase       :		Any 		=    torch.device(f'cuda:{gpu_id}'	)

										_UpperCAmelCase       :		str 		=    [
										    self.unet,
										    self.movq,
										]
										for cpu_offloaded_model in models:
														if cpu_offloaded_model is not None:
																		cpu_offload(A_						,					A_	)


						def        _UpperCAmelCase				(       self						,					A_=0	):







										'''simple docstring'''


										if is_accelerate_available() and is_accelerate_version(">="						,					"0.17.0.dev0"	):
														from accelerate import cpu_offload_with_hook
										else:
														raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher."	)

										_UpperCAmelCase       :		Optional[Any] 		=    torch.device(f'cuda:{gpu_id}'	)

										if self.device.type != "cpu":
														self.to("cpu"						,					silence_dtype_warnings=A_	)
														torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

										_UpperCAmelCase       :		Dict 		=    None
										for cpu_offloaded_model in [self.unet, self.movq]:
														_UpperCAmelCase       :		Union[str, Any] 		=    cpu_offload_with_hook(A_						,					A_						,					prev_module_hook=A_	)

										# We'll offload the last model manually.
										_UpperCAmelCase       :		Any 		=    hook


						@property
						# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										if not hasattr(self.unet						,					"_hf_hook"	):
														return self.device
										for module in self.unet.modules():
														if (
														    hasattr(A_						,					"_hf_hook"	)
														    and hasattr(module._hf_hook						,					"execution_device"	)
														    and module._hf_hook.execution_device is not None
														):
																		return torch.device(module._hf_hook.execution_device	)
										return self.device





						@torch.no_grad()
						@replace_example_docstring(A_	)
						def __call__(       self						,					A_						,					A_						,					A_ = 512						,					A_ = 512						,					A_ = 100						,					A_ = 4.0						,					A_ = 1						,					A_ = None						,					A_ = None						,					A_ = "pil"						,					A_ = True						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    self._execution_device

										_UpperCAmelCase       :		Optional[int] 		=    guidance_scale > 1.0

										if isinstance(A_						,					A_	):
														_UpperCAmelCase       :		str 		=    torch.cat(A_						,					dim=0	)
										_UpperCAmelCase       :		Dict 		=    image_embeds.shape[0] * num_images_per_prompt
										if isinstance(A_						,					A_	):
														_UpperCAmelCase       :		List[str] 		=    torch.cat(A_						,					dim=0	)

										if do_classifier_free_guidance:
														_UpperCAmelCase       :		Tuple 		=    image_embeds.repeat_interleave(A_						,					dim=0	)
														_UpperCAmelCase       :		Optional[int] 		=    negative_image_embeds.repeat_interleave(A_						,					dim=0	)

										_UpperCAmelCase       :		int 		=    torch.cat([negative_image_embeds, image_embeds]						,					dim=0	).to(dtype=self.unet.dtype						,					device=A_	)

										self.scheduler.set_timesteps(A_						,					device=A_	)
										_UpperCAmelCase       :		Union[str, Any] 		=    self.scheduler.timesteps

										_UpperCAmelCase       :		Any 		=    self.unet.config.in_channels

										_UpperCAmelCase       :		Optional[Any] 		=    downscale_height_and_width(A_						,					A_						,					self.movq_scale_factor	)

										# create initial latent
										_UpperCAmelCase       :		Tuple 		=    self.prepare_latents(
										    (batch_size, num_channels_latents, height, width)						,					image_embeds.dtype						,					A_						,					A_						,					A_						,					self.scheduler						,					)

										for i, t in enumerate(self.progress_bar(A_	)	):
														# expand the latents if we are doing classifier free guidance
														_UpperCAmelCase       :		Union[str, Any] 		=    torch.cat([latents] * 2	) if do_classifier_free_guidance else latents

														_UpperCAmelCase       :		List[Any] 		=    {"image_embeds": image_embeds}
														_UpperCAmelCase       :		Dict 		=    self.unet(
														    sample=A_						,					timestep=A_						,					encoder_hidden_states=A_						,					added_cond_kwargs=A_						,					return_dict=A_						,					)[0]

														if do_classifier_free_guidance:
																		_UpperCAmelCase       :		List[Any] 		=    noise_pred.split(latents.shape[1]						,					dim=1	)
																		_UpperCAmelCase       :		List[str] 		=    noise_pred.chunk(2	)
																		_UpperCAmelCase       :		str 		=    variance_pred.chunk(2	)
																		_UpperCAmelCase       :		int 		=    noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																		_UpperCAmelCase       :		Dict 		=    torch.cat([noise_pred, variance_pred_text]						,					dim=1	)

														if not (
														    hasattr(self.scheduler.config						,					"variance_type"	)
														    and self.scheduler.config.variance_type in ["learned", "learned_range"]
														):
																		_UpperCAmelCase       :		List[str] 		=    noise_pred.split(latents.shape[1]						,					dim=1	)

														# compute the previous noisy sample x_t -> x_t-1
														_UpperCAmelCase       :		Tuple 		=    self.scheduler.step(
														    A_						,					A_						,					A_						,					generator=A_						,					)[0]
										# post-processing
										_UpperCAmelCase       :		str 		=    self.movq.decode(A_						,					force_not_quantize=A_	)["sample"]

										if output_type not in ["pt", "np", "pil"]:
														raise ValueError(f'Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}'	)

										if output_type in ["np", "pil"]:
														_UpperCAmelCase       :		Any 		=    image * 0.5 + 0.5
														_UpperCAmelCase       :		str 		=    image.clamp(0						,					1	)
														_UpperCAmelCase       :		List[Any] 		=    image.cpu().permute(0						,					2						,					3						,					1	).float().numpy()

										if output_type == "pil":
														_UpperCAmelCase       :		Union[str, Any] 		=    self.numpy_to_pil(A_	)

										if not return_dict:
														return (image,)

										return ImagePipelineOutput(images=A_	)





import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY

@is_pipeline_test
class 				a  (	unittest.TestCase				):
						_lowercase    						=			MODEL_FOR_MASKED_LM_MAPPING
						_lowercase    						=			TF_MODEL_FOR_MASKED_LM_MAPPING


						def        _UpperCAmelCase				(       self	):







										'''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        _UpperCAmelCase				(       self	):







										'''simple docstring'''


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

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

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


						@require_torch
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    pipeline(task="fill-mask"						,					model="sshleifer/tiny-distilroberta-base"						,					top_k=2						,					framework="pt"	)

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

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

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

										_UpperCAmelCase       :		Optional[Any] 		=    unmasker("My name is <mask> <mask>"						,					top_k=2	)

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


						@require_torch_gpu
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    pipeline("fill-mask"						,					model="hf-internal-testing/tiny-random-distilbert"						,					device=0						,					framework="pt"	)

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

										_UpperCAmelCase       :		int 		=    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(A_						,					A_	)


						@slow
						@require_torch
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[Any] 		=    pipeline(task="fill-mask"						,					model="distilroberta-base"						,					top_k=2						,					framework="pt"	)
										self.run_large_test(A_	)


						@slow
						@require_tf
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    pipeline(task="fill-mask"						,					model="distilroberta-base"						,					top_k=2						,					framework="tf"	)
										self.run_large_test(A_	)


						def        _UpperCAmelCase				(       self						,					A_	):







										'''simple docstring'''


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

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


						@require_torch
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    pipeline(task="fill-mask"						,					model="sshleifer/tiny-distilroberta-base"						,					framework="pt"	)
										_UpperCAmelCase       :		Optional[int] 		=    None
										_UpperCAmelCase       :		List[Any] 		=    None
										self.run_pipeline_test(A_						,					[]	)


						@require_tf
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    pipeline(task="fill-mask"						,					model="sshleifer/tiny-distilroberta-base"						,					framework="tf"	)
										_UpperCAmelCase       :		Union[str, Any] 		=    None
										_UpperCAmelCase       :		str 		=    None
										self.run_pipeline_test(A_						,					[]	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										if tokenizer is None or tokenizer.mask_token_id is None:
														self.skipTest("The provided tokenizer has no mask token, (probably reformer or wav2vec2)"	)

										_UpperCAmelCase       :		Tuple 		=    FillMaskPipeline(model=A_						,					tokenizer=A_	)
										_UpperCAmelCase       :		Tuple 		=    [
										    f'This is another {tokenizer.mask_token} test',
										]
										return fill_masker, examples


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


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

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

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

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

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

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


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


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

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

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

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


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


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

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


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    tokenizer.get_vocab()
										_UpperCAmelCase       :		Union[str, Any] 		=    FillMaskPipeline(model=A_						,					tokenizer=A_	)

										# top_k=2, ntargets=3
										_UpperCAmelCase       :		Tuple 		=    sorted(vocab.keys()	)[:3]
										_UpperCAmelCase       :		Union[str, Any] 		=    fill_masker(f'This is a {tokenizer.mask_token}'						,					top_k=2						,					targets=A_	)

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


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


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

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





						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    FillMaskPipeline(model=A_						,					tokenizer=A_	)

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


