
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					):
   A__   :						Optional[Any]			=	MODEL_FOR_MASKED_LM_MAPPING
   A__   :						Union[str, Any]			=	TF_MODEL_FOR_MASKED_LM_MAPPING



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


          '''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						:   str					)						->      Any:


          '''simple docstring'''


          snake_case			:  List[Any]     							= pipeline(task="fill-mask"       ,		model="sshleifer/tiny-distilroberta-base"       ,		top_k=2       ,		framework="tf"					)
          snake_case			:  Union[str, Any]     							= unmasker("My name is <mask>"					)
          self.assertEqual(
              nested_simplify(snake_case__       ,		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"},
              ]       ,		)

          snake_case			:  Union[str, Any]     							= unmasker("The largest city in France is <mask>"					)
          self.assertEqual(
              nested_simplify(snake_case__       ,		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",
                  },
              ]       ,		)

          snake_case			:  List[Any]     							= unmasker("My name is <mask>"       ,		targets=[" Patrick", " Clara", " Teven"]       ,		top_k=3					)
          self.assertEqual(
              nested_simplify(snake_case__       ,		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						:   Dict					)						->      Union[str, Any]:


          '''simple docstring'''


          snake_case			:  Tuple     							= pipeline(task="fill-mask"       ,		model="sshleifer/tiny-distilroberta-base"       ,		top_k=2       ,		framework="pt"					)

          snake_case			:  Dict     							= unmasker("My name is <mask>"					)
          self.assertEqual(
              nested_simplify(snake_case__       ,		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"},
              ]       ,		)

          snake_case			:  Optional[Any]     							= unmasker("The largest city in France is <mask>"					)
          self.assertEqual(
              nested_simplify(snake_case__       ,		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"},
              ]       ,		)

          snake_case			:  int     							= unmasker("My name is <mask>"       ,		targets=[" Patrick", " Clara", " Teven"]       ,		top_k=3					)
          self.assertEqual(
              nested_simplify(snake_case__       ,		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"},
              ]       ,		)

          snake_case			:  Any     							= unmasker("My name is <mask> <mask>"       ,		top_k=2					)

          self.assertEqual(
              nested_simplify(snake_case__       ,		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						:   Dict					)						->      Dict:


          '''simple docstring'''


          snake_case			:  int     							= pipeline("fill-mask"       ,		model="hf-internal-testing/tiny-random-distilbert"       ,		device=0       ,		framework="pt"					)

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

          snake_case			:  Union[str, Any]     							= 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(snake_case__       ,		snake_case__					)



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


          '''simple docstring'''


          snake_case			:  Optional[int]     							= pipeline(task="fill-mask"       ,		model="distilroberta-base"       ,		top_k=2       ,		framework="pt"					)
          self.run_large_test(snake_case__					)



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


          '''simple docstring'''


          snake_case			:  Optional[Any]     							= pipeline(task="fill-mask"       ,		model="distilroberta-base"       ,		top_k=2       ,		framework="tf"					)
          self.run_large_test(snake_case__					)



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


          '''simple docstring'''


          snake_case			:  Any     							= unmasker("My name is <mask>"					)
          self.assertEqual(
              nested_simplify(snake_case__					)       ,		[
                  {"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"},
              ]       ,		)
          snake_case			:  List[str]     							= unmasker("The largest city in France is <mask>"					)
          self.assertEqual(
              nested_simplify(snake_case__					)       ,		[
                  {
                      "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",
                  },
              ]       ,		)

          snake_case			:  List[Any]     							= unmasker("My name is <mask>"       ,		targets=[" Patrick", " Clara", " Teven"]       ,		top_k=3					)
          self.assertEqual(
              nested_simplify(snake_case__					)       ,		[
                  {"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						:   Union[str, Any]					)						->      Dict:


          '''simple docstring'''


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



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


          '''simple docstring'''


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



   def 							_SCREAMING_SNAKE_CASE  (self						:   int       ,		snake_case__						:   Dict       ,		snake_case__						:   int       ,		snake_case__						:   Tuple					)						->      str:


          '''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)"					)

          snake_case			:  Union[str, Any]     							= FillMaskPipeline(model=snake_case__       ,		tokenizer=snake_case__					)
          snake_case			:  List[str]     							= [
              f"""This is another {tokenizer.mask_token} test""",
          ]
          return fill_masker, examples



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


          '''simple docstring'''


          snake_case			:  Union[str, Any]     							= fill_masker.tokenizer
          snake_case			:  Any     							= fill_masker.model

          snake_case			:  List[str]     							= fill_masker(
              f"""This is a {tokenizer.mask_token}"""       ,		)
          self.assertEqual(
              snake_case__       ,		[
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
              ]       ,		)

          snake_case			:  str     							= fill_masker([f"""This is a {tokenizer.mask_token}"""]					)
          self.assertEqual(
              snake_case__       ,		[
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
              ]       ,		)

          snake_case			:  Union[str, Any]     							= fill_masker([f"""This is a {tokenizer.mask_token}""", f"""Another {tokenizer.mask_token} great test."""]					)
          self.assertEqual(
              snake_case__       ,		[
                  [
                      {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                      {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                      {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                      {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                      {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  ],
                  [
                      {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                      {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                      {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                      {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                      {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  ],
              ]       ,		)

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

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



   def 							_SCREAMING_SNAKE_CASE  (self						:   List[Any]       ,		snake_case__						:   Any       ,		snake_case__						:   Optional[Any]					)						->      List[str]:


          '''simple docstring'''


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

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

          # Score equivalence
          snake_case			:  Any     							= fill_masker(f"""This is a {tokenizer.mask_token}"""       ,		targets=snake_case__					)
          snake_case			:  Optional[Any]     							= [top_mask["token_str"] for top_mask in outputs]
          snake_case			:  Optional[Any]     							= [top_mask["score"] for top_mask in outputs]

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

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



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


          '''simple docstring'''


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

          snake_case			:  List[str]     							= FillMaskPipeline(model=snake_case__       ,		tokenizer=snake_case__					)
          snake_case			:  int     							= fill_masker(f"""This is a {tokenizer.mask_token}"""       ,		top_k=2					)
          self.assertEqual(
              snake_case__       ,		[
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
                  {"sequence": ANY(snake_case__					), "score": ANY(snake_case__					), "token": ANY(snake_case__					), "token_str": ANY(snake_case__					)},
              ]       ,		)
          self.assertEqual(nested_simplify(snake_case__					)       ,		nested_simplify(snake_case__					)					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   List[str]       ,		snake_case__						:   List[Any]       ,		snake_case__						:   int					)						->      Optional[Any]:


          '''simple docstring'''


          snake_case			:  List[str]     							= tokenizer.get_vocab()
          snake_case			:  Optional[int]     							= FillMaskPipeline(model=snake_case__       ,		tokenizer=snake_case__					)

          # top_k=2, ntargets=3
          snake_case			:  Any     							= sorted(vocab.keys()					)[:3]
          snake_case			:  Tuple     							= fill_masker(f"""This is a {tokenizer.mask_token}"""       ,		top_k=2       ,		targets=snake_case__					)

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



   def 							_SCREAMING_SNAKE_CASE  (self						:   Any       ,		snake_case__						:   Any       ,		snake_case__						:   List[Any]					)						->      Optional[int]:


          '''simple docstring'''


          snake_case			:  Union[str, Any]     							= FillMaskPipeline(model=snake_case__       ,		tokenizer=snake_case__					)
          snake_case			:  str     							= tokenizer.get_vocab()
          # String duplicates + id duplicates
          snake_case			:  Dict     							= sorted(vocab.keys()					)[:3]
          snake_case			:  List[Any]     							= [targets[0], targets[1], targets[0], targets[2], targets[1]]
          snake_case			:  Dict     							= fill_masker(f"""My name is {tokenizer.mask_token}"""       ,		targets=snake_case__       ,		top_k=10					)

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



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


          '''simple docstring'''


          snake_case			:  str     							= FillMaskPipeline(model=snake_case__       ,		tokenizer=snake_case__					)

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







import argparse

import torch

from transformers import BertConfig, BertForPreTraining, load_tf_weights_in_bert
from transformers.utils import logging


logging.set_verbosity_info()


def      UpperCamelCase   (       __lowerCamelCase	:					List[str]							,   __lowerCamelCase	:					Union[str, Any]							,   __lowerCamelCase	:					Union[str, Any]       ):
       # Initialise PyTorch model
       snake_case			:  Optional[int]     							= BertConfig.from_json_file(__lowerCamelCase       )
       print(f"""Building PyTorch model from configuration: {config}"""       )
       snake_case			:  Any     							= BertForPreTraining(__lowerCamelCase       )

       # Load weights from tf checkpoint
       load_tf_weights_in_bert(__lowerCamelCase							,   __lowerCamelCase							,   __lowerCamelCase       )

       # Save pytorch-model
       print(f"""Save PyTorch model to {pytorch_dump_path}"""       )
       torch.save(model.state_dict()							,   __lowerCamelCase       )


if __name__ == "__main__":
   __lowerCamelCase				=			argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       """--tf_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
   )
   parser.add_argument(
       """--bert_config_file""",
       default=None,
       type=str,
       required=True,
       help=(
           """The config json file corresponding to the pre-trained BERT model. \n"""
           """This specifies the model architecture."""
       ),
   )
   parser.add_argument(
       """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
   )
   __lowerCamelCase				=			parser.parse_args()
   convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)





