

import copy
import os
from typing import TYPE_CHECKING, List, Union


if TYPE_CHECKING:
						pass

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase							:  str   	=						logging.get_logger(__name__)

lowerCAmelCase							:  Tuple   	=						{
    'kakaobrain/align-base': 'https://huggingface.co/kakaobrain/align-base/resolve/main/config.json',
}



class   __lowercase				( UpperCAmelCase_	):




			"""simple docstring"""







			_UpperCAmelCase   :							Optional[int]										=     '''align_text_model'''



			def __init__( self  :     Union[str, Any]   ,						lowerCAmelCase__  :     Tuple=3_0522   ,						lowerCAmelCase__  :     List[Any]=768   ,						lowerCAmelCase__  :     Optional[Any]=12   ,						lowerCAmelCase__  :     Optional[int]=12   ,						lowerCAmelCase__  :     List[Any]=3072   ,						lowerCAmelCase__  :     Optional[int]="gelu"   ,						lowerCAmelCase__  :     Optional[Any]=0.1   ,						lowerCAmelCase__  :     int=0.1   ,						lowerCAmelCase__  :     List[str]=512   ,						lowerCAmelCase__  :     Any=2   ,						lowerCAmelCase__  :     Optional[Any]=0.02   ,						lowerCAmelCase__  :     Optional[Any]=1E-12   ,						lowerCAmelCase__  :     Union[str, Any]=0   ,						lowerCAmelCase__  :     Tuple="absolute"   ,						lowerCAmelCase__  :     List[Any]=True   ,						**lowerCAmelCase__  :     str   ,						):
									super().__init__(**A__)

									SCREAMING_SNAKE_CASE_: List[Any]  		=       vocab_size
									SCREAMING_SNAKE_CASE_: List[str]  		=       hidden_size
									SCREAMING_SNAKE_CASE_: Any  		=       num_hidden_layers
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       num_attention_heads
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       hidden_act
									SCREAMING_SNAKE_CASE_: Tuple  		=       intermediate_size
									SCREAMING_SNAKE_CASE_: str  		=       hidden_dropout_prob
									SCREAMING_SNAKE_CASE_: List[Any]  		=       attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       max_position_embeddings
									SCREAMING_SNAKE_CASE_: str  		=       type_vocab_size
									SCREAMING_SNAKE_CASE_: List[Any]  		=       initializer_range
									SCREAMING_SNAKE_CASE_: int  		=       layer_norm_eps
									SCREAMING_SNAKE_CASE_: Tuple  		=       position_embedding_type
									SCREAMING_SNAKE_CASE_: int  		=       use_cache
									SCREAMING_SNAKE_CASE_: Tuple  		=       pad_token_id




			@classmethod
			def        _SCREAMING_SNAKE_CASE			( cls  :     List[Any]   ,						lowerCAmelCase__  :     Dict   ,						**lowerCAmelCase__  :     Any):
									cls._set_token_in_kwargs(A__)

									SCREAMING_SNAKE_CASE_: Dict  		=       cls.get_config_dict(A__   ,						**A__)

									# get the text config dict if we are loading from AlignConfig
									if config_dict.get("model_type") == "align":
															SCREAMING_SNAKE_CASE_: Any  		=       config_dict["""text_config"""]

									if "model_type" in config_dict and hasattr(cls   ,						"model_type") and config_dict["model_type"] != cls.model_type:
															logger.warning(
															    F"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
															    F"{cls.model_type}. This is not supported for all configurations of models and can yield errors.")

									return cls.from_dict(A__   ,						**A__)







