from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase      		=					{
    '''configuration_openai''': ['''OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''OpenAIGPTConfig'''],
    '''tokenization_openai''': ['''OpenAIGPTTokenizer'''],
}

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

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCAmelCase      		=					[
     '''OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''OpenAIGPTDoubleHeadsModel''',
     '''OpenAIGPTForSequenceClassification''',
     '''OpenAIGPTLMHeadModel''',
     '''OpenAIGPTModel''',
     '''OpenAIGPTPreTrainedModel''',
     '''load_tf_weights_in_openai_gpt''',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCAmelCase      		=					[
     '''TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''TFOpenAIGPTDoubleHeadsModel''',
     '''TFOpenAIGPTForSequenceClassification''',
     '''TFOpenAIGPTLMHeadModel''',
     '''TFOpenAIGPTMainLayer''',
     '''TFOpenAIGPTModel''',
     '''TFOpenAIGPTPreTrainedModel''',
 ]


if TYPE_CHECKING:
 from .configuration_openai import OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenAIGPTConfig
 from .tokenization_openai import OpenAIGPTTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_openai_fast import OpenAIGPTTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_openai import (
      OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
      OpenAIGPTDoubleHeadsModel,
      OpenAIGPTForSequenceClassification,
      OpenAIGPTLMHeadModel,
      OpenAIGPTModel,
      OpenAIGPTPreTrainedModel,
      load_tf_weights_in_openai_gpt,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_openai import (
      TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFOpenAIGPTDoubleHeadsModel,
      TFOpenAIGPTForSequenceClassification,
      TFOpenAIGPTLMHeadModel,
      TFOpenAIGPTMainLayer,
      TFOpenAIGPTModel,
      TFOpenAIGPTPreTrainedModel,
  )

else:
 import sys

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




import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel


def SCREAMING_SNAKE_CASE					(  lowercase_:						List[str]		,    lowercase_:						Dict		,    lowercase_:						List[Any]					):
       _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase  :							List[Any]	    =	pt_weights[F"{prefix}.in_proj_weight"].chunk(3		,    dim=0					)
       _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase  :							Dict	    =	pt_weights[F"{prefix}.in_proj_bias"].chunk(3		,    dim=0					)

       _UpperCamelCase  :							List[str]	    =	pt_weights[F"{prefix}.out_proj.weight"]
       _UpperCamelCase  :							Any	    =	pt_weights[F"{prefix}.out_proj.bias"]

       _UpperCamelCase  :							List[Any]	    =	q_proj
       _UpperCamelCase  :							Union[str, Any]	    =	q_proj_bias

       _UpperCamelCase  :							str	    =	k_proj
       _UpperCamelCase  :							List[str]	    =	k_proj_bias

       _UpperCamelCase  :							Dict	    =	v_proj
       _UpperCamelCase  :							str	    =	v_proj_bias

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


def SCREAMING_SNAKE_CASE					(  lowercase_:						List[str]		,    lowercase_:						List[Any]		,    lowercase_:						Any					):
       copy_linear(hf_mlp.fca		,    lowercase_		,    F"{prefix}.c_fc"					)
       copy_linear(hf_mlp.fca		,    lowercase_		,    F"{prefix}.c_proj"					)


def SCREAMING_SNAKE_CASE					(  lowercase_:						Optional[Any]		,    lowercase_:						List[str]		,    lowercase_:						Any					):
       _UpperCamelCase  :							Dict	    =	pt_weights[F"{prefix}.weight"].data
       _UpperCamelCase  :							List[Any]	    =	pt_weights[F"{prefix}.bias"].data


def SCREAMING_SNAKE_CASE					(  lowercase_:						Any		,    lowercase_:						Optional[int]		,    lowercase_:						Union[str, Any]					):
       # copy layer norms
       copy_linear(hf_layer.layer_norma		,    lowercase_		,    F"{prefix}.ln_1"					)
       copy_linear(hf_layer.layer_norma		,    lowercase_		,    F"{prefix}.ln_2"					)

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

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


def SCREAMING_SNAKE_CASE					(  lowercase_:						str		,    lowercase_:						str		,    lowercase_:						int					):
       for layer_id, hf_layer in enumerate(lowercase_					):
              copy_layer(lowercase_		,    lowercase_		,    F"{prefix}.{layer_id}"					)


def SCREAMING_SNAKE_CASE					(  lowercase_:						Union[str, Any]		,    lowercase_:						List[Any]					):
       # copy projection
       _UpperCamelCase  :							List[str]	    =	pt_weights['text_projection'].data.T

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


def SCREAMING_SNAKE_CASE					(  lowercase_:						str		,    lowercase_:						Any					):
       # copy projection
       _UpperCamelCase  :							List[Any]	    =	pt_weights['visual.proj'].data.T

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

       # copy embeddings
       _UpperCamelCase  :							List[Any]	    =	pt_weights['visual.conv1.weight'].data
       _UpperCamelCase  :							List[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		,    lowercase_		,    'visual.transformer.resblocks'					)


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

       _UpperCamelCase  :							List[Any]	    =	ChineseCLIPModel(lowercase_					).eval()

       _UpperCamelCase  :							List[Any]	    =	torch.load(lowercase_		,    map_location='cpu'					)['state_dict']
       _UpperCamelCase  :							Union[str, Any]	    =	{(name[7:] if name.startswith('module.'					) else name): value for name, value in pt_weights.items()}

       copy_text_model_and_projection(lowercase_		,    lowercase_					)
       copy_vision_model_and_projection(lowercase_		,    lowercase_					)
       _UpperCamelCase  :							List[Any]	    =	pt_weights['logit_scale'].data

       hf_model.save_pretrained(lowercase_					)


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!''')



