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


			"""simple docstring"""






			__a    :					Optional[Any]     					=		MODEL_FOR_MASKED_LM_MAPPING
			__a    :					Optional[int]     					=		TF_MODEL_FOR_MASKED_LM_MAPPING





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   str:

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

							'''simple docstring'''


							__lowercase					     =				pipeline(task='''fill-mask''' ,  model='''sshleifer/tiny-distilroberta-base''' ,  top_k=2 ,  framework='''tf'''		)
							__lowercase					     =				unmasker('''My name is <mask>'''		)
							self.assertEqual(
							    nested_simplify(lowerCAmelCase__ ,  decimals=6		) ,  [
							        {'''sequence''': '''My name is grouped''', '''score''': 2.1E-05, '''token''': 3_80_15, '''token_str''': ''' grouped'''},
							        {'''sequence''': '''My name is accuser''', '''score''': 2.1E-05, '''token''': 2_55_06, '''token_str''': ''' accuser'''},
							    ] ,  )

							__lowercase					     =				unmasker('''The largest city in France is <mask>'''		)
							self.assertEqual(
							    nested_simplify(lowerCAmelCase__ ,  decimals=6		) ,  [
							        {
							            '''sequence''': '''The largest city in France is grouped''',
							            '''score''': 2.1E-05,
							            '''token''': 3_80_15,
							            '''token_str''': ''' grouped''',
							        },
							        {
							            '''sequence''': '''The largest city in France is accuser''',
							            '''score''': 2.1E-05,
							            '''token''': 2_55_06,
							            '''token_str''': ''' accuser''',
							        },
							    ] ,  )

							__lowercase					     =				unmasker('''My name is <mask>''' ,  targets=[''' Patrick''', ''' Clara''', ''' Teven'''] ,  top_k=3		)
							self.assertEqual(
							    nested_simplify(lowerCAmelCase__ ,  decimals=6		) ,  [
							        {'''sequence''': '''My name is Clara''', '''score''': 2E-05, '''token''': 1_36_06, '''token_str''': ''' Clara'''},
							        {'''sequence''': '''My name is Patrick''', '''score''': 2E-05, '''token''': 34_99, '''token_str''': ''' Patrick'''},
							        {'''sequence''': '''My name is Te''', '''score''': 1.9E-05, '''token''': 29_41, '''token_str''': ''' Te'''},
							    ] ,  )





			@require_torch
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Optional[int]:

							'''simple docstring'''


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

							__lowercase					     =				unmasker('''My name is <mask>'''		)
							self.assertEqual(
							    nested_simplify(lowerCAmelCase__ ,  decimals=6		) ,  [
							        {'''sequence''': '''My name is Maul''', '''score''': 2.2E-05, '''token''': 3_56_76, '''token_str''': ''' Maul'''},
							        {'''sequence''': '''My name isELS''', '''score''': 2.2E-05, '''token''': 1_64_16, '''token_str''': '''ELS'''},
							    ] ,  )

							__lowercase					     =				unmasker('''The largest city in France is <mask>'''		)
							self.assertEqual(
							    nested_simplify(lowerCAmelCase__ ,  decimals=6		) ,  [
							        {
							            '''sequence''': '''The largest city in France is Maul''',
							            '''score''': 2.2E-05,
							            '''token''': 3_56_76,
							            '''token_str''': ''' Maul''',
							        },
							        {'''sequence''': '''The largest city in France isELS''', '''score''': 2.2E-05, '''token''': 1_64_16, '''token_str''': '''ELS'''},
							    ] ,  )

							__lowercase					     =				unmasker('''My name is <mask>''' ,  targets=[''' Patrick''', ''' Clara''', ''' Teven'''] ,  top_k=3		)
							self.assertEqual(
							    nested_simplify(lowerCAmelCase__ ,  decimals=6		) ,  [
							        {'''sequence''': '''My name is Patrick''', '''score''': 2.1E-05, '''token''': 34_99, '''token_str''': ''' Patrick'''},
							        {'''sequence''': '''My name is Te''', '''score''': 2E-05, '''token''': 29_41, '''token_str''': ''' Te'''},
							        {'''sequence''': '''My name is Clara''', '''score''': 2E-05, '''token''': 1_36_06, '''token_str''': ''' Clara'''},
							    ] ,  )

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

							self.assertEqual(
							    nested_simplify(lowerCAmelCase__ ,  decimals=6		) ,  [
							        [
							            {
							                '''score''': 2.2E-05,
							                '''token''': 3_56_76,
							                '''token_str''': ''' Maul''',
							                '''sequence''': '''<s>My name is Maul<mask></s>''',
							            },
							            {'''score''': 2.2E-05, '''token''': 1_64_16, '''token_str''': '''ELS''', '''sequence''': '''<s>My name isELS<mask></s>'''},
							        ],
							        [
							            {
							                '''score''': 2.2E-05,
							                '''token''': 3_56_76,
							                '''token_str''': ''' Maul''',
							                '''sequence''': '''<s>My name is<mask> Maul</s>''',
							            },
							            {'''score''': 2.2E-05, '''token''': 1_64_16, '''token_str''': '''ELS''', '''sequence''': '''<s>My name is<mask>ELS</s>'''},
							        ],
							    ] ,  )





			@require_torch_gpu
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Union[str, Any]:

							'''simple docstring'''


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

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

							__lowercase					     =				pipe('''Paris is the [MASK] of France.'''		)
							# We actually don't care about the result, we just want to make sure
							# it works, meaning the float16 tensor got casted back to float32
							# for postprocessing.
							self.assertIsInstance(lowerCAmelCase__ ,  lowerCAmelCase__		)





			@slow
			@require_torch
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   str:

							'''simple docstring'''


							__lowercase					     =				pipeline(task='''fill-mask''' ,  model='''distilroberta-base''' ,  top_k=2 ,  framework='''pt'''		)
							self.run_large_test(lowerCAmelCase__		)





			@slow
			@require_tf
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Any:

							'''simple docstring'''


							__lowercase					     =				pipeline(task='''fill-mask''' ,  model='''distilroberta-base''' ,  top_k=2 ,  framework='''tf'''		)
							self.run_large_test(lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__		)	->   Dict:

							'''simple docstring'''


							__lowercase					     =				unmasker('''My name is <mask>'''		)
							self.assertEqual(
							    nested_simplify(lowerCAmelCase__		) ,  [
							        {'''sequence''': '''My name is John''', '''score''': 0.008, '''token''': 6_10, '''token_str''': ''' John'''},
							        {'''sequence''': '''My name is Chris''', '''score''': 0.007, '''token''': 15_73, '''token_str''': ''' Chris'''},
							    ] ,  )
							__lowercase					     =				unmasker('''The largest city in France is <mask>'''		)
							self.assertEqual(
							    nested_simplify(lowerCAmelCase__		) ,  [
							        {
							            '''sequence''': '''The largest city in France is Paris''',
							            '''score''': 0.251,
							            '''token''': 22_01,
							            '''token_str''': ''' Paris''',
							        },
							        {
							            '''sequence''': '''The largest city in France is Lyon''',
							            '''score''': 0.214,
							            '''token''': 1_27_90,
							            '''token_str''': ''' Lyon''',
							        },
							    ] ,  )

							__lowercase					     =				unmasker('''My name is <mask>''' ,  targets=[''' Patrick''', ''' Clara''', ''' Teven'''] ,  top_k=3		)
							self.assertEqual(
							    nested_simplify(lowerCAmelCase__		) ,  [
							        {'''sequence''': '''My name is Patrick''', '''score''': 0.005, '''token''': 34_99, '''token_str''': ''' Patrick'''},
							        {'''sequence''': '''My name is Clara''', '''score''': 0.000, '''token''': 1_36_06, '''token_str''': ''' Clara'''},
							        {'''sequence''': '''My name is Te''', '''score''': 0.000, '''token''': 29_41, '''token_str''': ''' Te'''},
							    ] ,  )





			@require_torch
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   List[Any]:

							'''simple docstring'''


							__lowercase					     =				pipeline(task='''fill-mask''' ,  model='''sshleifer/tiny-distilroberta-base''' ,  framework='''pt'''		)
							__lowercase					     =				None
							__lowercase					     =				None
							self.run_pipeline_test(lowerCAmelCase__ ,  []		)





			@require_tf
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   List[Any]:

							'''simple docstring'''


							__lowercase					     =				pipeline(task='''fill-mask''' ,  model='''sshleifer/tiny-distilroberta-base''' ,  framework='''tf'''		)
							__lowercase					     =				None
							__lowercase					     =				None
							self.run_pipeline_test(lowerCAmelCase__ ,  []		)





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

							'''simple docstring'''


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

							__lowercase					     =				FillMaskPipeline(model=lowerCAmelCase__ ,  tokenizer=lowerCAmelCase__		)
							__lowercase					     =				[
							    F"This is another {tokenizer.mask_token} test",
							]
							return fill_masker, examples





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   Union[str, Any]:

							'''simple docstring'''


							__lowercase					     =				fill_masker.tokenizer
							__lowercase					     =				fill_masker.model

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

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

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

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

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





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   Union[str, Any]:

							'''simple docstring'''


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

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

							# Score equivalence
							__lowercase					     =				fill_masker(F"This is a {tokenizer.mask_token}" ,  targets=lowerCAmelCase__		)
							__lowercase					     =				[top_mask['''token_str'''] for top_mask in outputs]
							__lowercase					     =				[top_mask['''score'''] for top_mask in outputs]

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

							# Raises with invalid
							with self.assertRaises(lowerCAmelCase__		):
											__lowercase					     =				fill_masker(F"This is a {tokenizer.mask_token}" ,  targets=[]		)
							# For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
							if "" not in tokenizer.get_vocab():
											with self.assertRaises(lowerCAmelCase__		):
															__lowercase					     =				fill_masker(F"This is a {tokenizer.mask_token}" ,  targets=['''''']		)
											with self.assertRaises(lowerCAmelCase__		):
															__lowercase					     =				fill_masker(F"This is a {tokenizer.mask_token}" ,  targets=''''''		)





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

							'''simple docstring'''


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

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





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   Optional[int]:

							'''simple docstring'''


							__lowercase					     =				tokenizer.get_vocab()
							__lowercase					     =				FillMaskPipeline(model=lowerCAmelCase__ ,  tokenizer=lowerCAmelCase__		)

							# top_k=2, ntargets=3
							__lowercase					     =				sorted(vocab.keys()		)[:3]
							__lowercase					     =				fill_masker(F"This is a {tokenizer.mask_token}" ,  top_k=2 ,  targets=lowerCAmelCase__		)

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





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   Optional[int]:

							'''simple docstring'''


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

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





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   List[str]:

							'''simple docstring'''


							__lowercase					     =				FillMaskPipeline(model=lowerCAmelCase__ ,  tokenizer=lowerCAmelCase__		)

							__lowercase					     =				fill_masker(
							    F"This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}" ,  top_k=2		)
							self.assertEqual(
							    lowerCAmelCase__ ,  [
							        [
							            {'''sequence''': ANY(lowerCAmelCase__		), '''score''': ANY(lowerCAmelCase__		), '''token''': ANY(lowerCAmelCase__		), '''token_str''': ANY(lowerCAmelCase__		)},
							            {'''sequence''': ANY(lowerCAmelCase__		), '''score''': ANY(lowerCAmelCase__		), '''token''': ANY(lowerCAmelCase__		), '''token_str''': ANY(lowerCAmelCase__		)},
							        ],
							        [
							            {'''sequence''': ANY(lowerCAmelCase__		), '''score''': ANY(lowerCAmelCase__		), '''token''': ANY(lowerCAmelCase__		), '''token_str''': ANY(lowerCAmelCase__		)},
							            {'''sequence''': ANY(lowerCAmelCase__		), '''score''': ANY(lowerCAmelCase__		), '''token''': ANY(lowerCAmelCase__		), '''token_str''': ANY(lowerCAmelCase__		)},
							        ],
							        [
							            {'''sequence''': ANY(lowerCAmelCase__		), '''score''': ANY(lowerCAmelCase__		), '''token''': ANY(lowerCAmelCase__		), '''token_str''': ANY(lowerCAmelCase__		)},
							            {'''sequence''': ANY(lowerCAmelCase__		), '''score''': ANY(lowerCAmelCase__		), '''token''': ANY(lowerCAmelCase__		), '''token_str''': ANY(lowerCAmelCase__		)},
							        ],
							    ] ,  )
