

'''simple docstring'''
import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel


def  _lowerCAmelCase      (		_UpperCamelCase		:  Dict							,				_UpperCamelCase		:  List[str]							,				_UpperCamelCase		:  List[str]       )   ->    Optional[int]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE												=pt_weights[f"{prefix}.in_proj_weight"].chunk(3							,				dim=0       )
     _SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE					,		_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  _lowerCAmelCase      (		_UpperCamelCase		:  str							,				_UpperCamelCase		:  List[str]							,				_UpperCamelCase		:  Tuple       )   ->    Optional[int]:
     """simple docstring"""




     copy_linear(hf_mlp.fca							,				_UpperCamelCase							,				f"{prefix}.c_fc"       )
     copy_linear(hf_mlp.fca							,				_UpperCamelCase							,				f"{prefix}.c_proj"       )


def  _lowerCAmelCase      (		_UpperCamelCase		:  int							,				_UpperCamelCase		:  List[Any]							,				_UpperCamelCase		:  Dict       )   ->    Union[str, Any]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE												=pt_weights[f"{prefix}.weight"].data
     _SCREAMING_SNAKE_CASE												=pt_weights[f"{prefix}.bias"].data


def  _lowerCAmelCase      (		_UpperCamelCase		:  Any							,				_UpperCamelCase		:  List[Any]							,				_UpperCamelCase		:  str       )   ->    List[Any]:
     """simple docstring"""




     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      (		_UpperCamelCase		:  Optional[Any]							,				_UpperCamelCase		:  Any							,				_UpperCamelCase		:  Tuple       )   ->    Optional[Any]:
     """simple docstring"""




     for layer_id, hf_layer in enumerate(_UpperCamelCase       ):
          copy_layer(_UpperCamelCase							,				_UpperCamelCase							,				f"{prefix}.{layer_id}"       )


def  _lowerCAmelCase      (		_UpperCamelCase		:  List[str]							,				_UpperCamelCase		:  Union[str, Any]       )   ->    Optional[int]:
     """simple docstring"""




     _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  _lowerCAmelCase      (		_UpperCamelCase		:  int							,				_UpperCamelCase		:  Dict       )   ->    List[str]:
     """simple docstring"""




     _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  _lowerCAmelCase      (		_UpperCamelCase		:  Union[str, Any]							,				_UpperCamelCase		:  Tuple							,				_UpperCamelCase		:  Optional[Any]=None       )   ->    Any:
     """simple docstring"""




     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__":
     lowerCamelCase				:   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."
     )
     lowerCamelCase				:   List[str]      				=  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 TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCamelCase				:   str      				=  {"configuration_dinat": ["DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP", "DinatConfig"]}


try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCamelCase				:   Optional[int]      				=  [
         "DINAT_PRETRAINED_MODEL_ARCHIVE_LIST",
         "DinatForImageClassification",
         "DinatModel",
         "DinatPreTrainedModel",
         "DinatBackbone",
     ]

if TYPE_CHECKING:
     from .configuration_dinat import DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP, DinatConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_dinat import (
              DINAT_PRETRAINED_MODEL_ARCHIVE_LIST,
              DinatBackbone,
              DinatForImageClassification,
              DinatModel,
              DinatPreTrainedModel,
          )

else:
     import sys

     lowerCamelCase				:   Optional[Any]      				=  _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
