
import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel



def       _lowercase					(      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE )	->     int:
			UpperCamelCase__   ,			UpperCamelCase__   ,			UpperCamelCase__    :		Any     =				pt_weights[F"""{prefix}.in_proj_weight"""].chunk(3				,    dim=0 )
			UpperCamelCase__   ,			UpperCamelCase__   ,			UpperCamelCase__    :		Optional[int]     =				pt_weights[F"""{prefix}.in_proj_bias"""].chunk(3				,    dim=0 )

			UpperCamelCase__    :		Dict     =				pt_weights[F"""{prefix}.out_proj.weight"""]
			UpperCamelCase__    :		int     =				pt_weights[F"""{prefix}.out_proj.bias"""]

			UpperCamelCase__    :		Optional[Any]     =				q_proj
			UpperCamelCase__    :		int     =				q_proj_bias

			UpperCamelCase__    :		List[str]     =				k_proj
			UpperCamelCase__    :		Tuple     =				k_proj_bias

			UpperCamelCase__    :		Optional[Any]     =				v_proj
			UpperCamelCase__    :		Any     =				v_proj_bias

			UpperCamelCase__    :		List[str]     =				out_proj_weights
			UpperCamelCase__    :		List[Any]     =				out_proj_bias



def       _lowercase					(      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE )	->     str:
			copy_linear(hf_mlp.fca				,    __SCREAMING_SNAKE_CASE				,    F"""{prefix}.c_fc""" )
			copy_linear(hf_mlp.fca				,    __SCREAMING_SNAKE_CASE				,    F"""{prefix}.c_proj""" )



def       _lowercase					(      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE )	->     str:
			UpperCamelCase__    :		str     =				pt_weights[F"""{prefix}.weight"""].data
			UpperCamelCase__    :		Union[str, Any]     =				pt_weights[F"""{prefix}.bias"""].data



def       _lowercase					(      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE )	->     int:
			copy_linear(hf_layer.layer_norma				,    __SCREAMING_SNAKE_CASE				,    F"""{prefix}.ln_1""" )
			copy_linear(hf_layer.layer_norma				,    __SCREAMING_SNAKE_CASE				,    F"""{prefix}.ln_2""" )

			# copy MLP
			copy_mlp(hf_layer.mlp				,    __SCREAMING_SNAKE_CASE				,    F"""{prefix}.mlp""" )

			# copy attn
			copy_attn_layer(hf_layer.self_attn				,    __SCREAMING_SNAKE_CASE				,    F"""{prefix}.attn""" )



def       _lowercase					(      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE )	->     Dict:
			for layer_id, hf_layer in enumerate(__SCREAMING_SNAKE_CASE ):
						copy_layer(__SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE				,    F"""{prefix}.{layer_id}""" )



def       _lowercase					(      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE )	->     List[str]:
			UpperCamelCase__    :		Dict     =				pt_weights['text_projection'].data.T

			# copy text encoder
			for name, param in hf_model.text_model.named_parameters():
						UpperCamelCase__    :		List[Any]     =				pt_weights[F"""bert.{name}"""].data



def       _lowercase					(      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE )	->     Any:
			UpperCamelCase__    :		Optional[Any]     =				pt_weights['visual.proj'].data.T

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

			# copy embeddings
			UpperCamelCase__    :		Optional[Any]     =				pt_weights['visual.conv1.weight'].data
			UpperCamelCase__    :		Optional[Any]     =				pt_weights['visual.class_embedding'].data
			UpperCamelCase__    :		int     =				pt_weights['visual.positional_embedding'].data

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



@torch.no_grad()
def       _lowercase					(      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE=None )	->     List[str]:
			assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
			UpperCamelCase__    :		List[Any]     =				ChineseCLIPConfig.from_pretrained(__SCREAMING_SNAKE_CASE )

			UpperCamelCase__    :		Tuple     =				ChineseCLIPModel(__SCREAMING_SNAKE_CASE ).eval()

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

			copy_text_model_and_projection(__SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE )
			copy_vision_model_and_projection(__SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE )
			UpperCamelCase__    :		List[str]     =				pt_weights['logit_scale'].data

			hf_model.save_pretrained(__SCREAMING_SNAKE_CASE )


if __name__ == "__main__":
							UpperCAmelCase__				:  Tuple       					=				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.'''
							)
							UpperCAmelCase__				:  Optional[Any]       					=				parser.parse_args()

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





from __future__ import annotations

import re



def       _lowercase					(      __SCREAMING_SNAKE_CASE )	->     list[str]:

			def alphanum_key(__SCREAMING_SNAKE_CASE ):
						return [int(__SCREAMING_SNAKE_CASE ) if s.isdigit() else s.lower() for s in re.split('([0-9]+)'				,    __SCREAMING_SNAKE_CASE )]

			return sorted(__SCREAMING_SNAKE_CASE				,    key=__SCREAMING_SNAKE_CASE )


if __name__ == "__main__":
							import doctest

							doctest.testmod()



