
"""simple docstring"""




import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					SCREAMING_SNAKE_CASE           =      pt_weights[f'''{prefix}.in_proj_weight'''].chunk(3,	dim=0     )
					SCREAMING_SNAKE_CASE           =      pt_weights[f'''{prefix}.in_proj_bias'''].chunk(3,	dim=0     )

					SCREAMING_SNAKE_CASE           =      pt_weights[f'''{prefix}.out_proj.weight''']
					SCREAMING_SNAKE_CASE           =      pt_weights[f'''{prefix}.out_proj.bias''']

					SCREAMING_SNAKE_CASE           =      q_proj
					SCREAMING_SNAKE_CASE           =      q_proj_bias

					SCREAMING_SNAKE_CASE           =      k_proj
					SCREAMING_SNAKE_CASE           =      k_proj_bias

					SCREAMING_SNAKE_CASE           =      v_proj
					SCREAMING_SNAKE_CASE           =      v_proj_bias

					SCREAMING_SNAKE_CASE           =      out_proj_weights
					SCREAMING_SNAKE_CASE           =      out_proj_bias

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					copy_linear(hf_mlp.fca,	_UpperCAmelCase,	f'''{prefix}.c_fc'''     )
					copy_linear(hf_mlp.fca,	_UpperCAmelCase,	f'''{prefix}.c_proj'''     )

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					SCREAMING_SNAKE_CASE           =      pt_weights[f'''{prefix}.weight'''].data
					SCREAMING_SNAKE_CASE           =      pt_weights[f'''{prefix}.bias'''].data

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					copy_linear(hf_layer.layer_norma,	_UpperCAmelCase,	f'''{prefix}.ln_1'''     )
					copy_linear(hf_layer.layer_norma,	_UpperCAmelCase,	f'''{prefix}.ln_2'''     )

					# copy MLP
					copy_mlp(hf_layer.mlp,	_UpperCAmelCase,	f'''{prefix}.mlp'''     )

					# copy attn
					copy_attn_layer(hf_layer.self_attn,	_UpperCAmelCase,	f'''{prefix}.attn'''     )

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):




					for layer_id, hf_layer in enumerate(_UpperCAmelCase     ):
										copy_layer(_UpperCAmelCase,	_UpperCAmelCase,	f'''{prefix}.{layer_id}'''     )

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					SCREAMING_SNAKE_CASE           =      pt_weights['text_projection'].data.T

					# copy text encoder
					for name, param in hf_model.text_model.named_parameters():
										SCREAMING_SNAKE_CASE           =      pt_weights[f'''bert.{name}'''].data

def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_     ):
					SCREAMING_SNAKE_CASE           =      pt_weights['visual.proj'].data.T

					# copy layer norms
					copy_linear(hf_model.vision_model.pre_layrnorm,	_UpperCAmelCase,	'visual.ln_pre'     )
					copy_linear(hf_model.vision_model.post_layernorm,	_UpperCAmelCase,	'visual.ln_post'     )

					# copy embeddings
					SCREAMING_SNAKE_CASE           =      pt_weights['visual.conv1.weight'].data
					SCREAMING_SNAKE_CASE           =      pt_weights['visual.class_embedding'].data
					SCREAMING_SNAKE_CASE           =      pt_weights['visual.positional_embedding'].data

					# copy encoder
					copy_layers(hf_model.vision_model.encoder.layers,	_UpperCAmelCase,	'visual.transformer.resblocks'     )

@torch.no_grad()
def     UpperCamelCase_    (						SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_=None     ):
					assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
					SCREAMING_SNAKE_CASE           =      ChineseCLIPConfig.from_pretrained(_UpperCAmelCase     )

					SCREAMING_SNAKE_CASE           =      ChineseCLIPModel(_UpperCAmelCase     ).eval()

					SCREAMING_SNAKE_CASE           =      torch.load(_UpperCAmelCase,	map_location='cpu'     )['state_dict']
					SCREAMING_SNAKE_CASE           =      {(name[7:] if name.startswith('module.'     ) else name): value for name, value in pt_weights.items()}

					copy_text_model_and_projection(_UpperCAmelCase,	_UpperCAmelCase     )
					copy_vision_model_and_projection(_UpperCAmelCase,	_UpperCAmelCase     )
					SCREAMING_SNAKE_CASE           =      pt_weights['logit_scale'].data

					hf_model.save_pretrained(_UpperCAmelCase     )


if __name__ == "__main__":
	snake_case					  =						argparse.ArgumentParser()
	parser.add_argument(
	    '--pytorch_dump_folder_path',
	    default=None,
	    type=str,
	    help='Path to the output folder storing converted hf PyTorch model.',
	)
	parser.add_argument(
	    '--checkpoint_path', default=None, type=str, help='Path to original github format ChineseCLIP checkpoint.'
	)
	parser.add_argument(
	    '--config_path', default=None, required=True, type=str, help='Path to hf config.json of model to convert.'
	)
	snake_case					  =						parser.parse_args()

	convert_chinese_clip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)
	print('The conversion is finished!')







'''simple docstring'''


from typing import Callable




def       A							(      _UpperCAmelCase   :				Callable					)  ->				Dict:


  '''simple docstring'''




  def wrapper(_UpperCAmelCase   :				Optional[Any]					):
    __lowerCAmelCase		:    List[Any]			     =  function_with_docstring.__doc__
    return target_function

  return wrapper





