

import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel


def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )      ->     List[Any]:






    '''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       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )      ->     Optional[Any]:






    '''simple docstring'''
    copy_linear(hf_mlp.fca       ,   SCREAMING_SNAKE_CASE       ,   f'''{prefix}.c_fc'''      )
    copy_linear(hf_mlp.fca       ,   SCREAMING_SNAKE_CASE       ,   f'''{prefix}.c_proj'''      )


def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )      ->     str:






    '''simple docstring'''
    __UpperCAmelCase									=			pt_weights[f'''{prefix}.weight'''].data
    __UpperCAmelCase									=			pt_weights[f'''{prefix}.bias'''].data


def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )      ->     int:






    '''simple docstring'''
    copy_linear(hf_layer.layer_norma       ,   SCREAMING_SNAKE_CASE       ,   f'''{prefix}.ln_1'''      )
    copy_linear(hf_layer.layer_norma       ,   SCREAMING_SNAKE_CASE       ,   f'''{prefix}.ln_2'''      )

    # copy MLP
    copy_mlp(hf_layer.mlp       ,   SCREAMING_SNAKE_CASE       ,   f'''{prefix}.mlp'''      )

    # copy attn
    copy_attn_layer(hf_layer.self_attn       ,   SCREAMING_SNAKE_CASE       ,   f'''{prefix}.attn'''      )


def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )      ->     str:






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


def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )      ->     int:






    '''simple docstring'''
    __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       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )      ->     str:






    '''simple docstring'''
    __UpperCAmelCase									=			pt_weights['''visual.proj'''].data.T

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


@torch.no_grad()
def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE=None      )      ->     List[Any]:






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

    __UpperCAmelCase									=			ChineseCLIPModel(SCREAMING_SNAKE_CASE      ).eval()

    __UpperCAmelCase									=			torch.load(SCREAMING_SNAKE_CASE       ,   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(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )
    copy_vision_model_and_projection(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )
    __UpperCAmelCase									=			pt_weights['''logit_scale'''].data

    hf_model.save_pretrained(SCREAMING_SNAKE_CASE      )


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

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







from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import (
           TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
           TFXLMForMultipleChoice,
           TFXLMForQuestionAnsweringSimple,
           TFXLMForSequenceClassification,
           TFXLMForTokenClassification,
           TFXLMModel,
           TFXLMWithLMHeadModel,
           XLMConfig,
       )
class 					A_   :


    '''simple docstring'''






    def __init__(self     ,      lowercase__     ,      )     ->      Tuple:
        __UpperCAmelCase									=			parent
        __UpperCAmelCase									=			13
        __UpperCAmelCase									=			7
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			False
        __UpperCAmelCase									=			False
        __UpperCAmelCase									=			False
        __UpperCAmelCase									=			2
        __UpperCAmelCase									=			99
        __UpperCAmelCase									=			0
        __UpperCAmelCase									=			32
        __UpperCAmelCase									=			2
        __UpperCAmelCase									=			4
        __UpperCAmelCase									=			0.1
        __UpperCAmelCase									=			0.1
        __UpperCAmelCase									=			512
        __UpperCAmelCase									=			16
        __UpperCAmelCase									=			2
        __UpperCAmelCase									=			0.02
        __UpperCAmelCase									=			3
        __UpperCAmelCase									=			4
        __UpperCAmelCase									=			'''last'''
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			None
        __UpperCAmelCase									=			0


    def      lowerCAmelCase_						(self  )     ->      Optional[Any]:
        __UpperCAmelCase									=			ids_tensor([self.batch_size, self.seq_length]     ,      self.vocab_size  )
        __UpperCAmelCase									=			random_attention_mask([self.batch_size, self.seq_length]     ,      dtype=tf.floataa  )

        __UpperCAmelCase									=			None
        if self.use_input_lengths:
            __UpperCAmelCase									=			(
                ids_tensor([self.batch_size]     ,      vocab_size=2  ) + self.seq_length - 2
            )  # small variation of seq_length

        __UpperCAmelCase									=			None
        if self.use_token_type_ids:
            __UpperCAmelCase									=			ids_tensor([self.batch_size, self.seq_length]     ,      self.n_langs  )

        __UpperCAmelCase									=			None
        __UpperCAmelCase									=			None
        __UpperCAmelCase									=			None
        if self.use_labels:
            __UpperCAmelCase									=			ids_tensor([self.batch_size]     ,      self.type_sequence_label_size  )
            __UpperCAmelCase									=			ids_tensor([self.batch_size, self.seq_length]     ,      self.num_labels  )
            __UpperCAmelCase									=			ids_tensor([self.batch_size]     ,      2     ,      dtype=tf.floataa  )
            __UpperCAmelCase									=			ids_tensor([self.batch_size]     ,      self.num_choices  )

        __UpperCAmelCase									=			XLMConfig(
            vocab_size=self.vocab_size     ,      n_special=self.n_special     ,      emb_dim=self.hidden_size     ,      n_layers=self.num_hidden_layers     ,      n_heads=self.num_attention_heads     ,      dropout=self.hidden_dropout_prob     ,      attention_dropout=self.attention_probs_dropout_prob     ,      gelu_activation=self.gelu_activation     ,      sinusoidal_embeddings=self.sinusoidal_embeddings     ,      asm=self.asm     ,      causal=self.causal     ,      n_langs=self.n_langs     ,      max_position_embeddings=self.max_position_embeddings     ,      initializer_range=self.initializer_range     ,      summary_type=self.summary_type     ,      use_proj=self.use_proj     ,      bos_token_id=self.bos_token_id     ,      )

        return (
            config,
            input_ids,
            token_type_ids,
            input_lengths,
            sequence_labels,
            token_labels,
            is_impossible_labels,
            choice_labels,
            input_mask,
        )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      List[Any]:
        __UpperCAmelCase									=			TFXLMModel(config=lowercase__  )
        __UpperCAmelCase									=			{'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
        __UpperCAmelCase									=			model(lowercase__  )

        __UpperCAmelCase									=			[input_ids, input_mask]
        __UpperCAmelCase									=			model(lowercase__  )
        self.parent.assertEqual(result.last_hidden_state.shape     ,      (self.batch_size, self.seq_length, self.hidden_size)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      List[str]:
        __UpperCAmelCase									=			TFXLMWithLMHeadModel(lowercase__  )

        __UpperCAmelCase									=			{'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
        __UpperCAmelCase									=			model(lowercase__  )

        __UpperCAmelCase									=			outputs

        self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.seq_length, self.vocab_size)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      Optional[Any]:
        __UpperCAmelCase									=			TFXLMForQuestionAnsweringSimple(lowercase__  )

        __UpperCAmelCase									=			{'''input_ids''': input_ids, '''lengths''': input_lengths}

        __UpperCAmelCase									=			model(lowercase__  )

        self.parent.assertEqual(result.start_logits.shape     ,      (self.batch_size, self.seq_length)  )
        self.parent.assertEqual(result.end_logits.shape     ,      (self.batch_size, self.seq_length)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      Dict:
        __UpperCAmelCase									=			TFXLMForSequenceClassification(lowercase__  )

        __UpperCAmelCase									=			{'''input_ids''': input_ids, '''lengths''': input_lengths}

        __UpperCAmelCase									=			model(lowercase__  )

        self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.type_sequence_label_size)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      int:
        __UpperCAmelCase									=			self.num_labels
        __UpperCAmelCase									=			TFXLMForTokenClassification(config=lowercase__  )
        __UpperCAmelCase									=			{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        __UpperCAmelCase									=			model(lowercase__  )
        self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.seq_length, self.num_labels)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      Tuple:
        __UpperCAmelCase									=			self.num_choices
        __UpperCAmelCase									=			TFXLMForMultipleChoice(config=lowercase__  )
        __UpperCAmelCase									=			tf.tile(tf.expand_dims(lowercase__     ,      1  )     ,      (1, self.num_choices, 1)  )
        __UpperCAmelCase									=			tf.tile(tf.expand_dims(lowercase__     ,      1  )     ,      (1, self.num_choices, 1)  )
        __UpperCAmelCase									=			tf.tile(tf.expand_dims(lowercase__     ,      1  )     ,      (1, self.num_choices, 1)  )
        __UpperCAmelCase									=			{
            '''input_ids''': multiple_choice_inputs_ids,
            '''attention_mask''': multiple_choice_input_mask,
            '''token_type_ids''': multiple_choice_token_type_ids,
        }
        __UpperCAmelCase									=			model(lowercase__  )
        self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.num_choices)  )






    def      lowerCAmelCase_						(self  )     ->      str:
        __UpperCAmelCase									=			self.prepare_config_and_inputs()
        (
            (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     
        )									=			config_and_inputs
        __UpperCAmelCase									=			{
            '''input_ids''': input_ids,
            '''token_type_ids''': token_type_ids,
            '''langs''': token_type_ids,
            '''lengths''': input_lengths,
        }
        return config, inputs_dict






@require_tf
class 					A_   (      _a ,     _a ,     unittest.TestCase						):


    '''simple docstring'''






    a__     							=       (
        (
            TFXLMModel,
            TFXLMWithLMHeadModel,
            TFXLMForSequenceClassification,
            TFXLMForQuestionAnsweringSimple,
            TFXLMForTokenClassification,
            TFXLMForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    a__     							=       (
        (TFXLMWithLMHeadModel,) if is_tf_available() else ()
    )  # TODO (PVP): Check other models whether language generation is also applicable
    a__     							=       (
        {
            "feature-extraction": TFXLMModel,
            "fill-mask": TFXLMWithLMHeadModel,
            "question-answering": TFXLMForQuestionAnsweringSimple,
            "text-classification": TFXLMForSequenceClassification,
            "text-generation": TFXLMWithLMHeadModel,
            "token-classification": TFXLMForTokenClassification,
            "zero-shot": TFXLMForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    a__     							=       False
    a__     							=       False


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__  )     ->      Tuple:
        if (
            pipeline_test_casse_name == "QAPipelineTests"
            and tokenizer_name is not None
            and not tokenizer_name.endswith('''Fast'''  )
        ):
            # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
            # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
            # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
            return True

        return False


    def      lowerCAmelCase_						(self  )     ->      Tuple:
        __UpperCAmelCase									=			TFXLMModelTester(self  )
        __UpperCAmelCase									=			ConfigTester(self     ,      config_class=lowercase__     ,      emb_dim=37  )


    def      lowerCAmelCase_						(self  )     ->      str:
        self.config_tester.run_common_tests()


    def      lowerCAmelCase_						(self  )     ->      Optional[int]:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_model(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      Any:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_lm_head(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      Dict:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_qa(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      int:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_sequence_classif(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      Optional[Any]:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_for_token_classification(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      Tuple:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_for_multiple_choice(*lowercase__  )






    @slow
    def      lowerCAmelCase_						(self  )     ->      int:
        for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __UpperCAmelCase									=			TFXLMModel.from_pretrained(lowercase__  )
            self.assertIsNotNone(lowercase__  )






@require_tf
class 					A_   (      unittest.TestCase						):


    '''simple docstring'''






    @slow
    def      lowerCAmelCase_						(self  )     ->      Tuple:
        __UpperCAmelCase									=			TFXLMWithLMHeadModel.from_pretrained('''xlm-mlm-en-2048'''  )
        __UpperCAmelCase									=			tf.convert_to_tensor([[14, 447]]     ,      dtype=tf.intaa  )  # the president
        __UpperCAmelCase									=			[
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
        ]  # the president the president the president the president the president the president the president the president the president the president
        # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
        __UpperCAmelCase									=			model.generate(lowercase__     ,      do_sample=lowercase__  )
        self.assertListEqual(output_ids[0].numpy().tolist()     ,      lowercase__  )




