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							):
  SCREAMING_SNAKE_CASE_	:	Tuple						=MODEL_FOR_MASKED_LM_MAPPING
  SCREAMING_SNAKE_CASE_	:	int						=TF_MODEL_FOR_MASKED_LM_MAPPING






  def 		_a					(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 		_a					(self	)		->      Tuple:





       '''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-05, '''token''': 3_8015, '''token_str''': ''' grouped'''},
               {'''sequence''': '''My name is accuser''', '''score''': 2.1E-05, '''token''': 2_5506, '''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-05,
                   '''token''': 3_8015,
                   '''token_str''': ''' grouped''',
               },
               {
                   '''sequence''': '''The largest city in France is accuser''',
                   '''score''': 2.1E-05,
                   '''token''': 2_5506,
                   '''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-05, '''token''': 1_3606, '''token_str''': ''' Clara'''},
               {'''sequence''': '''My name is Patrick''', '''score''': 2E-05, '''token''': 3499, '''token_str''': ''' Patrick'''},
               {'''sequence''': '''My name is Te''', '''score''': 1.9E-05, '''token''': 2941, '''token_str''': ''' Te'''},
           ]   ,				)






  @require_torch
  def 		_a					(self	)		->      Tuple:





       '''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-05, '''token''': 3_5676, '''token_str''': ''' Maul'''},
               {'''sequence''': '''My name isELS''', '''score''': 2.2E-05, '''token''': 1_6416, '''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-05,
                   '''token''': 3_5676,
                   '''token_str''': ''' Maul''',
               },
               {'''sequence''': '''The largest city in France isELS''', '''score''': 2.2E-05, '''token''': 1_6416, '''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-05, '''token''': 3499, '''token_str''': ''' Patrick'''},
               {'''sequence''': '''My name is Te''', '''score''': 2E-05, '''token''': 2941, '''token_str''': ''' Te'''},
               {'''sequence''': '''My name is Clara''', '''score''': 2E-05, '''token''': 1_3606, '''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-05,
                       '''token''': 3_5676,
                       '''token_str''': ''' Maul''',
                       '''sequence''': '''<s>My name is Maul<mask></s>''',
                   },
                   {'''score''': 2.2E-05, '''token''': 1_6416, '''token_str''': '''ELS''', '''sequence''': '''<s>My name isELS<mask></s>'''},
               ],
               [
                   {
                       '''score''': 2.2E-05,
                       '''token''': 3_5676,
                       '''token_str''': ''' Maul''',
                       '''sequence''': '''<s>My name is<mask> Maul</s>''',
                   },
                   {'''score''': 2.2E-05, '''token''': 1_6416, '''token_str''': '''ELS''', '''sequence''': '''<s>My name is<mask>ELS</s>'''},
               ],
           ]   ,				)






  @require_torch_gpu
  def 		_a					(self	)		->      Tuple:





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





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





       '''simple docstring'''





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






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_	)		->      List[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''': 1_2790,
                   '''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''': 1_3606, '''token_str''': ''' Clara'''},
               {'''sequence''': '''My name is Te''', '''score''': 0.000, '''token''': 2941, '''token_str''': ''' Te'''},
           ]   ,				)






  @require_torch
  def 		_a					(self	)		->      Optional[Any]:





       '''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 		_a					(self	)		->      Union[str, Any]:





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





       '''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 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      str:





       '''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 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Optional[int]:





       '''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 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Union[str, Any]:





       '''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 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Dict:





       '''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 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Union[str, Any]:





       '''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_	)},
               ],
           ]   ,				)




import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to properly calculate the metrics on the
# validation dataset when in a distributed system, and builds off the
# `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


__magic_name__								=16
__magic_name__								=32







def  __UpperCamelCase    (						A ,							A = 16							):
     UpperCamelCase__										=      AutoTokenizer.from_pretrained('''bert-base-cased'''							)
     UpperCamelCase__										=      load_dataset('''glue''' ,							'''mrpc'''							)

     def tokenize_function(A							):
          # max_length=None => use the model max length (it's actually the default)
          UpperCamelCase__										=      tokenizer(examples['''sentence1'''] ,							examples['''sentence2'''] ,							truncation=A ,							max_length=A							)
          return outputs

     # Apply the method we just defined to all the examples in all the splits of the dataset
     # starting with the main process first:
     with accelerator.main_process_first():
          UpperCamelCase__										=      datasets.map(
              A ,							batched=A ,							remove_columns=['''idx''', '''sentence1''', '''sentence2'''] ,							)

     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
     UpperCamelCase__										=      tokenized_datasets.rename_column('''label''' ,							'''labels'''							)

     def collate_fn(A							):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          UpperCamelCase__										=      128 if accelerator.distributed_type == DistributedType.TPU else None
          # When using mixed precision we want round multiples of 8/16
          if accelerator.mixed_precision == "fp8":
               UpperCamelCase__										=      16
          elif accelerator.mixed_precision != "no":
               UpperCamelCase__										=      8
          else:
               UpperCamelCase__										=      None

          return tokenizer.pad(
              A ,							padding='''longest''' ,							max_length=A ,							pad_to_multiple_of=A ,							return_tensors='''pt''' ,							)

     # Instantiate dataloaders.
     UpperCamelCase__										=      DataLoader(
         tokenized_datasets['''train'''] ,							shuffle=A ,							collate_fn=A ,							batch_size=A							)
     UpperCamelCase__										=      DataLoader(
         tokenized_datasets['''validation'''] ,							shuffle=A ,							collate_fn=A ,							batch_size=A							)

     return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('''TESTING_MOCKED_DATALOADERS''', None) == "1":
     from accelerate.test_utils.training import mocked_dataloaders

     __magic_name__								=mocked_dataloaders  # noqa: F811







def  __UpperCamelCase    (						A ,							A							):
     # For testing only
     if os.environ.get('''TESTING_MOCKED_DATALOADERS''' ,							A							) == "1":
          UpperCamelCase__										=      2
     # Initialize accelerator
     UpperCamelCase__										=      Accelerator(cpu=args.cpu ,							mixed_precision=args.mixed_precision							)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     UpperCamelCase__										=      config['''lr''']
     UpperCamelCase__										=      int(config['''num_epochs''']							)
     UpperCamelCase__										=      int(config['''seed''']							)
     UpperCamelCase__										=      int(config['''batch_size''']							)

     UpperCamelCase__										=      evaluate.load('''glue''' ,							'''mrpc'''							)

     # If the batch size is too big we use gradient accumulation
     UpperCamelCase__										=      1
     if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
          UpperCamelCase__										=      batch_size // MAX_GPU_BATCH_SIZE
          UpperCamelCase__										=      MAX_GPU_BATCH_SIZE

     set_seed(A							)
     UpperCamelCase__						,       UpperCamelCase__										=      get_dataloaders(A ,							A							)
     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     UpperCamelCase__										=      AutoModelForSequenceClassification.from_pretrained('''bert-base-cased''' ,							return_dict=A							)

     # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
     # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
     # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
     UpperCamelCase__										=      model.to(accelerator.device							)

     # Instantiate optimizer
     UpperCamelCase__										=      AdamW(params=model.parameters() ,							lr=A							)

     # Instantiate scheduler
     UpperCamelCase__										=      get_linear_schedule_with_warmup(
         optimizer=A ,							num_warmup_steps=100 ,							num_training_steps=(len(A							) * num_epochs) // gradient_accumulation_steps ,							)

     # Prepare everything
     # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
     # prepare method.
     UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      accelerator.prepare(
         A ,							A ,							A ,							A ,							A							)

     # Now we train the model
     for epoch in range(A							):
          model.train()
          for step, batch in enumerate(A							):
               # We could avoid this line since we set the accelerator with `device_placement=True`.
               batch.to(accelerator.device							)
               UpperCamelCase__										=      model(**A							)
               UpperCamelCase__										=      outputs.loss
               UpperCamelCase__										=      loss / gradient_accumulation_steps
               accelerator.backward(A							)
               if step % gradient_accumulation_steps == 0:
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()

          model.eval()
          UpperCamelCase__										=      0
          for step, batch in enumerate(A							):
               # We could avoid this line since we set the accelerator with `device_placement=True`.
               batch.to(accelerator.device							)
               with torch.no_grad():
                    UpperCamelCase__										=      model(**A							)
               UpperCamelCase__										=      outputs.logits.argmax(dim=-1							)
               UpperCamelCase__						,       UpperCamelCase__										=      accelerator.gather((predictions, batch['''labels'''])							)
               # New Code #
               # First we check if it's a distributed system
               if accelerator.use_distributed:
                    # Then see if we're on the last batch of our eval dataloader
                    if step == len(A							) - 1:
                         # Last batch needs to be truncated on distributed systems as it contains additional samples
                         UpperCamelCase__										=      predictions[: len(eval_dataloader.dataset							) - samples_seen]
                         UpperCamelCase__										=      references[: len(eval_dataloader.dataset							) - samples_seen]
                    else:
                         # Otherwise we add the number of samples seen
                         samples_seen += references.shape[0]
            # All of this can be avoided if you use `Accelerator.gather_for_metrics` instead of `Accelerator.gather`:
            # accelerator.gather_for_metrics((predictions, batch["labels"]))
               metric.add_batch(
                   predictions=A ,							references=A ,							)

          UpperCamelCase__										=      metric.compute()
          # Use accelerator.print to print only on the main process.
          accelerator.print(f"epoch {epoch}:" ,							A							)







def  __UpperCamelCase    (						):
     UpperCamelCase__										=      argparse.ArgumentParser(description='''Simple example of training script.'''							)
     parser.add_argument(
         '''--mixed_precision''' ,							type=A ,							default=A ,							choices=['''no''', '''fp16''', '''bf16''', '''fp8'''] ,							help='''Whether to use mixed precision. Choose'''
         '''between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'''
         '''and an Nvidia Ampere GPU.''' ,							)
     parser.add_argument('''--cpu''' ,							action='''store_true''' ,							help='''If passed, will train on the CPU.'''							)
     UpperCamelCase__										=      parser.parse_args()
     UpperCamelCase__										=      {'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
     training_function(A ,							A							)


if __name__ == "__main__":
     main()



