




'''simple docstring'''




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

import torch_xla.distributed.xla_multiprocessing as xmp

def    _UpperCamelCase      (				):
 '''simple docstring'''





 UpperCAmelCase__										=    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=lowerCAmelCase__       ,       default=1       ,       help="""Number of TPU cores to use (1 or 8)."""       )

 # positional
 parser.add_argument(
     """training_script"""       ,       type=lowerCAmelCase__       ,       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=lowerCAmelCase__       )

 return parser.parse_args()

def    _UpperCamelCase      (				):
 '''simple docstring'''





 UpperCAmelCase__										=    parse_args()

 # Import training_script as a module.
 UpperCAmelCase__										=    Path(args.training_script       )
 sys.path.append(str(script_fpath.parent.resolve()       )       )
 UpperCAmelCase__										=    script_fpath.stem
 UpperCAmelCase__										=    importlib.import_module(lowerCAmelCase__       )

 # Patch sys.argv
 UpperCAmelCase__										=    [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()






'''simple docstring'''




import argparse

from transformers import TaConfig, TaForConditionalGeneration, load_tf_weights_in_ta
from transformers.utils import logging


logging.set_verbosity_info()

def    _UpperCamelCase      (				SCREAMING_SNAKE_CASE__			:      List[str]       ,						SCREAMING_SNAKE_CASE__			:      str       ,						SCREAMING_SNAKE_CASE__			:      Tuple       ):
 '''simple docstring'''





 UpperCAmelCase__										=    TaConfig.from_json_file(SCREAMING_SNAKE_CASE__       )
 print(F'''Building PyTorch model from configuration: {config}'''       )
 UpperCAmelCase__										=    TaForConditionalGeneration(SCREAMING_SNAKE_CASE__       )

 # Load weights from tf checkpoint
 load_tf_weights_in_ta(SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       )

 # Save pytorch-model
 print(F'''Save PyTorch model to {pytorch_dump_path}'''       )
 model.save_pretrained(SCREAMING_SNAKE_CASE__       )


if __name__ == "__main__":
    UpperCAmelCase_       				=      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(
        '--config_file',
        default=None,
        type=str,
        required=True,
        help=(
            'The config json file corresponding to the pre-trained T5 model. \nThis specifies the model architecture.'
        ),
    )
    parser.add_argument(
        '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
    )
    UpperCAmelCase_       				=      parser.parse_args()
    convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)
