


import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import (
    RobertaTokenizer,
    TrOCRConfig,
    TrOCRForCausalLM,
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    ViTConfig,
    ViTImageProcessor,
    ViTModel,
)
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase__    :Dict       =				logging.get_logger(__name__)

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


							'''simple docstring'''





							_UpperCAmelCase											=     []
							for i in range(encoder_config.num_hidden_layers			):
														# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
														rename_keys.append(
														    (F'''encoder.deit.blocks.{i}.norm1.weight''', F'''encoder.encoder.layer.{i}.layernorm_before.weight''')			)
														rename_keys.append((F'''encoder.deit.blocks.{i}.norm1.bias''', F'''encoder.encoder.layer.{i}.layernorm_before.bias''')			)
														rename_keys.append(
														    (F'''encoder.deit.blocks.{i}.attn.proj.weight''', F'''encoder.encoder.layer.{i}.attention.output.dense.weight''')			)
														rename_keys.append(
														    (F'''encoder.deit.blocks.{i}.attn.proj.bias''', F'''encoder.encoder.layer.{i}.attention.output.dense.bias''')			)
														rename_keys.append(
														    (F'''encoder.deit.blocks.{i}.norm2.weight''', F'''encoder.encoder.layer.{i}.layernorm_after.weight''')			)
														rename_keys.append((F'''encoder.deit.blocks.{i}.norm2.bias''', F'''encoder.encoder.layer.{i}.layernorm_after.bias''')			)
														rename_keys.append(
														    (F'''encoder.deit.blocks.{i}.mlp.fc1.weight''', F'''encoder.encoder.layer.{i}.intermediate.dense.weight''')			)
														rename_keys.append(
														    (F'''encoder.deit.blocks.{i}.mlp.fc1.bias''', F'''encoder.encoder.layer.{i}.intermediate.dense.bias''')			)
														rename_keys.append(
														    (F'''encoder.deit.blocks.{i}.mlp.fc2.weight''', F'''encoder.encoder.layer.{i}.output.dense.weight''')			)
														rename_keys.append((F'''encoder.deit.blocks.{i}.mlp.fc2.bias''', F'''encoder.encoder.layer.{i}.output.dense.bias''')			)

							# cls token, position embeddings and patch embeddings of encoder
							rename_keys.extend(
							    [
							        ('encoder.deit.cls_token', 'encoder.embeddings.cls_token'),
							        ('encoder.deit.pos_embed', 'encoder.embeddings.position_embeddings'),
							        ('encoder.deit.patch_embed.proj.weight', 'encoder.embeddings.patch_embeddings.projection.weight'),
							        ('encoder.deit.patch_embed.proj.bias', 'encoder.embeddings.patch_embeddings.projection.bias'),
							        ('encoder.deit.norm.weight', 'encoder.layernorm.weight'),
							        ('encoder.deit.norm.bias', 'encoder.layernorm.bias'),
							    ]			)

							return rename_keys

def 			lowerCAmelCase__					(     a__:						Tuple     ,    a__:						List[Any]			)    ->						Optional[Any]:


							'''simple docstring'''





							for i in range(encoder_config.num_hidden_layers			):
														# queries, keys and values (only weights, no biases)
														_UpperCAmelCase											=     state_dict.pop(F'''encoder.deit.blocks.{i}.attn.qkv.weight'''			)

														_UpperCAmelCase											=     in_proj_weight[
														    : encoder_config.hidden_size, :
														]
														_UpperCAmelCase											=     in_proj_weight[
														    encoder_config.hidden_size : encoder_config.hidden_size * 2, :
														]
														_UpperCAmelCase											=     in_proj_weight[
														    -encoder_config.hidden_size :, :
														]

def 			lowerCAmelCase__					(     a__:						Optional[int]     ,    a__:						List[str]     ,    a__:						List[Any]			)    ->						Tuple:


							'''simple docstring'''





							_UpperCAmelCase											=     dct.pop(a__			)
							_UpperCAmelCase											=     val

def 			lowerCAmelCase__					(     a__:						Tuple			)    ->						Union[str, Any]:


							'''simple docstring'''





							if "handwritten" in checkpoint_url:
														_UpperCAmelCase											=     'https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg'  # industry
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-12.jpg" # have
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-10.jpg" # let
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"  #
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122.jpg"
							elif "printed" in checkpoint_url or "stage1" in checkpoint_url:
														_UpperCAmelCase											=     'https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg'
							_UpperCAmelCase											=     Image.open(requests.get(a__     ,    stream=a__			).raw			).convert('RGB'			)
							return im

