





import shutil
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SpeechaTextTokenizer, is_speech_available
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_torch, require_torchaudio
from transformers.utils import FEATURE_EXTRACTOR_NAME

from .test_feature_extraction_speech_to_text import floats_list


if is_speech_available():
  from transformers import SpeechaTextFeatureExtractor, SpeechaTextProcessor


_A							= get_tests_dir("""fixtures/test_sentencepiece.model""")



@require_torch
@require_torchaudio
@require_sentencepiece
class 				__snake_case  (     unittest.TestCase						):


    '''simple docstring'''







    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= tempfile.mkdtemp()

        SCREAMING_SNAKE_CASE__       							= ['''<s>''', '''<pad>''', '''</s>''', '''<unk>''', '''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']
        SCREAMING_SNAKE_CASE__       							= dict(zip(A_      ,							range(len(A_    )    )    )    )
        SCREAMING_SNAKE_CASE__       							= Path(self.tmpdirname    )
        save_json(A_      ,							save_dir / VOCAB_FILES_NAMES['''vocab_file''']    )
        if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
            copyfile(A_      ,							save_dir / VOCAB_FILES_NAMES['''spm_file''']    )

        SCREAMING_SNAKE_CASE__       							= SpeechaTextTokenizer.from_pretrained(self.tmpdirname    )
        tokenizer.save_pretrained(self.tmpdirname    )

        SCREAMING_SNAKE_CASE__       							= {
            '''feature_size''': 24,
            '''num_mel_bins''': 24,
            '''padding_value''': 0.0,
            '''sampling_rate''': 1_60_00,
            '''return_attention_mask''': False,
            '''do_normalize''': True,
        }
        save_json(A_      ,							save_dir / FEATURE_EXTRACTOR_NAME    )



    def    lowercase_				(					self      ,							**A_    ):





        '''simple docstring'''

        return SpeechaTextTokenizer.from_pretrained(self.tmpdirname      ,							**A_    )



    def    lowercase_				(					self      ,							**A_    ):





        '''simple docstring'''

        return SpeechaTextFeatureExtractor.from_pretrained(self.tmpdirname      ,							**A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        shutil.rmtree(self.tmpdirname    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()
        SCREAMING_SNAKE_CASE__       							= self.get_feature_extractor()

        SCREAMING_SNAKE_CASE__       							= SpeechaTextProcessor(tokenizer=A_      ,							feature_extractor=A_    )

        processor.save_pretrained(self.tmpdirname    )
        SCREAMING_SNAKE_CASE__       							= SpeechaTextProcessor.from_pretrained(self.tmpdirname    )

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

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



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= SpeechaTextProcessor(
            tokenizer=self.get_tokenizer()      ,							feature_extractor=self.get_feature_extractor()    )
        processor.save_pretrained(self.tmpdirname    )

        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer(bos_token='''(BOS)'''      ,							eos_token='''(EOS)'''    )
        SCREAMING_SNAKE_CASE__       							= self.get_feature_extractor(do_normalize=A_      ,							padding_value=1.0    )

        SCREAMING_SNAKE_CASE__       							= SpeechaTextProcessor.from_pretrained(
            self.tmpdirname      ,							bos_token='''(BOS)'''      ,							eos_token='''(EOS)'''      ,							do_normalize=A_      ,							padding_value=1.0    )

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

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



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_feature_extractor()
        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()

        SCREAMING_SNAKE_CASE__       							= SpeechaTextProcessor(tokenizer=A_      ,							feature_extractor=A_    )

        SCREAMING_SNAKE_CASE__       							= floats_list((3, 10_00)    )

        SCREAMING_SNAKE_CASE__       							= feature_extractor(A_      ,							return_tensors='''np'''    )
        SCREAMING_SNAKE_CASE__       							= processor(A_      ,							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    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_feature_extractor()
        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()

        SCREAMING_SNAKE_CASE__       							= SpeechaTextProcessor(tokenizer=A_      ,							feature_extractor=A_    )

        SCREAMING_SNAKE_CASE__       							= '''This is a test string'''

        SCREAMING_SNAKE_CASE__       							= processor(text=A_    )

        SCREAMING_SNAKE_CASE__       							= tokenizer(A_    )

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



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_feature_extractor()
        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()

        SCREAMING_SNAKE_CASE__       							= SpeechaTextProcessor(tokenizer=A_      ,							feature_extractor=A_    )

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

        SCREAMING_SNAKE_CASE__       							= processor.batch_decode(A_    )
        SCREAMING_SNAKE_CASE__       							= tokenizer.batch_decode(A_    )

        self.assertListEqual(A_      ,							A_    )




    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_feature_extractor()
        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()

        SCREAMING_SNAKE_CASE__       							= SpeechaTextProcessor(tokenizer=A_      ,							feature_extractor=A_    )

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







import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	str:
    SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= pt_weights[f'''{prefix}.in_proj_weight'''].chunk(3   ,       dim=0  )
    SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= pt_weights[f'''{prefix}.in_proj_bias'''].chunk(3   ,       dim=0  )

    SCREAMING_SNAKE_CASE__       							= pt_weights[f'''{prefix}.out_proj.weight''']
    SCREAMING_SNAKE_CASE__       							= pt_weights[f'''{prefix}.out_proj.bias''']

    SCREAMING_SNAKE_CASE__       							= q_proj
    SCREAMING_SNAKE_CASE__       							= q_proj_bias

    SCREAMING_SNAKE_CASE__       							= k_proj
    SCREAMING_SNAKE_CASE__       							= k_proj_bias

    SCREAMING_SNAKE_CASE__       							= v_proj
    SCREAMING_SNAKE_CASE__       							= v_proj_bias

    SCREAMING_SNAKE_CASE__       							= out_proj_weights
    SCREAMING_SNAKE_CASE__       							= out_proj_bias


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	Dict:
    copy_linear(hf_mlp.fca   ,       lowerCAmelCase_   ,       f'''{prefix}.c_fc'''  )
    copy_linear(hf_mlp.fca   ,       lowerCAmelCase_   ,       f'''{prefix}.c_proj'''  )


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	str:
    SCREAMING_SNAKE_CASE__       							= pt_weights[f'''{prefix}.weight'''].data
    SCREAMING_SNAKE_CASE__       							= pt_weights[f'''{prefix}.bias'''].data


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	Dict:
    # 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        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	Optional[int]:
    for layer_id, hf_layer in enumerate(lowerCAmelCase_  ):
        copy_layer(lowerCAmelCase_   ,       lowerCAmelCase_   ,       f'''{prefix}.{layer_id}'''  )


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_  )							->	Any:
    # copy projection
    SCREAMING_SNAKE_CASE__       							= pt_weights['''text_projection'''].data.T

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


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_  )							->	List[str]:
    # copy projection
    SCREAMING_SNAKE_CASE__       							= 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
    SCREAMING_SNAKE_CASE__       							= pt_weights['''visual.conv1.weight'''].data
    SCREAMING_SNAKE_CASE__       							= pt_weights['''visual.class_embedding'''].data
    SCREAMING_SNAKE_CASE__       							= pt_weights['''visual.positional_embedding'''].data

    # copy encoder
    copy_layers(hf_model.vision_model.encoder.layers   ,       lowerCAmelCase_   ,       '''visual.transformer.resblocks'''  )


@torch.no_grad()
def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_=None  )							->	str:
    assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
    SCREAMING_SNAKE_CASE__       							= ChineseCLIPConfig.from_pretrained(lowerCAmelCase_  )

    SCREAMING_SNAKE_CASE__       							= ChineseCLIPModel(lowerCAmelCase_  ).eval()

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

    hf_model.save_pretrained(lowerCAmelCase_  )


if __name__ == "__main__":
  _A				:							Union[str, 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."""
  )
  _A				:							Union[str, Any]							= parser.parse_args()

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