import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SwiftFormerConfig,
    SwiftFormerForImageClassification,
    ViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__a    :	Optional[Any]         =    logging.get_logger(__name__)

__a    :	Optional[Any]         =    torch.device("""cpu""")





def 							UpperCAmelCase		(       ):







				"""simple docstring"""



				__lowercase					     =				'''http://images.cocodataset.org/val2017/000000039769.jpg'''
				__lowercase					     =				Image.open(requests.get(lowercase				,     stream=lowercase				).raw				)
				return im





def 							UpperCAmelCase		(       lowercase				):







				"""simple docstring"""



				if swiftformer_name == "swiftformer_xs":
								return torch.tensor([-2.1703E00, 2.1107E00, -2.0811E00, 8.8685E-01, 2.4360E-01]				)

				elif swiftformer_name == "swiftformer_s":
								return torch.tensor([3.9636E-01, 2.3478E-01, -1.6963E00, -1.7381E00, -8.6337E-01]				)

				elif swiftformer_name == "swiftformer_l1":
								return torch.tensor([-4.2768E-01, -4.7429E-01, -1.0897E00, -1.0248E00, 3.5523E-02]				)

				elif swiftformer_name == "swiftformer_l3":
								return torch.tensor([-2.5330E-01, 2.4211E-01, -6.0185E-01, -8.2789E-01, -6.0446E-02]				)





