


import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel

def   UpperCamelCase__   (    UpperCAmelCase_    ,     UpperCAmelCase_    ,     UpperCAmelCase_			)						->							Optional[Any]:




       '''simple docstring'''






       _lowercase       ,	_lowercase       ,	_lowercase					:				Union[str, Any]           =			pt_weights[F'{prefix}.in_proj_weight'].chunk(3    ,     dim=0			)
       _lowercase       ,	_lowercase       ,	_lowercase					:				int           =			pt_weights[F'{prefix}.in_proj_bias'].chunk(3    ,     dim=0			)

       _lowercase					:				Union[str, Any]           =			pt_weights[F'{prefix}.out_proj.weight']
       _lowercase					:				List[Any]           =			pt_weights[F'{prefix}.out_proj.bias']

       _lowercase					:				Dict           =			q_proj
       _lowercase					:				Any           =			q_proj_bias

       _lowercase					:				str           =			k_proj
       _lowercase					:				Tuple           =			k_proj_bias

       _lowercase					:				List[Any]           =			v_proj
       _lowercase					:				Tuple           =			v_proj_bias

       _lowercase					:				List[Any]           =			out_proj_weights
       _lowercase					:				Optional[Any]           =			out_proj_bias

def   UpperCamelCase__   (    UpperCAmelCase_    ,     UpperCAmelCase_    ,     UpperCAmelCase_			)						->							List[str]:




       '''simple docstring'''






       copy_linear(hf_mlp.fca    ,     UpperCAmelCase_    ,     F'{prefix}.c_fc'			)
       copy_linear(hf_mlp.fca    ,     UpperCAmelCase_    ,     F'{prefix}.c_proj'			)

def   UpperCamelCase__   (    UpperCAmelCase_    ,     UpperCAmelCase_    ,     UpperCAmelCase_			)						->							int:




       '''simple docstring'''






       _lowercase					:				str           =			pt_weights[F'{prefix}.weight'].data
       _lowercase					:				Union[str, Any]           =			pt_weights[F'{prefix}.bias'].data

def   UpperCamelCase__   (    UpperCAmelCase_    ,     UpperCAmelCase_    ,     UpperCAmelCase_			)						->							Union[str, 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   UpperCamelCase__   (    UpperCAmelCase_    ,     UpperCAmelCase_    ,     UpperCAmelCase_			)						->							Optional[Any]:




       '''simple docstring'''






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

def   UpperCamelCase__   (    UpperCAmelCase_    ,     UpperCAmelCase_			)						->							int:




       '''simple docstring'''






       _lowercase					:				Any           =			pt_weights['''text_projection'''].data.T

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

def   UpperCamelCase__   (    UpperCAmelCase_    ,     UpperCAmelCase_			)						->							Union[str, Any]:




       '''simple docstring'''






       _lowercase					:				List[Any]           =			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
       _lowercase					:				Union[str, Any]           =			pt_weights['''visual.conv1.weight'''].data
       _lowercase					:				Dict           =			pt_weights['''visual.class_embedding'''].data
       _lowercase					:				Optional[int]           =			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__   (    UpperCAmelCase_    ,     UpperCAmelCase_    ,     UpperCAmelCase_=None			)						->							Union[str, Any]:




       '''simple docstring'''






       assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
       _lowercase					:				Dict           =			ChineseCLIPConfig.from_pretrained(UpperCAmelCase_			)

       _lowercase					:				List[Any]           =			ChineseCLIPModel(UpperCAmelCase_			).eval()

       _lowercase					:				Dict           =			torch.load(UpperCAmelCase_    ,     map_location='''cpu'''			)['''state_dict''']
       _lowercase					:				Any           =			{(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_			)
       _lowercase					:				Union[str, Any]           =			pt_weights['''logit_scale'''].data

       hf_model.save_pretrained(UpperCAmelCase_			)


if __name__ == "__main__":
   UpperCamelCase__				       =			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__				       =			parser.parse_args()

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



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


UpperCamelCase__				       =			{
    'configuration_realm': ['REALM_PRETRAINED_CONFIG_ARCHIVE_MAP', 'RealmConfig'],
    'tokenization_realm': ['RealmTokenizer'],
}

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCamelCase__				       =			['RealmTokenizerFast']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCamelCase__				       =			[
       'REALM_PRETRAINED_MODEL_ARCHIVE_LIST',
       'RealmEmbedder',
       'RealmForOpenQA',
       'RealmKnowledgeAugEncoder',
       'RealmPreTrainedModel',
       'RealmReader',
       'RealmScorer',
       'load_tf_weights_in_realm',
   ]
   UpperCamelCase__				       =			['RealmRetriever']


if TYPE_CHECKING:
   from .configuration_realm import REALM_PRETRAINED_CONFIG_ARCHIVE_MAP, RealmConfig
   from .tokenization_realm import RealmTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_realm import RealmTokenizerFast

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_realm import (
          REALM_PRETRAINED_MODEL_ARCHIVE_LIST,
          RealmEmbedder,
          RealmForOpenQA,
          RealmKnowledgeAugEncoder,
          RealmPreTrainedModel,
          RealmReader,
          RealmScorer,
          load_tf_weights_in_realm,
      )
      from .retrieval_realm import RealmRetriever


else:
   import sys

   UpperCamelCase__				       =			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)