


import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

lowerCAmelCase__    :List[Any]       =				[
    # (stable-diffusion, HF Diffusers)
    ('''time_embed.0.weight''', '''time_embedding.linear_1.weight'''),
    ('''time_embed.0.bias''', '''time_embedding.linear_1.bias'''),
    ('''time_embed.2.weight''', '''time_embedding.linear_2.weight'''),
    ('''time_embed.2.bias''', '''time_embedding.linear_2.bias'''),
    ('''input_blocks.0.0.weight''', '''conv_in.weight'''),
    ('''input_blocks.0.0.bias''', '''conv_in.bias'''),
    ('''out.0.weight''', '''conv_norm_out.weight'''),
    ('''out.0.bias''', '''conv_norm_out.bias'''),
    ('''out.2.weight''', '''conv_out.weight'''),
    ('''out.2.bias''', '''conv_out.bias'''),
]

lowerCAmelCase__    :int       =				[
    # (stable-diffusion, HF Diffusers)
    ('''in_layers.0''', '''norm1'''),
    ('''in_layers.2''', '''conv1'''),
    ('''out_layers.0''', '''norm2'''),
    ('''out_layers.3''', '''conv2'''),
    ('''emb_layers.1''', '''time_emb_proj'''),
    ('''skip_connection''', '''conv_shortcut'''),
]

lowerCAmelCase__    :Optional[int]       =				[]
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
		# loop over downblocks/upblocks

		for j in range(2):
				# loop over resnets/attentions for downblocks
				lowerCAmelCase__    :Dict       =				f'''down_blocks.{i}.resnets.{j}.'''
				lowerCAmelCase__    :Any       =				f'''input_blocks.{3*i + j + 1}.0.'''
				unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

				if i < 3:
						# no attention layers in down_blocks.3
						lowerCAmelCase__    :str       =				f'''down_blocks.{i}.attentions.{j}.'''
						lowerCAmelCase__    :Optional[int]       =				f'''input_blocks.{3*i + j + 1}.1.'''
						unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

		for j in range(3):
				# loop over resnets/attentions for upblocks
				lowerCAmelCase__    :Optional[int]       =				f'''up_blocks.{i}.resnets.{j}.'''
				lowerCAmelCase__    :List[Any]       =				f'''output_blocks.{3*i + j}.0.'''
				unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

				if i > 0:
						# no attention layers in up_blocks.0
						lowerCAmelCase__    :int       =				f'''up_blocks.{i}.attentions.{j}.'''
						lowerCAmelCase__    :Dict       =				f'''output_blocks.{3*i + j}.1.'''
						unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

		if i < 3:
				# no downsample in down_blocks.3
				lowerCAmelCase__    :List[Any]       =				f'''down_blocks.{i}.downsamplers.0.conv.'''
				lowerCAmelCase__    :Tuple       =				f'''input_blocks.{3*(i+1)}.0.op.'''
				unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

				# no upsample in up_blocks.3
				lowerCAmelCase__    :Optional[Any]       =				f'''up_blocks.{i}.upsamplers.0.'''
				lowerCAmelCase__    :Optional[Any]       =				f'''output_blocks.{3*i + 2}.{1 if i == 0 else 2}.'''
				unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

lowerCAmelCase__    :Optional[int]       =				'''mid_block.attentions.0.'''
lowerCAmelCase__    :Union[str, Any]       =				'''middle_block.1.'''
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
		lowerCAmelCase__    :Union[str, Any]       =				f'''mid_block.resnets.{j}.'''
		lowerCAmelCase__    :List[Any]       =				f'''middle_block.{2*j}.'''
		unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))

def 			lowerCAmelCase__					(     a__:						Any			)    ->						Any:


							'''simple docstring'''





							_UpperCAmelCase											=     {k: k for k in unet_state_dict.keys()}
							for sd_name, hf_name in unet_conversion_map:
														_UpperCAmelCase											=     sd_name
							for k, v in mapping.items():
														if "resnets" in k:
																					for sd_part, hf_part in unet_conversion_map_resnet:
																												_UpperCAmelCase											=     v.replace(a__     ,    a__			)
																					_UpperCAmelCase											=     v
							for k, v in mapping.items():
														for sd_part, hf_part in unet_conversion_map_layer:
																					_UpperCAmelCase											=     v.replace(a__     ,    a__			)
														_UpperCAmelCase											=     v
							_UpperCAmelCase											=     {v: unet_state_dict[k] for k, v in mapping.items()}
							return new_state_dict


