
import argparse
import json

from transformers import GPTNeoConfig, GPTNeoForCausalLM, load_tf_weights_in_gpt_neo
from transformers.utils import logging


logging.set_verbosity_info()




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  Dict		,      SCREAMING_SNAKE_CASE_  :  Optional[int]		,      SCREAMING_SNAKE_CASE_  :  Tuple					)     ->    Dict:
						# Initialise PyTorch model
						_lowercase									=				json.load(open(SCREAMING_SNAKE_CASE_		,      """r"""					)					)
						_lowercase									=				GPTNeoConfig(
						    hidden_size=config_json["""n_embd"""]		,      num_layers=config_json["""n_layer"""]		,      num_heads=config_json["""n_head"""]		,      attention_types=config_json["""attention_types"""]		,      max_position_embeddings=config_json["""n_positions"""]		,      resid_dropout=config_json["""res_dropout"""]		,      embed_dropout=config_json["""embed_dropout"""]		,      attention_dropout=config_json["""attn_dropout"""]		,      )
						print(f"""Building PyTorch model from configuration: {config}"""					)
						_lowercase									=				GPTNeoForCausalLM(SCREAMING_SNAKE_CASE_					)

						# Load weights from tf checkpoint
						load_tf_weights_in_gpt_neo(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__":
		A		:   List[str]         =    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 mesh-tf 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		:   Dict         =    parser.parse_args()
		convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)

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


A		:   List[Any]         =    logging.get_logger(__name__)

A		:   str         =    {
    '''MCG-NJU/videomae-base''': '''https://huggingface.co/MCG-NJU/videomae-base/resolve/main/config.json''',
}



class       a_   ( _a						):
			a  : str				      =       '''videomae'''






			def __init__(			self							, __UpperCamelCase=224							, __UpperCamelCase=16							, __UpperCamelCase=3							, __UpperCamelCase=16							, __UpperCamelCase=2							, __UpperCamelCase=768							, __UpperCamelCase=12							, __UpperCamelCase=12							, __UpperCamelCase=3_072							, __UpperCamelCase="gelu"							, __UpperCamelCase=0.0							, __UpperCamelCase=0.0							, __UpperCamelCase=0.02							, __UpperCamelCase=1E-12							, __UpperCamelCase=True							, __UpperCamelCase=True							, __UpperCamelCase=6							, __UpperCamelCase=384							, __UpperCamelCase=4							, __UpperCamelCase=1_536							, __UpperCamelCase=True							, **__UpperCamelCase							, ):
									super().__init__(**__UpperCamelCase     )

									_lowercase									=				image_size
									_lowercase									=				patch_size
									_lowercase									=				num_channels
									_lowercase									=				num_frames
									_lowercase									=				tubelet_size

									_lowercase									=				hidden_size
									_lowercase									=				num_hidden_layers
									_lowercase									=				num_attention_heads
									_lowercase									=				intermediate_size
									_lowercase									=				hidden_act
									_lowercase									=				hidden_dropout_prob
									_lowercase									=				attention_probs_dropout_prob
									_lowercase									=				initializer_range
									_lowercase									=				layer_norm_eps
									_lowercase									=				qkv_bias
									_lowercase									=				use_mean_pooling

									_lowercase									=				decoder_num_attention_heads
									_lowercase									=				decoder_hidden_size
									_lowercase									=				decoder_num_hidden_layers
									_lowercase									=				decoder_intermediate_size
									_lowercase									=				norm_pix_loss