def 							UpperCAmelCase		(       lowercase				,     lowercase				,     lowercase				):







				"""simple docstring"""



				__lowercase					     =				dct.pop(lowercase				)
				__lowercase					     =				val





def 							UpperCAmelCase		(       lowercase				):







				"""simple docstring"""



				__lowercase					     =				[]
				for k in state_dict.keys():
								__lowercase					     =				k
								if ".pwconv" in k:
												__lowercase					     =				k_new.replace('''.pwconv'''				,     '''.point_wise_conv'''				)
								if ".dwconv" in k:
												__lowercase					     =				k_new.replace('''.dwconv'''				,     '''.depth_wise_conv'''				)
								if ".Proj." in k:
												__lowercase					     =				k_new.replace('''.Proj.'''				,     '''.proj.'''				)
								if "patch_embed" in k_new:
												__lowercase					     =				k_new.replace('''patch_embed'''				,     '''swiftformer.patch_embed.patch_embedding'''				)
								if "network" in k_new:
												__lowercase					     =				k_new.split('''.'''				)
												if ls[2].isdigit():
																__lowercase					     =				'''swiftformer.encoder.network.''' + ls[1] + '''.blocks.''' + ls[2] + '''.''' + '''.'''.join(ls[3:]				)
												else:
																__lowercase					     =				k_new.replace('''network'''				,     '''swiftformer.encoder.network'''				)
								rename_keys.append((k, k_new)				)
				return rename_keys





