



import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel







def       lowerCamelCase_  (    _a       :						List[str]   ,		_a       :						Optional[int]   ,		_a       :						Optional[Any]    ):



    '''simple docstring'''




    UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_     :       List[Any]  		=					pt_weights[F'''{prefix}.in_proj_weight'''].chunk(3   ,		dim=0    )
    UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_     :       Optional[Any]  		=					pt_weights[F'''{prefix}.in_proj_bias'''].chunk(3   ,		dim=0    )

    UpperCAmelCase_     :       Tuple  		=					pt_weights[F'''{prefix}.out_proj.weight''']
    UpperCAmelCase_     :       Any  		=					pt_weights[F'''{prefix}.out_proj.bias''']

    UpperCAmelCase_     :       Optional[int]  		=					q_proj
    UpperCAmelCase_     :       Dict  		=					q_proj_bias

    UpperCAmelCase_     :       List[Any]  		=					k_proj
    UpperCAmelCase_     :       Optional[int]  		=					k_proj_bias

    UpperCAmelCase_     :       Union[str, Any]  		=					v_proj
    UpperCAmelCase_     :       Optional[Any]  		=					v_proj_bias

    UpperCAmelCase_     :       Union[str, Any]  		=					out_proj_weights
    UpperCAmelCase_     :       Optional[int]  		=					out_proj_bias







def       lowerCamelCase_  (    _a       :						Dict   ,		_a       :						Optional[int]   ,		_a       :						Tuple    ):



    '''simple docstring'''




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







def       lowerCamelCase_  (    _a       :						Optional[Any]   ,		_a       :						List[Any]   ,		_a       :						Optional[int]    ):



    '''simple docstring'''




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







def       lowerCamelCase_  (    _a       :						Union[str, Any]   ,		_a       :						Union[str, Any]   ,		_a       :						Optional[Any]    ):



    '''simple docstring'''




    copy_linear(hf_layer.layer_norma   ,		_a   ,		F'''{prefix}.ln_1'''    )
    copy_linear(hf_layer.layer_norma   ,		_a   ,		F'''{prefix}.ln_2'''    )

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

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







def       lowerCamelCase_  (    _a       :						Any   ,		_a       :						Union[str, Any]   ,		_a       :						Union[str, Any]    ):



    '''simple docstring'''




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







def       lowerCamelCase_  (    _a       :						Tuple   ,		_a       :						Optional[int]    ):



    '''simple docstring'''




    UpperCAmelCase_     :       List[Any]  		=					pt_weights["""text_projection"""].data.T

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







def       lowerCamelCase_  (    _a       :						int   ,		_a       :						Tuple    ):



    '''simple docstring'''




    UpperCAmelCase_     :       Tuple  		=					pt_weights["""visual.proj"""].data.T

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

    # copy embeddings
    UpperCAmelCase_     :       str  		=					pt_weights["""visual.conv1.weight"""].data
    UpperCAmelCase_     :       List[str]  		=					pt_weights["""visual.class_embedding"""].data
    UpperCAmelCase_     :       Dict  		=					pt_weights["""visual.positional_embedding"""].data

    # copy encoder
    copy_layers(hf_model.vision_model.encoder.layers   ,		_a   ,		"""visual.transformer.resblocks"""    )







@torch.no_grad()
def       lowerCamelCase_  (    _a       :						Any   ,		_a       :						Dict   ,		_a       :						Optional[int]=None    ):



    '''simple docstring'''




    assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
    UpperCAmelCase_     :       Optional[int]  		=					ChineseCLIPConfig.from_pretrained(_a    )

    UpperCAmelCase_     :       str  		=					ChineseCLIPModel(_a    ).eval()

    UpperCAmelCase_     :       Optional[Any]  		=					torch.load(_a   ,		map_location="""cpu"""    )["""state_dict"""]
    UpperCAmelCase_     :       List[Any]  		=					{(name[7:] if name.startswith("""module."""    ) else name): value for name, value in pt_weights.items()}

    copy_text_model_and_projection(_a   ,		_a    )
    copy_vision_model_and_projection(_a   ,		_a    )
    UpperCAmelCase_     :       Optional[Any]  		=					pt_weights["""logit_scale"""].data

    hf_model.save_pretrained(_a    )


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



import random
import shutil
import tempfile
import unittest

import numpy as np

from transformers import TaTokenizer, TaTokenizerFast
from transformers.testing_utils import require_sentencepiece, require_torch
from transformers.utils.import_utils import is_speech_available, is_torch_available


if is_torch_available():
    pass

if is_speech_available():
    from transformers import EncodecFeatureExtractor, MusicgenProcessor


lowerCamelCase_					  =   random.Random()



def   UpperCamelCase(   lowercase_     ,   lowercase_=1.0     ,   lowercase_=None     ,   lowercase_=None				)      ->							List[str]:





  '''simple docstring'''





  if rng is None:
    snake_case_					= global_rng

  snake_case_					= []
  for batch_idx in range(shape[0]				):
    values.append([]				)
    for _ in range(shape[1]				):
      values[-1].append(rng.random() * scale				)

  return values







