



import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel


def 						UpperCAmelCase__   (   UpperCAmelCase__			:str			,       UpperCAmelCase__			:List[str]			,       UpperCAmelCase__			:Dict       ):






 '''simple docstring'''

 a							=		pt_weights[F"""{prefix}.in_proj_weight"""].chunk(3			,       dim=0       )
 a							=		pt_weights[F"""{prefix}.in_proj_bias"""].chunk(3			,       dim=0       )

 a							=		pt_weights[F"""{prefix}.out_proj.weight"""]
 a							=		pt_weights[F"""{prefix}.out_proj.bias"""]

 a							=		q_proj
 a							=		q_proj_bias

 a							=		k_proj
 a							=		k_proj_bias

 a							=		v_proj
 a							=		v_proj_bias

 a							=		out_proj_weights
 a							=		out_proj_bias


def 						UpperCAmelCase__   (   UpperCAmelCase__			:List[str]			,       UpperCAmelCase__			:Optional[int]			,       UpperCAmelCase__			:Union[str, Any]       ):






 '''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__			:Optional[Any]			,       UpperCAmelCase__			:List[str]			,       UpperCAmelCase__			:str       ):






 '''simple docstring'''

 a							=		pt_weights[F"""{prefix}.weight"""].data
 a							=		pt_weights[F"""{prefix}.bias"""].data


def 						UpperCAmelCase__   (   UpperCAmelCase__			:Tuple			,       UpperCAmelCase__			:List[str]			,       UpperCAmelCase__			:List[str]       ):






 '''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__			:List[str]			,       UpperCAmelCase__			:List[str]			,       UpperCAmelCase__			:List[str]       ):






 '''simple docstring'''

 for layer_id, hf_layer in enumerate(_UpperCAmelCase       ):
  copy_layer(_UpperCAmelCase			,       _UpperCAmelCase			,       F"""{prefix}.{layer_id}"""       )


def 						UpperCAmelCase__   (   UpperCAmelCase__			:str			,       UpperCAmelCase__			:Any       ):






 '''simple docstring'''

 a							=		pt_weights['text_projection'].data.T

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


def 						UpperCAmelCase__   (   UpperCAmelCase__			:int			,       UpperCAmelCase__			:Tuple       ):






 '''simple docstring'''

 a							=		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
 a							=		pt_weights['visual.conv1.weight'].data
 a							=		pt_weights['visual.class_embedding'].data
 a							=		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__			:Optional[Any]			,       UpperCAmelCase__			:Dict			,       UpperCAmelCase__			:Dict=None       ):






 '''simple docstring'''

 assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
 a							=		ChineseCLIPConfig.from_pretrained(_UpperCAmelCase       )

 a							=		ChineseCLIPModel(_UpperCAmelCase       ).eval()

 a							=		torch.load(_UpperCAmelCase			,       map_location="cpu"       )['state_dict']
 a							=		{(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       )
 a							=		pt_weights['logit_scale'].data

 hf_model.save_pretrained(_UpperCAmelCase       )


if __name__ == "__main__":
     A_		:  Optional[int]										=  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_		:  Dict										=  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 multiprocessing
import os
from typing import BinaryIO, Optional, Union

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.csv.csv import Csv
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader

class 					_UpperCAmelCase    (  UpperCAmelCase__						):



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_ = None   , snake_case_ = None   , snake_case_ = None   , snake_case_ = False   , snake_case_ = False   , snake_case_ = None   , **snake_case_   , ):

        """simple docstring"""

        super().__init__(
            snake_case_   , split=snake_case_   , features=snake_case_   , cache_dir=snake_case_   , keep_in_memory=snake_case_   , streaming=snake_case_   , num_proc=snake_case_   , **snake_case_   , )
        A_    :			int									=    path_or_paths if isinstance(snake_case_   , snake_case_    ) else {self.split: path_or_paths}
        A_    :			List[str]									=    Csv(
            cache_dir=snake_case_   , data_files=snake_case_   , features=snake_case_   , **snake_case_   , )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        if self.streaming:
          A_    :			List[str]									=    self.builder.as_streaming_dataset(split=self.split    )
        # Build regular (map-style) dataset
        else:
          A_    :			List[Any]									=    None
          A_    :			Optional[int]									=    None
          A_    :			Optional[int]									=    None
          A_    :			List[str]									=    None

          self.builder.download_and_prepare(
              download_config=snake_case_   , download_mode=snake_case_   , verification_mode=snake_case_   , base_path=snake_case_   , num_proc=self.num_proc   , )
          A_    :			List[Any]									=    self.builder.as_dataset(
              split=self.split   , verification_mode=snake_case_   , in_memory=self.keep_in_memory    )
        return dataset


class 					_UpperCAmelCase    :



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_   , snake_case_ = None   , snake_case_ = None   , **snake_case_   , ):

        """simple docstring"""

        if num_proc is not None and num_proc <= 0:
          raise ValueError(F"""num_proc {num_proc} must be an integer > 0."""    )

        A_    :			Optional[int]									=    dataset
        A_    :			List[str]									=    path_or_buf
        A_    :			Optional[int]									=    batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
        A_    :			Optional[int]									=    num_proc
        A_    :			Optional[int]									=    'utf-8'
        A_    :			int									=    to_csv_kwargs

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Union[str, Any]									=    self.to_csv_kwargs.pop('path_or_buf'   , snake_case_    )
        A_    :			str									=    self.to_csv_kwargs.pop('header'   , snake_case_    )
        A_    :			Optional[Any]									=    self.to_csv_kwargs.pop('index'   , snake_case_    )

        if isinstance(self.path_or_buf   , (str, bytes, os.PathLike)    ):
          with open(self.path_or_buf   , 'wb+'    ) as buffer:
            A_    :			Optional[Any]									=    self._write(file_obj=snake_case_   , header=snake_case_   , index=snake_case_   , **self.to_csv_kwargs    )
        else:
          A_    :			int									=    self._write(file_obj=self.path_or_buf   , header=snake_case_   , index=snake_case_   , **self.to_csv_kwargs    )
        return written

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_		,			A_		,			A_		,			A_    :			Union[str, Any]									=    args

        A_    :			Dict									=    query_table(
            table=self.dataset.data   , key=slice(snake_case_   , offset + self.batch_size    )   , indices=self.dataset._indices   , )
        A_    :			List[str]									=    batch.to_pandas().to_csv(
            path_or_buf=snake_case_   , header=header if (offset == 0) else False   , index=snake_case_   , **snake_case_    )
        return csv_str.encode(self.encoding    )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_   , snake_case_   , **snake_case_    ):

        """simple docstring"""

        A_    :			Any									=    0

        if self.num_proc is None or self.num_proc == 1:
          for offset in logging.tqdm(
              range(0   , len(self.dataset    )   , self.batch_size    )   , unit='ba'   , disable=not logging.is_progress_bar_enabled()   , desc='Creating CSV from Arrow format'   , ):
            A_    :			Optional[int]									=    self._batch_csv((offset, header, index, to_csv_kwargs)    )
            written += file_obj.write(snake_case_    )

        else:
          A_		,			A_    :			Optional[int]									=    len(self.dataset    ), self.batch_size
          with multiprocessing.Pool(self.num_proc    ) as pool:
            for csv_str in logging.tqdm(
                pool.imap(
                    self._batch_csv   , [(offset, header, index, to_csv_kwargs) for offset in range(0   , snake_case_   , snake_case_    )]   , )   , total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size   , unit='ba'   , disable=not logging.is_progress_bar_enabled()   , desc='Creating CSV from Arrow format'   , ):
              written += file_obj.write(snake_case_    )

        return written