@torch.no_grad()
def 							UpperCAmelCase		(       lowercase				,     lowercase				,     lowercase				):







				"""simple docstring"""



				__lowercase					     =				SwiftFormerConfig()

				# dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
				__lowercase					     =				1000
				__lowercase					     =				'''huggingface/label-files'''
				__lowercase					     =				'''imagenet-1k-id2label.json'''
				__lowercase					     =				json.load(open(hf_hub_download(lowercase				,     lowercase				,     repo_type='''dataset'''				)				,     '''r'''				)				)
				__lowercase					     =				{int(lowercase				): v for k, v in idalabel.items()}
				__lowercase					     =				idalabel
				__lowercase					     =				{v: k for k, v in idalabel.items()}

				# size of the architecture
				if swiftformer_name == "swiftformer_xs":
								__lowercase					     =				[3, 3, 6, 4]
								__lowercase					     =				[48, 56, 112, 220]

				elif swiftformer_name == "swiftformer_s":
								__lowercase					     =				[3, 3, 9, 6]
								__lowercase					     =				[48, 64, 168, 224]

				elif swiftformer_name == "swiftformer_l1":
								__lowercase					     =				[4, 3, 10, 5]
								__lowercase					     =				[48, 96, 192, 384]

				elif swiftformer_name == "swiftformer_l3":
								__lowercase					     =				[4, 4, 12, 6]
								__lowercase					     =				[64, 128, 320, 512]

				# load state_dict of original model, remove and rename some keys
				if original_ckpt:
								if original_ckpt.startswith('''https'''				):
												__lowercase					     =				torch.hub.load_state_dict_from_url(lowercase				,     map_location='''cpu'''				,     check_hash=lowercase				)
								else:
												__lowercase					     =				torch.load(lowercase				,     map_location='''cpu'''				)
				__lowercase					     =				checkpoint

				__lowercase					     =				create_rename_keys(lowercase				)
				for rename_key_src, rename_key_dest in rename_keys:
								rename_key(lowercase				,     lowercase				,     lowercase				)

				# load HuggingFace model
				__lowercase					     =				SwiftFormerForImageClassification(lowercase				).eval()
				hf_model.load_state_dict(lowercase				)

				# prepare test inputs
				__lowercase					     =				prepare_img()
				__lowercase					     =				ViTImageProcessor.from_pretrained('''preprocessor_config'''				)
				__lowercase					     =				processor(images=lowercase				,     return_tensors='''pt'''				)

				# compare outputs from both models
				__lowercase					     =				get_expected_output(lowercase				)
				__lowercase					     =				hf_model(inputs['''pixel_values''']				).logits

				assert hf_logits.shape == torch.Size([1, 1000]				)
				assert torch.allclose(hf_logits[0, 0:5]				,     lowercase				,     atol=1E-3				)

				Path(lowercase				).mkdir(exist_ok=lowercase				)
				print(F"Saving model {swiftformer_name} to {pytorch_dump_folder_path}"				)
				hf_model.save_pretrained(lowercase				)


if __name__ == "__main__":
		__a    :	Optional[int]         =    argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    """--swiftformer_name""",
		    default="""swiftformer_xs""",
		    choices=["""swiftformer_xs""", """swiftformer_s""", """swiftformer_l1""", """swiftformer_l3"""],
		    type=str,
		    help="""Name of the SwiftFormer model you'd like to convert.""",
		)
		parser.add_argument(
		    """--pytorch_dump_folder_path""",
		    default="""./converted_outputs/""",
		    type=str,
		    help="""Path to the output PyTorch model directory.""",
		)
		parser.add_argument("""--original_ckpt""", default=None, type=str, help="""Path to the original model checkpoint.""")

		__a    :	Optional[int]         =    parser.parse_args()
		convert_swiftformer_checkpoint(args.swiftformer_name, args.pytorch_dump_folder_path, args.original_ckpt)