class   __lowercase				( UpperCAmelCase_	):




			"""simple docstring"""







			_UpperCAmelCase   :							Union[str, Any]										=     '''align_vision_model'''



			def __init__( self  :     List[str]   ,						lowerCAmelCase__  :     Optional[Any] = 3   ,						lowerCAmelCase__  :     str = 600   ,						lowerCAmelCase__  :     Optional[Any] = 2.0   ,						lowerCAmelCase__  :     Dict = 3.1   ,						lowerCAmelCase__  :     Tuple = 8   ,						lowerCAmelCase__  :     List[str] = [3, 3, 5, 3, 5, 5, 3]   ,						lowerCAmelCase__  :     List[str] = [32, 16, 24, 40, 80, 112, 192]   ,						lowerCAmelCase__  :     List[Any] = [16, 24, 40, 80, 112, 192, 320]   ,						lowerCAmelCase__  :     Union[str, Any] = []   ,						lowerCAmelCase__  :     Union[str, Any] = [1, 2, 2, 2, 1, 2, 1]   ,						lowerCAmelCase__  :     Tuple = [1, 2, 2, 3, 3, 4, 1]   ,						lowerCAmelCase__  :     Any = [1, 6, 6, 6, 6, 6, 6]   ,						lowerCAmelCase__  :     str = 0.25   ,						lowerCAmelCase__  :     Dict = "swish"   ,						lowerCAmelCase__  :     Union[str, Any] = 2560   ,						lowerCAmelCase__  :     int = "mean"   ,						lowerCAmelCase__  :     List[str] = 0.02   ,						lowerCAmelCase__  :     Optional[Any] = 0.001   ,						lowerCAmelCase__  :     Union[str, Any] = 0.99   ,						lowerCAmelCase__  :     Union[str, Any] = 0.2   ,						**lowerCAmelCase__  :     int   ,						):
									super().__init__(**A__)

									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       num_channels
									SCREAMING_SNAKE_CASE_: Dict  		=       image_size
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       width_coefficient
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       depth_coefficient
									SCREAMING_SNAKE_CASE_: str  		=       depth_divisor
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       kernel_sizes
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       in_channels
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       out_channels
									SCREAMING_SNAKE_CASE_: List[Any]  		=       depthwise_padding
									SCREAMING_SNAKE_CASE_: str  		=       strides
									SCREAMING_SNAKE_CASE_: int  		=       num_block_repeats
									SCREAMING_SNAKE_CASE_: List[Any]  		=       expand_ratios
									SCREAMING_SNAKE_CASE_: Dict  		=       squeeze_expansion_ratio
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       hidden_act
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       hidden_dim
									SCREAMING_SNAKE_CASE_: List[Any]  		=       pooling_type
									SCREAMING_SNAKE_CASE_: List[Any]  		=       initializer_range
									SCREAMING_SNAKE_CASE_: int  		=       batch_norm_eps
									SCREAMING_SNAKE_CASE_: int  		=       batch_norm_momentum
									SCREAMING_SNAKE_CASE_: Dict  		=       drop_connect_rate
									SCREAMING_SNAKE_CASE_: List[str]  		=       sum(A__) * 4




			@classmethod
			def        _SCREAMING_SNAKE_CASE			( cls  :     Union[str, Any]   ,						lowerCAmelCase__  :     Dict   ,						**lowerCAmelCase__  :     Optional[int]):
									cls._set_token_in_kwargs(A__)

									SCREAMING_SNAKE_CASE_: List[Any]  		=       cls.get_config_dict(A__   ,						**A__)

									# get the vision config dict if we are loading from AlignConfig
									if config_dict.get("model_type") == "align":
															SCREAMING_SNAKE_CASE_: Tuple  		=       config_dict["""vision_config"""]

									if "model_type" in config_dict and hasattr(cls   ,						"model_type") and config_dict["model_type"] != cls.model_type:
															logger.warning(
															    F"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
															    F"{cls.model_type}. This is not supported for all configurations of models and can yield errors.")

									return cls.from_dict(A__   ,						**A__)







class   __lowercase				( UpperCAmelCase_	):




			"""simple docstring"""







			_UpperCAmelCase   :							Optional[Any]										=     '''align'''
			_UpperCAmelCase   :							Any										=     True



			def __init__( self  :     Dict   ,						lowerCAmelCase__  :     Optional[Any]=None   ,						lowerCAmelCase__  :     Optional[Any]=None   ,						lowerCAmelCase__  :     Tuple=640   ,						lowerCAmelCase__  :     List[str]=1.0   ,						lowerCAmelCase__  :     List[str]=0.02   ,						**lowerCAmelCase__  :     List[str]   ,						):
									super().__init__(**A__)

									if text_config is None:
															SCREAMING_SNAKE_CASE_: Dict  		=       {}
															logger.info("text_config is None. Initializing the AlignTextConfig with default values.")

									if vision_config is None:
															SCREAMING_SNAKE_CASE_: str  		=       {}
															logger.info("vision_config is None. Initializing the AlignVisionConfig with default values.")

									SCREAMING_SNAKE_CASE_: Dict  		=       AlignTextConfig(**A__)
									SCREAMING_SNAKE_CASE_: List[Any]  		=       AlignVisionConfig(**A__)

									SCREAMING_SNAKE_CASE_: int  		=       projection_dim
									SCREAMING_SNAKE_CASE_: List[Any]  		=       temperature_init_value
									SCREAMING_SNAKE_CASE_: str  		=       initializer_range



			@classmethod
			def        _SCREAMING_SNAKE_CASE			( cls  :     Optional[Any]   ,						lowerCAmelCase__  :     str   ,						lowerCAmelCase__  :     Optional[int]   ,						**lowerCAmelCase__  :     List[Any]):
									return cls(text_config=text_config.to_dict()   ,						vision_config=vision_config.to_dict()   ,						**A__)




			def        _SCREAMING_SNAKE_CASE			( self  :     Optional[int]):
									SCREAMING_SNAKE_CASE_: int  		=       copy.deepcopy(self.__dict__)
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       self.text_config.to_dict()
									SCREAMING_SNAKE_CASE_: Any  		=       self.vision_config.to_dict()
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       self.__class__.model_type
									return output









import argparse

import torch

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


logging.set_verbosity_info()

def 						UpperCamelCase					(lowercase_: Optional[int]    ,      lowercase_: Union[str, Any]    ,      lowercase_: Optional[Any]					)     ->      Tuple:
	# Initialise PyTorch model
	A__							:    str        =		AlbertConfig.from_json_file(lowercase_					)
	print(f"""Building PyTorch model from configuration: {config}"""					)
	A__							:    List[Any]        =		AlbertForPreTraining(lowercase_					)

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

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


if __name__ == "__main__":
	A_       : 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.'
	)
	A_       : int         =			parser.parse_args()
	convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.albert_config_file, args.pytorch_dump_path)



