


import argparse

import torch

from ...utils import logging
from . import AlbertConfig, AlbertForPreTraining, load_tf_weights_in_albert


logging.set_verbosity_info()

def 			lowerCAmelCase__					(     a__:						Optional[int]     ,    a__:						List[Any]     ,    a__:						Any			)    ->						int:


							'''simple docstring'''





							_UpperCAmelCase											=     AlbertConfig.from_json_file(a__			)
							print(F'''Building PyTorch model from configuration: {config}'''			)
							_UpperCAmelCase											=     AlbertForPreTraining(a__			)

							# Load weights from tf checkpoint
							load_tf_weights_in_albert(a__     ,    a__     ,    a__			)

							# Save pytorch-model
							print(F'''Save PyTorch model to {pytorch_dump_path}'''			)
							torch.save(model.state_dict()     ,    a__			)


if __name__ == "__main__":
		lowerCAmelCase__    :Optional[int]       =				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(
		    '''--albert_config_file''',
		    default=None,
		    type=str,
		    required=True,
		    help=(
		        '''The config json file corresponding to the pre-trained ALBERT 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__    :str       =				parser.parse_args()
		convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.albert_config_file, args.pytorch_dump_path)




from .imports import is_rich_available


if is_rich_available():
		from rich.traceback import install

		install(show_locals=False)

else:
		raise ModuleNotFoundError('''To use the rich extension, install rich with `pip install rich`''')
