

'''simple docstring'''




import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel


def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   Tuple       ,   SCREAMING_SNAKE_CASE      :   Tuple       ,   SCREAMING_SNAKE_CASE      :   Optional[int]						):
  UpperCAmelCase													=	pt_weights[f'''{prefix}.in_proj_weight'''].chunk(3       ,   dim=0						)
  UpperCAmelCase													=	pt_weights[f'''{prefix}.in_proj_bias'''].chunk(3       ,   dim=0						)

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

  UpperCAmelCase													=	q_proj
  UpperCAmelCase													=	q_proj_bias

  UpperCAmelCase													=	k_proj
  UpperCAmelCase													=	k_proj_bias

  UpperCAmelCase													=	v_proj
  UpperCAmelCase													=	v_proj_bias

  UpperCAmelCase													=	out_proj_weights
  UpperCAmelCase													=	out_proj_bias


def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   Optional[Any]       ,   SCREAMING_SNAKE_CASE      :   List[Any]       ,   SCREAMING_SNAKE_CASE      :   List[str]						):
  copy_linear(hf_mlp.fca       ,   UpperCAmelCase_       ,   f'''{prefix}.c_fc'''						)
  copy_linear(hf_mlp.fca       ,   UpperCAmelCase_       ,   f'''{prefix}.c_proj'''						)


def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   Dict       ,   SCREAMING_SNAKE_CASE      :   Dict       ,   SCREAMING_SNAKE_CASE      :   List[str]						):
  UpperCAmelCase													=	pt_weights[f'''{prefix}.weight'''].data
  UpperCAmelCase													=	pt_weights[f'''{prefix}.bias'''].data


def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   Any       ,   SCREAMING_SNAKE_CASE      :   Optional[int]       ,   SCREAMING_SNAKE_CASE      :   Optional[Any]						):
  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       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   List[Any]       ,   SCREAMING_SNAKE_CASE      :   List[str]       ,   SCREAMING_SNAKE_CASE      :   List[Any]						):






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


def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   str       ,   SCREAMING_SNAKE_CASE      :   Dict						):
  UpperCAmelCase													=	pt_weights['''text_projection'''].data.T

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


def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   List[str]       ,   SCREAMING_SNAKE_CASE      :   Union[str, Any]						):
  UpperCAmelCase													=	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
  UpperCAmelCase													=	pt_weights['''visual.conv1.weight'''].data
  UpperCAmelCase													=	pt_weights['''visual.class_embedding'''].data
  UpperCAmelCase													=	pt_weights['''visual.positional_embedding'''].data

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