# ================#
# VAE Conversion #
# ================#

lowerCAmelCase__    :List[str]       =				[
    # (stable-diffusion, HF Diffusers)
    ('''nin_shortcut''', '''conv_shortcut'''),
    ('''norm_out''', '''conv_norm_out'''),
    ('''mid.attn_1.''', '''mid_block.attentions.0.'''),
]

for i in range(4):
		# down_blocks have two resnets
		for j in range(2):
				lowerCAmelCase__    :Tuple       =				f'''encoder.down_blocks.{i}.resnets.{j}.'''
				lowerCAmelCase__    :str       =				f'''encoder.down.{i}.block.{j}.'''
				vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

		if i < 3:
				lowerCAmelCase__    :Any       =				f'''down_blocks.{i}.downsamplers.0.'''
				lowerCAmelCase__    :Union[str, Any]       =				f'''down.{i}.downsample.'''
				vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

				lowerCAmelCase__    :List[Any]       =				f'''up_blocks.{i}.upsamplers.0.'''
				lowerCAmelCase__    :List[str]       =				f'''up.{3-i}.upsample.'''
				vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

		# up_blocks have three resnets
		# also, up blocks in hf are numbered in reverse from sd
		for j in range(3):
				lowerCAmelCase__    :str       =				f'''decoder.up_blocks.{i}.resnets.{j}.'''
				lowerCAmelCase__    :Optional[int]       =				f'''decoder.up.{3-i}.block.{j}.'''
				vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
		lowerCAmelCase__    :Optional[Any]       =				f'''mid_block.resnets.{i}.'''
		lowerCAmelCase__    :int       =				f'''mid.block_{i+1}.'''
		vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


lowerCAmelCase__    :Tuple       =				[
    # (stable-diffusion, HF Diffusers)
    ('''norm.''', '''group_norm.'''),
    ('''q.''', '''query.'''),
    ('''k.''', '''key.'''),
    ('''v.''', '''value.'''),
    ('''proj_out.''', '''proj_attn.'''),
]

def 			lowerCAmelCase__					(     a__:						Tuple			)    ->						Tuple:


							'''simple docstring'''





							return w.reshape(*w.shape     ,    1     ,    1			)

def 			lowerCAmelCase__					(     a__:						int			)    ->						int:


							'''simple docstring'''





							_UpperCAmelCase											=     {k: k for k in vae_state_dict.keys()}
							for k, v in mapping.items():
														for sd_part, hf_part in vae_conversion_map:
																					_UpperCAmelCase											=     v.replace(a__     ,    a__			)
														_UpperCAmelCase											=     v
							for k, v in mapping.items():
														if "attentions" in k:
																					for sd_part, hf_part in vae_conversion_map_attn:
																												_UpperCAmelCase											=     v.replace(a__     ,    a__			)
																					_UpperCAmelCase											=     v
							_UpperCAmelCase											=     {v: vae_state_dict[k] for k, v in mapping.items()}
							_UpperCAmelCase											=     ['q', 'k', 'v', 'proj_out']
							for k, v in new_state_dict.items():
														for weight_name in weights_to_convert:
																					if F'''mid.attn_1.{weight_name}.weight''' in k:
																												print(F'''Reshaping {k} for SD format'''			)
																												_UpperCAmelCase											=     reshape_weight_for_sd(a__			)
							return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