@require_torch
@require_sentencepiece
class 					__lowerCamelCase							(      unittest.TestCase      ):







      def       lowerCAmelCase_       (	self  )		->			Union[str, Any]:
        snake_case_					= """facebook/musicgen-small"""
        snake_case_					= tempfile.mkdtemp()







      def       lowerCAmelCase_       (	self				,    **lowerCamelCase  )		->			int:
        return TaTokenizer.from_pretrained(self.checkpoint				,    **lowerCamelCase  )







      def       lowerCAmelCase_       (	self				,    **lowerCamelCase  )		->			Dict:
        return EncodecFeatureExtractor.from_pretrained(self.checkpoint				,    **lowerCamelCase  )







      def       lowerCAmelCase_       (	self  )		->			Tuple:
        shutil.rmtree(self.tmpdirname  )







      def       lowerCAmelCase_       (	self  )		->			Union[str, Any]:
        snake_case_					= self.get_tokenizer()
        snake_case_					= self.get_feature_extractor()

        snake_case_					= MusicgenProcessor(tokenizer=lowerCamelCase				,    feature_extractor=lowerCamelCase  )

        processor.save_pretrained(self.tmpdirname  )
        snake_case_					= MusicgenProcessor.from_pretrained(self.tmpdirname  )

        self.assertEqual(processor.tokenizer.get_vocab()				,    tokenizer.get_vocab()  )
        self.assertIsInstance(processor.tokenizer				,    lowerCamelCase  )

        self.assertEqual(processor.feature_extractor.to_json_string()				,    feature_extractor.to_json_string()  )
        self.assertIsInstance(processor.feature_extractor				,    lowerCamelCase  )







      def       lowerCAmelCase_       (	self  )		->			List[str]:
        snake_case_					= MusicgenProcessor(tokenizer=self.get_tokenizer()				,    feature_extractor=self.get_feature_extractor()  )
        processor.save_pretrained(self.tmpdirname  )

        snake_case_					= self.get_tokenizer(bos_token="""(BOS)"""				,    eos_token="""(EOS)"""  )
        snake_case_					= self.get_feature_extractor(do_normalize=lowerCamelCase				,    padding_value=1.0  )

        snake_case_					= MusicgenProcessor.from_pretrained(
            self.tmpdirname				,    bos_token="""(BOS)"""				,    eos_token="""(EOS)"""				,    do_normalize=lowerCamelCase				,    padding_value=1.0  )

        self.assertEqual(processor.tokenizer.get_vocab()				,    tokenizer_add_kwargs.get_vocab()  )
        self.assertIsInstance(processor.tokenizer				,    lowerCamelCase  )

        self.assertEqual(processor.feature_extractor.to_json_string()				,    feature_extractor_add_kwargs.to_json_string()  )
        self.assertIsInstance(processor.feature_extractor				,    lowerCamelCase  )







      def       lowerCAmelCase_       (	self  )		->			Tuple:
        snake_case_					= self.get_feature_extractor()
        snake_case_					= self.get_tokenizer()

        snake_case_					= MusicgenProcessor(tokenizer=lowerCamelCase				,    feature_extractor=lowerCamelCase  )

        snake_case_					= floats_list((3, 1000)  )

        snake_case_					= feature_extractor(lowerCamelCase				,    return_tensors="""np"""  )
        snake_case_					= processor(lowerCamelCase				,    return_tensors="""np"""  )

        for key in input_feat_extract.keys():
          self.assertAlmostEqual(input_feat_extract[key].sum()				,    input_processor[key].sum()				,    delta=1e-2  )







      def       lowerCAmelCase_       (	self  )		->			List[str]:
        snake_case_					= self.get_feature_extractor()
        snake_case_					= self.get_tokenizer()

        snake_case_					= MusicgenProcessor(tokenizer=lowerCamelCase				,    feature_extractor=lowerCamelCase  )

        snake_case_					= """This is a test string"""

        snake_case_					= processor(text=lowerCamelCase  )

        snake_case_					= tokenizer(lowerCamelCase  )

        for key in encoded_tok.keys():
          self.assertListEqual(encoded_tok[key]				,    encoded_processor[key]  )







      def       lowerCAmelCase_       (	self  )		->			int:
        snake_case_					= self.get_feature_extractor()
        snake_case_					= self.get_tokenizer()

        snake_case_					= MusicgenProcessor(tokenizer=lowerCamelCase				,    feature_extractor=lowerCamelCase  )

        snake_case_					= [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

        snake_case_					= processor.batch_decode(sequences=lowerCamelCase  )
        snake_case_					= tokenizer.batch_decode(lowerCamelCase  )

        self.assertListEqual(lowerCamelCase				,    lowerCamelCase  )







      def       lowerCAmelCase_       (	self  )		->			Optional[int]:
        snake_case_					= self.get_feature_extractor()
        snake_case_					= self.get_tokenizer()

        snake_case_					= MusicgenProcessor(tokenizer=lowerCamelCase				,    feature_extractor=lowerCamelCase  )

        self.assertListEqual(
            processor.model_input_names				,    feature_extractor.model_input_names				,    msg="""`processor` and `feature_extractor` model input names do not match"""				,    )







      def       lowerCAmelCase_       (	self  )		->			Optional[Any]:
        snake_case_					= self.get_feature_extractor(padding_side="""left"""  )
        snake_case_					= self.get_tokenizer()

        snake_case_					= MusicgenProcessor(tokenizer=lowerCamelCase				,    feature_extractor=lowerCamelCase  )

        snake_case_					= [floats_list((1, x)  )[0] for x in range(5				,    20				,    5  )]
        snake_case_					= processor(lowerCamelCase  ).padding_mask

        snake_case_					= np.asarray(floats_list((3, 20)  )  )[:, None, :]
        snake_case_					= processor.batch_decode(lowerCamelCase				,    padding_mask=lowerCamelCase  )

        self.assertIsInstance(lowerCamelCase				,    lowerCamelCase  )

        for audio in decoded_audios:
          self.assertIsInstance(lowerCamelCase				,    np.ndarray  )

        self.assertTrue(decoded_audios[0].shape == (1, 10)  )
        self.assertTrue(decoded_audios[1].shape == (1, 15)  )
        self.assertTrue(decoded_audios[2].shape == (1, 20)  )