




"""simple docstring"""





from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
					if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.25.0""")):
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					from ...utils.dummy_torch_and_transformers_objects import UnCLIPImageVariationPipeline, UnCLIPPipeline
else:
					from .pipeline_unclip import UnCLIPPipeline
					from .pipeline_unclip_image_variation import UnCLIPImageVariationPipeline
					from .text_proj import UnCLIPTextProjModel








import argparse
import os

import torch

from transformers.utils import WEIGHTS_NAME


a__		     =     ['''small''', '''medium''', '''large''']

a__		     =     '''lm_head.decoder.weight'''
a__		     =     '''lm_head.weight'''





def 				A__ (snake_case   :		str		,	snake_case   :		str				)     ->						Union[str, Any]:
					__UpperCamelCase					:	Union[str, Any]             =      torch.load(snake_case				)
					__UpperCamelCase					:	int             =      d.pop(snake_case				)
					os.makedirs(snake_case		,	exist_ok=snake_case				)
					torch.save(snake_case		,	os.path.join(snake_case		,	snake_case				)				)


if __name__ == "__main__":
	a__		     =     argparse.ArgumentParser()
	parser.add_argument('''--dialogpt_path''', default='''.''', type=str)
	a__		     =     parser.parse_args()
	for MODEL in DIALOGPT_MODELS:
		a__		     =     os.path.join(args.dialogpt_path, f"{MODEL}_ft.pkl")
		a__		     =     f"./DialoGPT-{MODEL}"
		convert_dialogpt_checkpoint(
		    checkpoint_path,
		    pytorch_dump_folder_path,
		)