lowerCAmelCase__    :Union[str, Any]       =				[
    # (stable-diffusion, HF Diffusers)
    ('''resblocks.''', '''text_model.encoder.layers.'''),
    ('''ln_1''', '''layer_norm1'''),
    ('''ln_2''', '''layer_norm2'''),
    ('''.c_fc.''', '''.fc1.'''),
    ('''.c_proj.''', '''.fc2.'''),
    ('''.attn''', '''.self_attn'''),
    ('''ln_final.''', '''transformer.text_model.final_layer_norm.'''),
    ('''token_embedding.weight''', '''transformer.text_model.embeddings.token_embedding.weight'''),
    ('''positional_embedding''', '''transformer.text_model.embeddings.position_embedding.weight'''),
]
lowerCAmelCase__    :Optional[Any]       =				{re.escape(x[1]): x[0] for x in textenc_conversion_lst}
lowerCAmelCase__    :List[str]       =				re.compile('''|'''.join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
lowerCAmelCase__    :str       =				{'''q''': 0, '''k''': 1, '''v''': 2}

def 			lowerCAmelCase__					(     a__:						List[Any]			)    ->						List[str]:


							'''simple docstring'''





							_UpperCAmelCase											=     {}
							_UpperCAmelCase											=     {}
							_UpperCAmelCase											=     {}
							for k, v in text_enc_dict.items():
														if (
														    k.endswith('.self_attn.q_proj.weight'			)
														    or k.endswith('.self_attn.k_proj.weight'			)
														    or k.endswith('.self_attn.v_proj.weight'			)
														):
																					_UpperCAmelCase											=     k[: -len('.q_proj.weight'			)]
																					_UpperCAmelCase											=     k[-len('q_proj.weight'			)]
																					if k_pre not in capture_qkv_weight:
																												_UpperCAmelCase											=     [None, None, None]
																					_UpperCAmelCase											=     v
																					continue

														if (
														    k.endswith('.self_attn.q_proj.bias'			)
														    or k.endswith('.self_attn.k_proj.bias'			)
														    or k.endswith('.self_attn.v_proj.bias'			)
														):
																					_UpperCAmelCase											=     k[: -len('.q_proj.bias'			)]
																					_UpperCAmelCase											=     k[-len('q_proj.bias'			)]
																					if k_pre not in capture_qkv_bias:
																												_UpperCAmelCase											=     [None, None, None]
																					_UpperCAmelCase											=     v
																					continue

														_UpperCAmelCase											=     textenc_pattern.sub(lambda a__			: protected[re.escape(m.group(0			)			)]     ,    a__			)
														_UpperCAmelCase											=     v

							for k_pre, tensors in capture_qkv_weight.items():
														if None in tensors:
																					raise Exception('CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'			)
														_UpperCAmelCase											=     textenc_pattern.sub(lambda a__			: protected[re.escape(m.group(0			)			)]     ,    a__			)
														_UpperCAmelCase											=     torch.cat(a__			)

							for k_pre, tensors in capture_qkv_bias.items():
														if None in tensors:
																					raise Exception('CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'			)
														_UpperCAmelCase											=     textenc_pattern.sub(lambda a__			: protected[re.escape(m.group(0			)			)]     ,    a__			)
														_UpperCAmelCase											=     torch.cat(a__			)

							return new_state_dict

def 			lowerCAmelCase__					(     a__:						List[str]			)    ->						str:


							'''simple docstring'''





							return text_enc_dict


if __name__ == "__main__":
		lowerCAmelCase__    :Tuple       =				argparse.ArgumentParser()

		parser.add_argument('''--model_path''', default=None, type=str, required=True, help='''Path to the model to convert.''')
		parser.add_argument('''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the output model.''')
		parser.add_argument('''--half''', action='''store_true''', help='''Save weights in half precision.''')
		parser.add_argument(
		    '''--use_safetensors''', action='''store_true''', help='''Save weights use safetensors, default is ckpt.'''
		)

		lowerCAmelCase__    :Tuple       =				parser.parse_args()

		assert args.model_path is not None, "Must provide a model path!"

		assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

		# Path for safetensors
		lowerCAmelCase__    :Optional[Any]       =				osp.join(args.model_path, '''unet''', '''diffusion_pytorch_model.safetensors''')
		lowerCAmelCase__    :int       =				osp.join(args.model_path, '''vae''', '''diffusion_pytorch_model.safetensors''')
		lowerCAmelCase__    :List[Any]       =				osp.join(args.model_path, '''text_encoder''', '''model.safetensors''')

		# Load models from safetensors if it exists, if it doesn't pytorch
		if osp.exists(unet_path):
				lowerCAmelCase__    :List[str]       =				load_file(unet_path, device='''cpu''')
		else:
				lowerCAmelCase__    :Union[str, Any]       =				osp.join(args.model_path, '''unet''', '''diffusion_pytorch_model.bin''')
				lowerCAmelCase__    :Optional[int]       =				torch.load(unet_path, map_location='''cpu''')

		if osp.exists(vae_path):
				lowerCAmelCase__    :Dict       =				load_file(vae_path, device='''cpu''')
		else:
				lowerCAmelCase__    :List[str]       =				osp.join(args.model_path, '''vae''', '''diffusion_pytorch_model.bin''')
				lowerCAmelCase__    :List[str]       =				torch.load(vae_path, map_location='''cpu''')

		if osp.exists(text_enc_path):
				lowerCAmelCase__    :List[str]       =				load_file(text_enc_path, device='''cpu''')
		else:
				lowerCAmelCase__    :List[Any]       =				osp.join(args.model_path, '''text_encoder''', '''pytorch_model.bin''')
				lowerCAmelCase__    :Tuple       =				torch.load(text_enc_path, map_location='''cpu''')

		# Convert the UNet model
		lowerCAmelCase__    :Optional[Any]       =				convert_unet_state_dict(unet_state_dict)
		lowerCAmelCase__    :Tuple       =				{'''model.diffusion_model.''' + k: v for k, v in unet_state_dict.items()}

		# Convert the VAE model
		lowerCAmelCase__    :str       =				convert_vae_state_dict(vae_state_dict)
		lowerCAmelCase__    :Any       =				{'''first_stage_model.''' + k: v for k, v in vae_state_dict.items()}

		# Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
		lowerCAmelCase__    :List[str]       =				'''text_model.encoder.layers.22.layer_norm2.bias''' in text_enc_dict

		if is_vaa_model:
				# Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
				lowerCAmelCase__    :Tuple       =				{'''transformer.''' + k: v for k, v in text_enc_dict.items()}
				lowerCAmelCase__    :Tuple       =				convert_text_enc_state_dict_vaa(text_enc_dict)
				lowerCAmelCase__    :Union[str, Any]       =				{'''cond_stage_model.model.''' + k: v for k, v in text_enc_dict.items()}
		else:
				lowerCAmelCase__    :Union[str, Any]       =				convert_text_enc_state_dict(text_enc_dict)
				lowerCAmelCase__    :List[Any]       =				{'''cond_stage_model.transformer.''' + k: v for k, v in text_enc_dict.items()}

		# Put together new checkpoint
		lowerCAmelCase__    :Optional[int]       =				{**unet_state_dict, **vae_state_dict, **text_enc_dict}
		if args.half:
				lowerCAmelCase__    :Optional[int]       =				{k: v.half() for k, v in state_dict.items()}

		if args.use_safetensors:
				save_file(state_dict, args.checkpoint_path)
		else:
				lowerCAmelCase__    :int       =				{'''state_dict''': state_dict}
				torch.save(state_dict, args.checkpoint_path)




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     ):
		_a      :    int           =      MODEL_FOR_MASKED_LM_MAPPING
		_a      :    Tuple           =      TF_MODEL_FOR_MASKED_LM_MAPPING
		def       UpperCAmelCase__							(     self		)					->			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       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





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

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

									_UpperCAmelCase											=     unmasker('My name is <mask>'     ,       targets=[' Patrick', ' Clara', ' Teven']     ,       top_k=3		)
									self.assertEqual(
									    nested_simplify(_SCREAMING_SNAKE_CASE     ,       decimals=6		)     ,       [
									        {'sequence': 'My name is Clara', 'score': 2e-0_5, 'token': 13606, 'token_str': ' Clara'},
									        {'sequence': 'My name is Patrick', 'score': 2e-0_5, 'token': 3499, 'token_str': ' Patrick'},
									        {'sequence': 'My name is Te', 'score': 1.9e-0_5, 'token': 2941, 'token_str': ' Te'},
									    ]     ,       )
		@require_torch
		def       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





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

									_UpperCAmelCase											=     unmasker('My name is <mask>'		)
									self.assertEqual(
									    nested_simplify(_SCREAMING_SNAKE_CASE     ,       decimals=6		)     ,       [
									        {'sequence': 'My name is Maul', 'score': 2.2e-0_5, 'token': 35676, 'token_str': ' Maul'},
									        {'sequence': 'My name isELS', 'score': 2.2e-0_5, 'token': 16416, 'token_str': 'ELS'},
									    ]     ,       )

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

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

									_UpperCAmelCase											=     unmasker('My name is <mask> <mask>'     ,       top_k=2		)

									self.assertEqual(
									    nested_simplify(_SCREAMING_SNAKE_CASE     ,       decimals=6		)     ,       [
									        [
									            {
									                'score': 2.2e-0_5,
									                'token': 35676,
									                'token_str': ' Maul',
									                'sequence': '<s>My name is Maul<mask></s>',
									            },
									            {'score': 2.2e-0_5, 'token': 16416, 'token_str': 'ELS', 'sequence': '<s>My name isELS<mask></s>'},
									        ],
									        [
									            {
									                'score': 2.2e-0_5,
									                'token': 35676,
									                'token_str': ' Maul',
									                'sequence': '<s>My name is<mask> Maul</s>',
									            },
									            {'score': 2.2e-0_5, 'token': 16416, 'token_str': 'ELS', 'sequence': '<s>My name is<mask>ELS</s>'},
									        ],
									    ]     ,       )
		@require_torch_gpu
		def       UpperCAmelCase__							(     self		)					->			List[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     pipeline('fill-mask'     ,       model='hf-internal-testing/tiny-random-distilbert'     ,       device=0     ,       framework='pt'		)

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

									_UpperCAmelCase											=     pipe('Paris is the [MASK] of France.'		)
									# We actually don't care about the result, we just want to make sure
									# it works, meaning the float16 tensor got casted back to float32
									# for postprocessing.
									self.assertIsInstance(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
		@slow
		@require_torch
		def       UpperCAmelCase__							(     self		)					->			Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase											=     pipeline(task='fill-mask'     ,       model='distilroberta-base'     ,       top_k=2     ,       framework='pt'		)
									self.run_large_test(_SCREAMING_SNAKE_CASE		)
		@slow
		@require_tf
		def       UpperCAmelCase__							(     self		)					->			Dict:




									"""simple docstring"""





									_UpperCAmelCase											=     pipeline(task='fill-mask'     ,       model='distilroberta-base'     ,       top_k=2     ,       framework='tf'		)
									self.run_large_test(_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE		)					->			Any:




									"""simple docstring"""





									_UpperCAmelCase											=     unmasker('My name is <mask>'		)
									self.assertEqual(
									    nested_simplify(_SCREAMING_SNAKE_CASE		)     ,       [
									        {'sequence': 'My name is John', 'score': 0.008, 'token': 610, 'token_str': ' John'},
									        {'sequence': 'My name is Chris', 'score': 0.007, 'token': 1573, 'token_str': ' Chris'},
									    ]     ,       )
									_UpperCAmelCase											=     unmasker('The largest city in France is <mask>'		)
									self.assertEqual(
									    nested_simplify(_SCREAMING_SNAKE_CASE		)     ,       [
									        {
									            'sequence': 'The largest city in France is Paris',
									            'score': 0.251,
									            'token': 2201,
									            'token_str': ' Paris',
									        },
									        {
									            'sequence': 'The largest city in France is Lyon',
									            'score': 0.214,
									            'token': 12790,
									            'token_str': ' Lyon',
									        },
									    ]     ,       )

									_UpperCAmelCase											=     unmasker('My name is <mask>'     ,       targets=[' Patrick', ' Clara', ' Teven']     ,       top_k=3		)
									self.assertEqual(
									    nested_simplify(_SCREAMING_SNAKE_CASE		)     ,       [
									        {'sequence': 'My name is Patrick', 'score': 0.005, 'token': 3499, 'token_str': ' Patrick'},
									        {'sequence': 'My name is Clara', 'score': 0.000, 'token': 13606, 'token_str': ' Clara'},
									        {'sequence': 'My name is Te', 'score': 0.000, 'token': 2941, 'token_str': ' Te'},
									    ]     ,       )
		@require_torch
		def       UpperCAmelCase__							(     self		)					->			Dict:




									"""simple docstring"""





									_UpperCAmelCase											=     pipeline(task='fill-mask'     ,       model='sshleifer/tiny-distilroberta-base'     ,       framework='pt'		)
									_UpperCAmelCase											=     None
									_UpperCAmelCase											=     None
									self.run_pipeline_test(_SCREAMING_SNAKE_CASE     ,       []		)
		@require_tf
		def       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     pipeline(task='fill-mask'     ,       model='sshleifer/tiny-distilroberta-base'     ,       framework='tf'		)
									_UpperCAmelCase											=     None
									_UpperCAmelCase											=     None
									self.run_pipeline_test(_SCREAMING_SNAKE_CASE     ,       []		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			List[Any]:




									"""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											=     FillMaskPipeline(model=_SCREAMING_SNAKE_CASE     ,       tokenizer=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    f'''This is another {tokenizer.mask_token} test''',
									]
									return fill_masker, examples
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			List[str]:




									"""simple docstring"""





									_UpperCAmelCase											=     fill_masker.tokenizer
									_UpperCAmelCase											=     fill_masker.model

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

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

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

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

									self.run_test_top_k(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
									self.run_test_targets(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
									self.run_test_top_k_targets(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
									self.fill_mask_with_duplicate_targets_and_top_k(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
									self.fill_mask_with_multiple_masks(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Any:




									"""simple docstring"""





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

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

									# Score equivalence
									_UpperCAmelCase											=     fill_masker(f'''This is a {tokenizer.mask_token}'''     ,       targets=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [top_mask['token_str'] for top_mask in outputs]
									_UpperCAmelCase											=     [top_mask['score'] for top_mask in outputs]

									# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
									if set(_SCREAMING_SNAKE_CASE		) == set(_SCREAMING_SNAKE_CASE		):
																_UpperCAmelCase											=     fill_masker(f'''This is a {tokenizer.mask_token}'''     ,       targets=_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase											=     [top_mask['score'] for top_mask in unmasked_targets]
																self.assertEqual(nested_simplify(_SCREAMING_SNAKE_CASE		)     ,       nested_simplify(_SCREAMING_SNAKE_CASE		)		)

									# Raises with invalid
									with self.assertRaises(_SCREAMING_SNAKE_CASE		):
																_UpperCAmelCase											=     fill_masker(f'''This is a {tokenizer.mask_token}'''     ,       targets=[]		)
									# For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
									if "" not in tokenizer.get_vocab():
																with self.assertRaises(_SCREAMING_SNAKE_CASE		):
																							_UpperCAmelCase											=     fill_masker(f'''This is a {tokenizer.mask_token}'''     ,       targets=['']		)
																with self.assertRaises(_SCREAMING_SNAKE_CASE		):
																							_UpperCAmelCase											=     fill_masker(f'''This is a {tokenizer.mask_token}'''     ,       targets=''		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Union[str, Any]:




									"""simple docstring"""





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

									_UpperCAmelCase											=     FillMaskPipeline(model=_SCREAMING_SNAKE_CASE     ,       tokenizer=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     fill_masker(f'''This is a {tokenizer.mask_token}'''     ,       top_k=2		)
									self.assertEqual(
									    _SCREAMING_SNAKE_CASE     ,       [
									        {'sequence': ANY(_SCREAMING_SNAKE_CASE		), 'score': ANY(_SCREAMING_SNAKE_CASE		), 'token': ANY(_SCREAMING_SNAKE_CASE		), 'token_str': ANY(_SCREAMING_SNAKE_CASE		)},
									        {'sequence': ANY(_SCREAMING_SNAKE_CASE		), 'score': ANY(_SCREAMING_SNAKE_CASE		), 'token': ANY(_SCREAMING_SNAKE_CASE		), 'token_str': ANY(_SCREAMING_SNAKE_CASE		)},
									    ]     ,       )
									self.assertEqual(nested_simplify(_SCREAMING_SNAKE_CASE		)     ,       nested_simplify(_SCREAMING_SNAKE_CASE		)		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     tokenizer.get_vocab()
									_UpperCAmelCase											=     FillMaskPipeline(model=_SCREAMING_SNAKE_CASE     ,       tokenizer=_SCREAMING_SNAKE_CASE		)

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

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




									"""simple docstring"""





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

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






		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase											=     FillMaskPipeline(model=_SCREAMING_SNAKE_CASE     ,       tokenizer=_SCREAMING_SNAKE_CASE		)

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