@torch.no_grad()
def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   Optional[Any]       ,   SCREAMING_SNAKE_CASE      :   int       ,   SCREAMING_SNAKE_CASE      :   int=None						):
  assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
  UpperCAmelCase													=	ChineseCLIPConfig.from_pretrained(UpperCAmelCase_						)

  UpperCAmelCase													=	ChineseCLIPModel(UpperCAmelCase_						).eval()

  UpperCAmelCase													=	torch.load(UpperCAmelCase_       ,   map_location='cpu'						)['''state_dict''']
  UpperCAmelCase													=	{(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_						)
  UpperCAmelCase													=	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				: Tuple 							=						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

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
  import torch

if is_vision_available():
  from PIL import Image

  from transformers import MobileNetVaImageProcessor



class        lowercase_	(		unittest.TestCase ):

  '''simple docstring'''





  def __init__(     self   , a_   , a_=7   , a_=3   , a_=1_8   , a_=3_0   , a_=4_0_0   , a_=True   , a_=None   , a_=True   , a_=None   , )							->				List[str]:


    """simple docstring"""
    UpperCAmelCase													=	size if size is not None else {'shortest_edge': 2_0}
    UpperCAmelCase													=	crop_size if crop_size is not None else {'height': 1_8, 'width': 1_8}
    UpperCAmelCase													=	parent
    UpperCAmelCase													=	batch_size
    UpperCAmelCase													=	num_channels
    UpperCAmelCase													=	image_size
    UpperCAmelCase													=	min_resolution
    UpperCAmelCase													=	max_resolution
    UpperCAmelCase													=	do_resize
    UpperCAmelCase													=	size
    UpperCAmelCase													=	do_center_crop
    UpperCAmelCase													=	crop_size





  def 						snake_case_			(     self     )							->				Optional[int]:


    """simple docstring"""
    return {
        "do_resize": self.do_resize,
        "size": self.size,
        "do_center_crop": self.do_center_crop,
        "crop_size": self.crop_size,
    }




@require_torch
@require_vision
class        lowercase_	(		a ,       unittest.TestCase ):

  '''simple docstring'''





  __lowerCAmelCase					: Union[str, Any]      =   MobileNetVaImageProcessor if is_vision_available() else None
  def 						snake_case_			(     self     )							->				Dict:


    """simple docstring"""
    UpperCAmelCase													=	MobileNetVaImageProcessingTester(self     )
  @property
  def 						snake_case_			(     self     )							->				List[str]:


    """simple docstring"""
    return self.image_processor_tester.prepare_image_processor_dict()
  def 						snake_case_			(     self     )							->				Optional[int]:


    """simple docstring"""
    UpperCAmelCase													=	self.image_processing_class(**self.image_processor_dict     )
    self.assertTrue(hasattr(a_   , 'do_resize'     )     )
    self.assertTrue(hasattr(a_   , 'size'     )     )
    self.assertTrue(hasattr(a_   , 'do_center_crop'     )     )
    self.assertTrue(hasattr(a_   , 'center_crop'     )     )
  def 						snake_case_			(     self     )							->				str:


    """simple docstring"""
    UpperCAmelCase													=	self.image_processing_class.from_dict(self.image_processor_dict     )
    self.assertEqual(image_processor.size   , {'shortest_edge': 2_0}     )
    self.assertEqual(image_processor.crop_size   , {'height': 1_8, 'width': 1_8}     )

    UpperCAmelCase													=	self.image_processing_class.from_dict(self.image_processor_dict   , size=4_2   , crop_size=8_4     )
    self.assertEqual(image_processor.size   , {'shortest_edge': 4_2}     )
    self.assertEqual(image_processor.crop_size   , {'height': 8_4, 'width': 8_4}     )
  def 						snake_case_			(     self     )							->				List[str]:


    """simple docstring"""
    pass
  def 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    # Initialize image_processing
    UpperCAmelCase													=	self.image_processing_class(**self.image_processor_dict     )
    # create random PIL images
    UpperCAmelCase													=	prepare_image_inputs(self.image_processor_tester   , equal_resolution=a_     )
    for image in image_inputs:
      self.assertIsInstance(a_   , Image.Image     )

    # Test not batched input
    UpperCAmelCase													=	image_processing(image_inputs[0]   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            1,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )

    # Test batched
    UpperCAmelCase													=	image_processing(a_   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )
  def 						snake_case_			(     self     )							->				str:


    """simple docstring"""
    # Initialize image_processing
    UpperCAmelCase													=	self.image_processing_class(**self.image_processor_dict     )
    # create random numpy tensors
    UpperCAmelCase													=	prepare_image_inputs(self.image_processor_tester   , equal_resolution=a_   , numpify=a_     )
    for image in image_inputs:
      self.assertIsInstance(a_   , np.ndarray     )

    # Test not batched input
    UpperCAmelCase													=	image_processing(image_inputs[0]   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            1,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )

    # Test batched
    UpperCAmelCase													=	image_processing(a_   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )





  def 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    # Initialize image_processing
    UpperCAmelCase													=	self.image_processing_class(**self.image_processor_dict     )
    # create random PyTorch tensors
    UpperCAmelCase													=	prepare_image_inputs(self.image_processor_tester   , equal_resolution=a_   , torchify=a_     )
    for image in image_inputs:
      self.assertIsInstance(a_   , torch.Tensor     )

    # Test not batched input
    UpperCAmelCase													=	image_processing(image_inputs[0]   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            1,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )

    # Test batched
    UpperCAmelCase													=	image_processing(a_   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )


