import argparse

import torch

from transformers import RemBertConfig, RemBertModel, load_tf_weights_in_rembert
from transformers.utils import logging


logging.set_verbosity_info()
def lowerCamelCase  (       SCREAMING_SNAKE_CASE  ,      SCREAMING_SNAKE_CASE  ,      SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







    __UpperCamelCase					:Optional[Any]        =		RemBertConfig.from_json_file(lowercase__     )
    print('''Building PyTorch model from configuration: {}'''.format(str(lowercase__     )     )     )
    __UpperCamelCase					:List[str]        =		RemBertModel(lowercase__     )

    # Load weights from tf checkpoint
    load_tf_weights_in_rembert(lowercase__  ,      lowercase__  ,      lowercase__     )

    # Save pytorch-model
    print('''Save PyTorch model to {}'''.format(lowercase__     )     )
    torch.save(model.state_dict()  ,      lowercase__     )


if __name__ == "__main__":
    __lowercase										= 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(
        '''--rembert_config_file''',
        default=None,
        type=str,
        required=True,
        help=(
            '''The config json file corresponding to the pre-trained RemBERT 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.'''
    )
    __lowercase										= parser.parse_args()
    convert_rembert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.rembert_config_file, args.pytorch_dump_path)






import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp



def lowerCamelCase  (       ):
    '''simple docstring'''







    __UpperCamelCase					:int        =		ArgumentParser(
        description=(
            '''PyTorch TPU distributed training launch '''
            '''helper utility that will spawn up '''
            '''multiple distributed processes'''
        )     )

    # Optional arguments for the launch helper
    parser.add_argument('''--num_cores'''  ,      type=SCREAMING_SNAKE_CASE  ,      default=1  ,      help='''Number of TPU cores to use (1 or 8).'''     )

    # positional
    parser.add_argument(
        '''training_script'''  ,      type=SCREAMING_SNAKE_CASE  ,      help=(
            '''The full path to the single TPU training '''
            '''program/script to be launched in parallel, '''
            '''followed by all the arguments for the '''
            '''training script'''
        )  ,      )

    # rest from the training program
    parser.add_argument('''training_script_args'''  ,      nargs=SCREAMING_SNAKE_CASE     )

    return parser.parse_args()



def lowerCamelCase  (       ):
    '''simple docstring'''







    __UpperCamelCase					:Tuple        =		parse_args()

    # Import training_script as a module.
    __UpperCamelCase					:Dict        =		Path(args.training_script     )
    sys.path.append(str(script_fpath.parent.resolve()     )     )
    __UpperCamelCase					:Any        =		script_fpath.stem
    __UpperCamelCase					:Union[str, Any]        =		importlib.import_module(SCREAMING_SNAKE_CASE     )

    # Patch sys.argv
    __UpperCamelCase					:Dict        =		[args.training_script] + args.training_script_args + ['''--tpu_num_cores''', str(args.num_cores     )]
    xmp.spawn(mod._mp_fn  ,      args=()  ,      nprocs=args.num_cores     )


if __name__ == "__main__":
    main()





