






"""simple docstring"""







import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase			,				UpperCAmelCase			,				UpperCAmelCase					=	pt_weights[F'''{prefix}.in_proj_weight'''].chunk(3		,					dim=0						)
       UpperCAmelCase			,				UpperCAmelCase			,				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				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       copy_linear(hf_mlp.fca		,					lowerCAmelCase		,					F'''{prefix}.c_fc'''						)
       copy_linear(hf_mlp.fca		,					lowerCAmelCase		,					F'''{prefix}.c_proj'''						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	pt_weights[F'''{prefix}.weight'''].data
       UpperCAmelCase					=	pt_weights[F'''{prefix}.bias'''].data






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       # copy layer norms
       copy_linear(hf_layer.layer_norma		,					lowerCAmelCase		,					F'''{prefix}.ln_1'''						)
       copy_linear(hf_layer.layer_norma		,					lowerCAmelCase		,					F'''{prefix}.ln_2'''						)

       # copy MLP
       copy_mlp(hf_layer.mlp		,					lowerCAmelCase		,					F'''{prefix}.mlp'''						)

       # copy attn
       copy_attn_layer(hf_layer.self_attn		,					lowerCAmelCase		,					F'''{prefix}.attn'''						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







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






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       # copy projection
       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				(					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       # copy projection
       UpperCAmelCase					=	pt_weights["""visual.proj"""].data.T

       # copy layer norms
       copy_linear(hf_model.vision_model.pre_layrnorm		,					lowerCAmelCase		,					"""visual.ln_pre"""						)
       copy_linear(hf_model.vision_model.post_layernorm		,					lowerCAmelCase		,					"""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		,					lowerCAmelCase		,					"""visual.transformer.resblocks"""						)






@torch.no_grad()
def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase=None						):







       '''simple docstring'''
       assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
       UpperCAmelCase					=	ChineseCLIPConfig.from_pretrained(lowerCAmelCase						)

       UpperCAmelCase					=	ChineseCLIPModel(lowerCAmelCase						).eval()

       UpperCAmelCase					=	torch.load(lowerCAmelCase		,					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(lowerCAmelCase		,					lowerCAmelCase						)
       copy_vision_model_and_projection(lowerCAmelCase		,					lowerCAmelCase						)
       UpperCAmelCase					=	pt_weights["""logit_scale"""].data

       hf_model.save_pretrained(lowerCAmelCase						)


if __name__ == "__main__":
 lowerCAmelCase_		:			Optional[Any]				   =					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_		:			Optional[int]				   =					parser.parse_args()

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







import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import transformers
from transformers import (
    AutoConfig,
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    DataCollatorWithPadding,
    HfArgumentParser,
    SquadDataset,
    Trainer,
    TrainingArguments,
)
from transformers import SquadDataTrainingArguments as DataTrainingArguments
from transformers.trainer_utils import is_main_process


_snake_case				= logging.getLogger(__name__)





@dataclass
class 					lowerCAmelCase							:
  __lowerCamelCase		=			field(
      metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}					)
  __lowerCamelCase		=			field(
      default=lowercase_							,							metadata={'help': 'Pretrained config name or path if not the same as model_name'}					)
  __lowerCamelCase		=			field(
      default=lowercase_							,							metadata={'help': 'Pretrained tokenizer name or path if not the same as model_name'}					)
  __lowerCamelCase		=			field(default=lowercase_							,							metadata={'help': 'Set this flag to use fast tokenization.'}					)
  # If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,
  # or just modify its tokenizer_config.json.
  __lowerCamelCase		=			field(
      default=lowercase_							,							metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}							,							)


def _A   (    ):
   # See all possible arguments in src/transformers/training_args.py
   # or by passing the --help flag to this script.
   # We now keep distinct sets of args, for a cleaner separation of concerns.

   lowercase__    =  HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)       )

   if len(sys.argv       ) == 2 and sys.argv[1].endswith(".json"       ):
      # If we pass only one argument to the script and it's the path to a json file,
      # let's parse it to get our arguments.
      lowercase__   ,							lowercase__   ,							lowercase__    =  parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]       )       )
   else:
      lowercase__   ,							lowercase__   ,							lowercase__    =  parser.parse_args_into_dataclasses()

   if (
       os.path.exists(training_args.output_dir       )
       and os.listdir(training_args.output_dir       )
       and training_args.do_train
       and not training_args.overwrite_output_dir
   ):
      raise ValueError(
          f'''Output directory ({training_args.output_dir}) already exists and is not empty. Use'''
          " --overwrite_output_dir to overcome."       )

   # Setup logging
   logging.basicConfig(
       format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"  ,					datefmt="%m/%d/%Y %H:%M:%S"  ,					level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN  ,					)
   logger.warning(
       "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s"  ,					training_args.local_rank  ,					training_args.device  ,					training_args.n_gpu  ,					bool(training_args.local_rank != -1       )  ,					training_args.fpaa  ,					)
   # Set the verbosity to info of the Transformers logger (on main process only):
   if is_main_process(training_args.local_rank       ):
      transformers.utils.logging.set_verbosity_info()
      transformers.utils.logging.enable_default_handler()
      transformers.utils.logging.enable_explicit_format()
   logger.info("Training/evaluation parameters %s"  ,					__magic_name__       )

   # Prepare Question-Answering task
   # Load pretrained model and tokenizer
   #
   # Distributed training:
   # The .from_pretrained methods guarantee that only one local process can concurrently
   # download model & vocab.

   lowercase__    =  AutoConfig.from_pretrained(
       model_args.config_name if model_args.config_name else model_args.model_name_or_path  ,					cache_dir=model_args.cache_dir  ,					)
   lowercase__    =  AutoTokenizer.from_pretrained(
       model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path  ,					cache_dir=model_args.cache_dir  ,					use_fast=__magic_name__  ,					)
   lowercase__    =  AutoModelForQuestionAnswering.from_pretrained(
       model_args.model_name_or_path  ,					from_tf=bool(".ckpt" in model_args.model_name_or_path       )  ,					config=__magic_name__  ,					cache_dir=model_args.cache_dir  ,					)

   # Get datasets
   lowercase__    =  hasattr(model.config  ,					"lang2id"       )
   lowercase__    =  (
       SquadDataset(
           __magic_name__  ,					tokenizer=__magic_name__  ,					is_language_sensitive=__magic_name__  ,					cache_dir=model_args.cache_dir       )
       if training_args.do_train
       else None
   )
   lowercase__    =  (
       SquadDataset(
           __magic_name__  ,					tokenizer=__magic_name__  ,					mode="dev"  ,					is_language_sensitive=__magic_name__  ,					cache_dir=model_args.cache_dir  ,					)
       if training_args.do_eval
       else None
   )

   # Data collator
   lowercase__    =  DataCollatorWithPadding(__magic_name__  ,					pad_to_multiple_of=8       ) if training_args.fpaa else None

   # Initialize our Trainer
   lowercase__    =  Trainer(
       model=__magic_name__  ,					args=__magic_name__  ,					train_dataset=__magic_name__  ,					eval_dataset=__magic_name__  ,					data_collator=__magic_name__  ,					)

   # Training
   if training_args.do_train:
      trainer.train(
          model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path       ) else None       )
      trainer.save_model()
      # For convenience, we also re-save the tokenizer to the same directory,
      # so that you can share your model easily on huggingface.co/models =)
      if trainer.is_world_master():
         tokenizer.save_pretrained(training_args.output_dir       )


def _A   (    __magic_name__       ):
   # For xla_spawn (TPUs)
   main()


if __name__ == "__main__":
  main()






