






'''simple docstring'''







import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			Tuple:
							lowerCamelCase__    ,						lowerCamelCase__    ,						lowerCamelCase__   :      int			    =			pt_weights[f'''{prefix}.in_proj_weight'''].chunk(3							,     dim=0 )
							lowerCamelCase__    ,						lowerCamelCase__    ,						lowerCamelCase__   :      Optional[Any]			    =			pt_weights[f'''{prefix}.in_proj_bias'''].chunk(3							,     dim=0 )

							lowerCamelCase__   :      List[str]			    =			pt_weights[f'''{prefix}.out_proj.weight''']
							lowerCamelCase__   :      Dict			    =			pt_weights[f'''{prefix}.out_proj.bias''']

							lowerCamelCase__   :      List[Any]			    =			q_proj
							lowerCamelCase__   :      Optional[Any]			    =			q_proj_bias

							lowerCamelCase__   :      Optional[int]			    =			k_proj
							lowerCamelCase__   :      int			    =			k_proj_bias

							lowerCamelCase__   :      int			    =			v_proj
							lowerCamelCase__   :      Any			    =			v_proj_bias

							lowerCamelCase__   :      Optional[Any]			    =			out_proj_weights
							lowerCamelCase__   :      Any			    =			out_proj_bias





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			Optional[Any]:
							copy_linear(hf_mlp.fca							,     UpperCamelCase							,     f'''{prefix}.c_fc''' )
							copy_linear(hf_mlp.fca							,     UpperCamelCase							,     f'''{prefix}.c_proj''' )





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			Tuple:
							lowerCamelCase__   :      int			    =			pt_weights[f'''{prefix}.weight'''].data
							lowerCamelCase__   :      List[str]			    =			pt_weights[f'''{prefix}.bias'''].data





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			Union[str, Any]:
							# copy layer norms
							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 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			List[str]:
							for layer_id, hf_layer in enumerate(UpperCamelCase ):
														copy_layer(UpperCamelCase							,     UpperCamelCase							,     f'''{prefix}.{layer_id}''' )





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase ) ->			List[Any]:
							# copy projection
							lowerCamelCase__   :      int			    =			pt_weights["""text_projection"""].data.T

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





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase ) ->			Union[str, Any]:
							# copy projection
							lowerCamelCase__   :      Dict			    =			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
							lowerCamelCase__   :      int			    =			pt_weights["""visual.conv1.weight"""].data
							lowerCamelCase__   :      Optional[Any]			    =			pt_weights["""visual.class_embedding"""].data
							lowerCamelCase__   :      List[Any]			    =			pt_weights["""visual.positional_embedding"""].data

							# copy encoder
							copy_layers(hf_model.vision_model.encoder.layers							,     UpperCamelCase							,     """visual.transformer.resblocks""" )





@torch.no_grad()
def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase=None ) ->			str:
							assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
							lowerCamelCase__   :      Dict			    =			ChineseCLIPConfig.from_pretrained(UpperCamelCase )

							lowerCamelCase__   :      str			    =			ChineseCLIPModel(UpperCamelCase ).eval()

							lowerCamelCase__   :      Dict			    =			torch.load(UpperCamelCase							,     map_location="""cpu""" )["""state_dict"""]
							lowerCamelCase__   :      int			    =			{(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 )
							lowerCamelCase__   :      Tuple			    =			pt_weights["""logit_scale"""].data

							hf_model.save_pretrained(UpperCamelCase )


if __name__ == "__main__":
			_A      :     Dict       	=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.'''
			)
			_A      :     Dict       	=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'''







import unittest
from pathlib import Path

from transformers import is_vision_available, load_tool
from transformers.testing_utils import get_tests_dir

from .test_tools_common import ToolTesterMixin


if is_vision_available():
			from PIL import Image

class _lowercase      (							unittest.TestCase							, _lowercase						):




							def 						lowerCamelCase_							(				self:							Tuple	):
														lowerCamelCase__   :      Optional[Any]			    =			load_tool("""image-captioning"""	)
														self.tool.setup()
														lowerCamelCase__   :      Optional[Any]			    =			load_tool("""image-captioning"""     ,				remote=UpperCamelCase__	)




							def 						lowerCamelCase_							(				self:							int	):
														lowerCamelCase__   :      List[Any]			    =			Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""	)	) / """000000039769.png"""	)
														lowerCamelCase__   :      Optional[Any]			    =			self.tool(UpperCamelCase__	)
														self.assertEqual(UpperCamelCase__     ,				"""two cats sleeping on a couch"""	)




							def 						lowerCamelCase_							(				self:							Union[str, Any]	):
														lowerCamelCase__   :      Tuple			    =			Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""	)	) / """000000039769.png"""	)
														lowerCamelCase__   :      Dict			    =			self.remote_tool(UpperCamelCase__	)
														self.assertEqual(UpperCamelCase__     ,				"""two cats sleeping on a couch"""	)




							def 						lowerCamelCase_							(				self:							str	):
														lowerCamelCase__   :      Optional[int]			    =			Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""	)	) / """000000039769.png"""	)
														lowerCamelCase__   :      str			    =			self.tool(image=UpperCamelCase__	)
														self.assertEqual(UpperCamelCase__     ,				"""two cats sleeping on a couch"""	)




							def 						lowerCamelCase_							(				self:							Union[str, Any]	):
														lowerCamelCase__   :      Optional[int]			    =			Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""	)	) / """000000039769.png"""	)
														lowerCamelCase__   :      Optional[int]			    =			self.remote_tool(image=UpperCamelCase__	)
														self.assertEqual(UpperCamelCase__     ,				"""two cats sleeping on a couch"""	)

