



import argparse

import torch

from transformers import ImageGPTConfig, ImageGPTForCausalLM, load_tf_weights_in_imagegpt
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()




def 	_UpperCAmelCase       (   snake_case			,		snake_case			,		snake_case    ):





					"""simple docstring"""

					_lowerCAmelCase      			=					{"""small""": (5_12, 8, 24), """medium""": (10_24, 8, 36), """large""": (15_36, 16, 48)}
					_lowerCAmelCase			,			_lowerCAmelCase			,			_lowerCAmelCase      			=					MODELS[model_size]  # set model hyperparameters
					_lowerCAmelCase      			=					ImageGPTConfig(n_embd=snake_case			,		n_layer=snake_case			,		n_head=snake_case    )
					_lowerCAmelCase      			=					ImageGPTForCausalLM(snake_case    )

					# Load weights from numpy
					load_tf_weights_in_imagegpt(snake_case			,		snake_case			,		snake_case    )

					# Save pytorch-model
					_lowerCAmelCase      			=					pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
					_lowerCAmelCase      			=					pytorch_dump_folder_path + """/""" + CONFIG_NAME
					print(F'Save PyTorch model to {pytorch_weights_dump_path}'    )
					torch.save(model.state_dict()			,		snake_case    )
					print(F'Save configuration file to {pytorch_config_dump_path}'    )
					with open(snake_case			,		"""w"""			,		encoding="""utf-8"""    ) as f:
										f.write(config.to_json_string()    )


if __name__ == "__main__":
	A__			   = argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    """--imagegpt_checkpoint_path""",
	    default=None,
	    type=str,
	    required=True,
	    help="""Path to the TensorFlow checkpoint path.""",
	)
	parser.add_argument(
	    """--model_size""",
	    default=None,
	    type=str,
	    required=True,
	    help="""Size of the model (can be either 'small', 'medium' or 'large').""",
	)
	parser.add_argument(
	    """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
	)
	A__			   = parser.parse_args()
	convert_imagegpt_checkpoint_to_pytorch(
	    args.imagegpt_checkpoint_path, args.model_size, args.pytorch_dump_folder_path
	)








from __future__ import annotations


def 	_UpperCAmelCase       (   snake_case    ):





					"""simple docstring"""

					_lowerCAmelCase      			=					sorted(snake_case    )
					_lowerCAmelCase      			=					len(snake_case    )
					_lowerCAmelCase      			=					length >> 1
					return (
					    (sorted_list[mid_index] + sorted_list[mid_index - 1]) / 2
					    if length % 2 == 0
					    else sorted_list[mid_index]
					)


def 	_UpperCAmelCase       (   ):





					"""simple docstring"""

					import doctest

					doctest.testmod()


if __name__ == "__main__":
	main()