@torch.no_grad()
def 			lowerCAmelCase__					(     a__:						Optional[Any]     ,    a__:						Dict			)    ->						Optional[int]:


							'''simple docstring'''





							_UpperCAmelCase											=     ViTConfig(image_size=3_8_4     ,    qkv_bias=a__			)
							_UpperCAmelCase											=     TrOCRConfig()

							# size of the architecture
							if "base" in checkpoint_url:
														_UpperCAmelCase											=     7_6_8
							elif "large" in checkpoint_url:
														# use ViT-large encoder
														_UpperCAmelCase											=     1_0_2_4
														_UpperCAmelCase											=     4_0_9_6
														_UpperCAmelCase											=     2_4
														_UpperCAmelCase											=     1_6
														_UpperCAmelCase											=     1_0_2_4
							else:
														raise ValueError('Should either find \'base\' or \'large\' in checkpoint URL'			)

							# the large-printed + stage1 checkpoints uses sinusoidal position embeddings, no layernorm afterwards
							if "large-printed" in checkpoint_url or "stage1" in checkpoint_url:
														_UpperCAmelCase											=     False
														_UpperCAmelCase											=     'relu'
														_UpperCAmelCase											=     1_0_2_4
														_UpperCAmelCase											=     True
														_UpperCAmelCase											=     False
														_UpperCAmelCase											=     False

							# load HuggingFace model
							_UpperCAmelCase											=     ViTModel(a__     ,    add_pooling_layer=a__			)
							_UpperCAmelCase											=     TrOCRForCausalLM(a__			)
							_UpperCAmelCase											=     VisionEncoderDecoderModel(encoder=a__     ,    decoder=a__			)
							model.eval()

							# load state_dict of original model, rename some keys
							_UpperCAmelCase											=     torch.hub.load_state_dict_from_url(a__     ,    map_location='cpu'     ,    check_hash=a__			)['model']

							_UpperCAmelCase											=     create_rename_keys(a__     ,    a__			)
							for src, dest in rename_keys:
														rename_key(a__     ,    a__     ,    a__			)
							read_in_q_k_v(a__     ,    a__			)

							# remove parameters we don't need
							del state_dict["encoder.deit.head.weight"]
							del state_dict["encoder.deit.head.bias"]
							del state_dict["decoder.version"]

							# add prefix to decoder keys
							for key, val in state_dict.copy().items():
														_UpperCAmelCase											=     state_dict.pop(a__			)
														if key.startswith('decoder'			) and "output_projection" not in key:
																					_UpperCAmelCase											=     val
														else:
																					_UpperCAmelCase											=     val

    # load state dict
							model.load_state_dict(a__			)

							# Check outputs on an image
							_UpperCAmelCase											=     ViTImageProcessor(size=encoder_config.image_size			)
							_UpperCAmelCase											=     RobertaTokenizer.from_pretrained('roberta-large'			)
							_UpperCAmelCase											=     TrOCRProcessor(a__     ,    a__			)

							_UpperCAmelCase											=     processor(images=prepare_img(a__			)     ,    return_tensors='pt'			).pixel_values

							# verify logits
							_UpperCAmelCase											=     torch.tensor([[model.config.decoder.decoder_start_token_id]]			)
							_UpperCAmelCase											=     model(pixel_values=a__     ,    decoder_input_ids=a__			)
							_UpperCAmelCase											=     outputs.logits

							_UpperCAmelCase											=     torch.Size([1, 1, 5_0_2_6_5]			)
							if "trocr-base-handwritten" in checkpoint_url:
														_UpperCAmelCase											=     torch.tensor(
														    [-1.4_502, -4.6_683, -0.5_347, -2.9_291, 9.1_435, -3.0_571, 8.9_764, 1.7_560, 8.7_358, -1.5_311]			)
							elif "trocr-large-handwritten" in checkpoint_url:
														_UpperCAmelCase											=     torch.tensor(
														    [-2.6_437, -1.3_129, -2.2_596, -5.3_455, 6.3_539, 1.7_604, 5.4_991, 1.4_702, 5.6_113, 2.0_170]			)
							elif "trocr-base-printed" in checkpoint_url:
														_UpperCAmelCase											=     torch.tensor(
														    [-5.6_816, -5.8_388, 1.1_398, -6.9_034, 6.8_505, -2.4_393, 1.2_284, -1.0_232, -1.9_661, -3.9_210]			)
							elif "trocr-large-printed" in checkpoint_url:
														_UpperCAmelCase											=     torch.tensor(
														    [-6.0_162, -7.0_959, 4.4_155, -5.1_063, 7.0_468, -3.1_631, 2.6_466, -0.3_081, -0.8_106, -1.7_535]			)

							if "stage1" not in checkpoint_url:
														assert logits.shape == expected_shape, "Shape of logits not as expected"
														assert torch.allclose(logits[0, 0, :1_0]     ,    a__     ,    atol=1e-3			), "First elements of logits not as expected"

							Path(a__			).mkdir(exist_ok=a__			)
							print(F'''Saving model to {pytorch_dump_folder_path}'''			)
							model.save_pretrained(a__			)
							print(F'''Saving processor to {pytorch_dump_folder_path}'''			)
							processor.save_pretrained(a__			)


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

		parser.add_argument(
		    '''--checkpoint_url''',
		    default='''https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt''',
		    type=str,
		    help='''URL to the original PyTorch checkpoint (.pth file).''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
		)
		lowerCAmelCase__    :Optional[int]       =				parser.parse_args()
		convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_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		)},
									        ],
									    ]     ,       )
