





import unittest

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           CpmAntConfig,
           CpmAntForCausalLM,
           CpmAntModel,
           CpmAntTokenizer,
       )





@require_torch
class A_      :




   def __init__(self			:int					,     _UpperCamelCase			:Union[str, Any]					,     _UpperCamelCase			:str=2					,     _UpperCamelCase			:Union[str, Any]=8					,     _UpperCamelCase			:Optional[int]=True					,     _UpperCamelCase			:Dict=False					,     _UpperCamelCase			:Tuple=False					,     _UpperCamelCase			:Optional[int]=False					,     _UpperCamelCase			:Union[str, Any]=False					,     _UpperCamelCase			:int=99					,     _UpperCamelCase			:List[Any]=32					,     _UpperCamelCase			:int=3					,     _UpperCamelCase			:List[str]=4					,     _UpperCamelCase			:Union[str, Any]=37					,     _UpperCamelCase			:int=32					,     _UpperCamelCase			:str=128					,     _UpperCamelCase			:Any=8					,     _UpperCamelCase			:int=8					,     _UpperCamelCase			:Any=8					,     _UpperCamelCase			:Tuple=1.0					,     _UpperCamelCase			:int=True					,     )->      str:
         __A      =       parent
         __A      =       batch_size
         __A      =       seq_length
         __A      =       is_training
         __A      =       use_token_type_ids
         __A      =       use_input_mask
         __A      =       use_labels
         __A      =       use_mc_token_ids
         __A      =       vocab_size
         __A      =       hidden_size
         __A      =       num_hidden_layers
         __A      =       num_attention_heads
         __A      =       intermediate_size
         __A      =       num_buckets
         __A      =       max_distance
         __A      =       prompt_length
         __A      =       prompt_types
         __A      =       segment_types
         __A      =       init_std
         __A      =       return_dict




   def _lowerCAmelCase   (self			:List[str]				)->      Optional[int]:
         __A      =       {}
         __A      =       ids_tensor([self.batch_size, self.seq_length]					,     self.vocab_size				).type(torch.intaa				)
         __A      =       False

         __A      =       self.get_config()

         return (config, input_ids)




   def _lowerCAmelCase   (self			:Union[str, Any]				)->      Optional[int]:
         return CpmAntConfig(
             vocab_size=self.vocab_size					,     hidden_size=self.hidden_size					,     num_hidden_layers=self.num_hidden_layers					,     num_attention_heads=self.num_attention_heads					,     dim_ff=self.intermediate_size					,     position_bias_num_buckets=self.num_buckets					,     position_bias_max_distance=self.max_distance					,     prompt_types=self.prompt_types					,     prompt_length=self.prompt_length					,     segment_types=self.segment_types					,     use_cache=_UpperCamelCase					,     init_std=self.init_std					,     return_dict=self.return_dict					,     )




   def _lowerCAmelCase   (self			:Dict					,     _UpperCamelCase			:Optional[int]					,     _UpperCamelCase			:str					,     *_UpperCamelCase			:Optional[Any]				)->      Any:
         __A      =       CpmAntModel(config=_UpperCamelCase				)
         model.to(_UpperCamelCase				)
         model.eval()

         __A      =       model(**_UpperCamelCase				).last_hidden_state

         self.parent.assertEqual(hidden_states.shape					,     (self.batch_size, self.seq_length, config.hidden_size)				)




   def _lowerCAmelCase   (self			:str					,     _UpperCamelCase			:List[str]					,     _UpperCamelCase			:int					,     *_UpperCamelCase			:Optional[Any]				)->      Union[str, Any]:
         __A      =       CpmAntForCausalLM(_UpperCamelCase				)
         model.to(_UpperCamelCase				)
         __A      =       input_ids['''input_ids'''].to(_UpperCamelCase				)
         model.eval()

         __A      =       model(**_UpperCamelCase				)
         self.parent.assertEqual(
             model_output.logits.shape					,     (self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)					,     )






   def _lowerCAmelCase   (self			:str				)->      Optional[Any]:
         __A     ,		__A      =       self.prepare_config_and_inputs()
         return config, inputs_dict





@require_torch
class A_      (    UpperCAmelCase_				,   UpperCAmelCase_				,   unittest.TestCase     ):
   lowerCAmelCase__       =				(CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
   lowerCAmelCase__       =				(
       {"""feature-extraction""": CpmAntModel, """text-generation""": CpmAntForCausalLM} if is_torch_available() else {}
   )

   lowerCAmelCase__       =				False
   lowerCAmelCase__       =				False
   lowerCAmelCase__       =				False
   lowerCAmelCase__       =				False
   lowerCAmelCase__       =				False




   def _lowerCAmelCase   (self			:str				)->      List[Any]:
         __A      =       CpmAntModelTester(self				)
         __A      =       ConfigTester(self					,     config_class=_UpperCamelCase				)




   def _lowerCAmelCase   (self			:Any				)->      Tuple:
         self.config_tester.create_and_test_config_common_properties()
         self.config_tester.create_and_test_config_to_json_string()
         self.config_tester.create_and_test_config_to_json_file()
         self.config_tester.create_and_test_config_from_and_save_pretrained()
         self.config_tester.check_config_can_be_init_without_params()
         self.config_tester.check_config_arguments_init()




   def _lowerCAmelCase   (self			:Dict				)->      Optional[int]:
         unittest.skip('''CPMAnt doesn\'t support input_embeds.'''				)(self.test_inputs_embeds				)




   def _lowerCAmelCase   (self			:Optional[int]				)->      Tuple:
         unittest.skip(
             '''CPMAnt doesn\'t support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model.'''				)(self.test_retain_grad_hidden_states_attentions				)




   def _lowerCAmelCase   (self			:str				)->      Optional[int]:
         __A     ,		__A      =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_cpmant_model(_UpperCamelCase					,     _UpperCamelCase				)






   def _lowerCAmelCase   (self			:Dict				)->      Dict:
         __A     ,		__A      =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_lm_head_model(_UpperCamelCase					,     _UpperCamelCase				)





@require_torch
class A_      (    unittest.TestCase     ):




   @tooslow
   def _lowerCAmelCase   (self			:List[Any]				)->      Optional[Any]:
         __A      =       '''今天天气真好！'''
         __A      =       '''openbmb/cpm-ant-10b'''
         __A      =       CpmAntModel.from_pretrained(_UpperCamelCase				)
         __A      =       CpmAntTokenizer.from_pretrained(_UpperCamelCase				)
         __A      =       tokenizer(_UpperCamelCase					,     return_tensors='''pt'''				)
         __A      =       model(**_UpperCamelCase				).last_hidden_state

         __A      =       torch.tensor(
             [[[6.1_7_0_8, 5.9_2_4_4, 1.0_8_3_5], [6.5_2_0_7, 6.2_8_9_3, -1_1.3_3_2_4], [-1.0_1_0_7, -0.0_5_7_6, -5.9_5_7_7]]]					,     )
         self.assertTrue(torch.allclose(hidden_states[:, :3, :3]					,     _UpperCamelCase					,     atol=1e-2				)				)





@require_torch
class A_      (    unittest.TestCase     ):




   @tooslow
   def _lowerCAmelCase   (self			:Optional[Any]				)->      Any:
         __A      =       '''今天天气真好！'''
         __A      =       '''openbmb/cpm-ant-10b'''
         __A      =       CpmAntForCausalLM.from_pretrained(_UpperCamelCase				)
         __A      =       CpmAntTokenizer.from_pretrained(_UpperCamelCase				)
         __A      =       tokenizer(_UpperCamelCase					,     return_tensors='''pt'''				)
         __A      =       model(**_UpperCamelCase				).logits

         __A      =       torch.tensor(
             [[[-6.4_2_6_7, -6.4_0_8_3, -6.3_9_5_8], [-5.8_8_0_2, -5.9_4_4_7, -5.7_8_1_1], [-5.3_8_9_6, -5.4_8_2_0, -5.4_2_9_5]]]					,     )
         self.assertTrue(torch.allclose(hidden_states[:, :3, :3]					,     _UpperCamelCase					,     atol=1e-2				)				)




   @tooslow
   def _lowerCAmelCase   (self			:int				)->      int:
         __A      =       '''openbmb/cpm-ant-10b'''
         __A      =       CpmAntForCausalLM.from_pretrained(_UpperCamelCase				)
         __A      =       CpmAntTokenizer.from_pretrained(_UpperCamelCase				)
         __A      =       '''今天天气不错，'''
         __A      =       '''今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的'''
         __A      =       tokenizer(_UpperCamelCase					,     return_tensors='''pt'''				)
         __A      =       model.generate(**_UpperCamelCase				)
         __A      =       tokenizer.batch_decode(_UpperCamelCase				)
         self.assertEqual(_UpperCamelCase					,     _UpperCamelCase				)






   @tooslow
   def _lowerCAmelCase   (self			:List[str]				)->      Dict:
         __A      =       '''openbmb/cpm-ant-10b'''
         __A      =       CpmAntForCausalLM.from_pretrained(_UpperCamelCase				)
         __A      =       CpmAntTokenizer.from_pretrained(_UpperCamelCase				)
         __A      =       ['''今天天气不错，''', '''新年快乐，万事如意！''']
         __A      =       [
             '''今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的''',
             '''新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的''',
         ]
         __A      =       tokenizer(_UpperCamelCase					,     return_tensors='''pt'''					,     padding=_UpperCamelCase				)
         __A      =       model.generate(**_UpperCamelCase				)
         __A      =       tokenizer.batch_decode(_UpperCamelCase				)
         self.assertEqual(_UpperCamelCase					,     _UpperCamelCase				)









import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import TimesformerConfig, TimesformerForVideoClassification, VideoMAEImageProcessor



def _a       (			lowerCamelCase:     str				)			->		List[str]:







      '''simple docstring'''






      __A      =       TimesformerConfig()

      if "large" in model_name:
            __A      =       96

      if "hr" in model_name:
            __A      =       16
            __A      =       4_48

      __A      =       '''huggingface/label-files'''
      if "k400" in model_name:
            __A      =       4_00
            __A      =       '''kinetics400-id2label.json'''
      elif "k600" in model_name:
            __A      =       6_00
            __A      =       '''kinetics600-id2label.json'''
      elif "ssv2" in model_name:
            __A      =       1_74
            __A      =       '''something-something-v2-id2label.json'''
      else:
            raise ValueError('''Model name should either contain \'k400\', \'k600\' or \'ssv2\'.'''				)
      __A      =       json.load(open(hf_hub_download(lowerCamelCase      ,       lowerCamelCase      ,       repo_type='''dataset'''				)      ,       '''r'''				)				)
      __A      =       {int(lowerCamelCase				): v for k, v in idalabel.items()}
      __A      =       idalabel
      __A      =       {v: k for k, v in idalabel.items()}

      return config



def _a       (			lowerCamelCase:     Optional[Any]				)			->		List[str]:







      '''simple docstring'''






      if "encoder." in name:
            __A      =       name.replace('''encoder.'''      ,       ''''''				)
      if "cls_token" in name:
            __A      =       name.replace('''cls_token'''      ,       '''timesformer.embeddings.cls_token'''				)
      if "pos_embed" in name:
            __A      =       name.replace('''pos_embed'''      ,       '''timesformer.embeddings.position_embeddings'''				)
      if "time_embed" in name:
            __A      =       name.replace('''time_embed'''      ,       '''timesformer.embeddings.time_embeddings'''				)
      if "patch_embed.proj" in name:
            __A      =       name.replace('''patch_embed.proj'''      ,       '''timesformer.embeddings.patch_embeddings.projection'''				)
      if "patch_embed.norm" in name:
            __A      =       name.replace('''patch_embed.norm'''      ,       '''timesformer.embeddings.norm'''				)
      if "blocks" in name:
            __A      =       name.replace('''blocks'''      ,       '''timesformer.encoder.layer'''				)
      if "attn.proj" in name:
            __A      =       name.replace('''attn.proj'''      ,       '''attention.output.dense'''				)
      if "attn" in name and "bias" not in name and "temporal" not in name:
            __A      =       name.replace('''attn'''      ,       '''attention.self'''				)
      if "attn" in name and "temporal" not in name:
            __A      =       name.replace('''attn'''      ,       '''attention.attention'''				)
      if "temporal_norm1" in name:
            __A      =       name.replace('''temporal_norm1'''      ,       '''temporal_layernorm'''				)
      if "temporal_attn.proj" in name:
            __A      =       name.replace('''temporal_attn'''      ,       '''temporal_attention.output.dense'''				)
      if "temporal_fc" in name:
            __A      =       name.replace('''temporal_fc'''      ,       '''temporal_dense'''				)
      if "norm1" in name and "temporal" not in name:
            __A      =       name.replace('''norm1'''      ,       '''layernorm_before'''				)
      if "norm2" in name:
            __A      =       name.replace('''norm2'''      ,       '''layernorm_after'''				)
      if "mlp.fc1" in name:
            __A      =       name.replace('''mlp.fc1'''      ,       '''intermediate.dense'''				)
      if "mlp.fc2" in name:
            __A      =       name.replace('''mlp.fc2'''      ,       '''output.dense'''				)
      if "norm.weight" in name and "fc" not in name and "temporal" not in name:
            __A      =       name.replace('''norm.weight'''      ,       '''timesformer.layernorm.weight'''				)
      if "norm.bias" in name and "fc" not in name and "temporal" not in name:
            __A      =       name.replace('''norm.bias'''      ,       '''timesformer.layernorm.bias'''				)
      if "head" in name:
            __A      =       name.replace('''head'''      ,       '''classifier'''				)

      return name



def _a       (			lowerCamelCase:     str      ,       lowerCamelCase:     Tuple				)			->		Tuple:







      '''simple docstring'''






      for key in orig_state_dict.copy().keys():
            __A      =       orig_state_dict.pop(lowerCamelCase				)

            if key.startswith('''model.'''				):
                  __A      =       key.replace('''model.'''      ,       ''''''				)

            if "qkv" in key:
                  __A      =       key.split('''.'''				)
                  __A      =       int(key_split[1]				)
                  __A      =       '''timesformer.encoder.layer.'''
                  if "temporal" in key:
                        __A      =       '''.temporal_attention.attention.qkv.'''
                  else:
                        __A      =       '''.attention.attention.qkv.'''
                  if "weight" in key:
                        __A      =       val
                  else:
                        __A      =       val
            else:
                  __A      =       val

      return orig_state_dict



def _a       (			)			->		str:







      '''simple docstring'''






      __A      =       hf_hub_download(
          repo_id='''hf-internal-testing/spaghetti-video'''      ,       filename='''eating_spaghetti.npy'''      ,       repo_type='''dataset'''				)
      __A      =       np.load(lowerCamelCase				)
      return list(lowerCamelCase				)



def _a       (			lowerCamelCase:     List[str]      ,       lowerCamelCase:     Tuple      ,       lowerCamelCase:     Optional[Any]      ,       lowerCamelCase:     Tuple				)			->		List[Any]:







      '''simple docstring'''






      __A      =       get_timesformer_config(lowerCamelCase				)

      __A      =       TimesformerForVideoClassification(lowerCamelCase				)

      # download original checkpoint, hosted on Google Drive
      __A      =       '''pytorch_model.bin'''
      gdown.cached_download(lowerCamelCase      ,       lowerCamelCase      ,       quiet=lowerCamelCase				)
      __A      =       torch.load(lowerCamelCase      ,       map_location='''cpu'''				)
      if "model" in files:
            __A      =       files['''model''']
      elif "module" in files:
            __A      =       files['''module''']
      else:
            __A      =       files['''model_state''']
      __A      =       convert_state_dict(lowerCamelCase      ,       lowerCamelCase				)

      model.load_state_dict(lowerCamelCase				)
      model.eval()

      # verify model on basic input
      __A      =       VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]      ,       image_std=[0.5, 0.5, 0.5]				)
      __A      =       prepare_video()
      __A      =       image_processor(video[:8]      ,       return_tensors='''pt'''				)

      __A      =       model(**lowerCamelCase				)
      __A      =       outputs.logits

      __A      =       [
          # Kinetics-400 checkpoints (hr = high resolution input of 448px instead of 224px)
          '''timesformer-base-finetuned-k400''',
          '''timesformer-large-finetuned-k400''',
          '''timesformer-hr-finetuned-k400''',
          # Kinetics-600 checkpoints (hr = high resolution input of 448px instead of 224px)
          '''timesformer-base-finetuned-k600''',
          '''timesformer-large-finetuned-k600''',
          '''timesformer-hr-finetuned-k600''',
          # Something-Something-v2 checkpoints (hr = high resolution input of 448px instead of 224px)
          '''timesformer-base-finetuned-ssv2''',
          '''timesformer-large-finetuned-ssv2''',
          '''timesformer-hr-finetuned-ssv2''',
      ]

      # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
      if model_name == "timesformer-base-finetuned-k400":
            __A      =       torch.Size([1, 4_00]				)
            __A      =       torch.tensor([-0.3_016, -0.7_713, -0.4_205]				)
      elif model_name == "timesformer-base-finetuned-k600":
            __A      =       torch.Size([1, 6_00]				)
            __A      =       torch.tensor([-0.7_267, -0.7_466, 3.2_404]				)
      elif model_name == "timesformer-base-finetuned-ssv2":
            __A      =       torch.Size([1, 1_74]				)
            __A      =       torch.tensor([-0.9_059, 0.6_433, -3.1_457]				)
      elif model_name == "timesformer-large-finetuned-k400":
            __A      =       torch.Size([1, 4_00]				)
            __A      =       torch.tensor([0, 0, 0]				)
      elif model_name == "timesformer-large-finetuned-k600":
            __A      =       torch.Size([1, 6_00]				)
            __A      =       torch.tensor([0, 0, 0]				)
      elif model_name == "timesformer-large-finetuned-ssv2":
            __A      =       torch.Size([1, 1_74]				)
            __A      =       torch.tensor([0, 0, 0]				)
      elif model_name == "timesformer-hr-finetuned-k400":
            __A      =       torch.Size([1, 4_00]				)
            __A      =       torch.tensor([-0.9_617, -3.7_311, -3.7_708]				)
      elif model_name == "timesformer-hr-finetuned-k600":
            __A      =       torch.Size([1, 6_00]				)
            __A      =       torch.tensor([2.5_273, 0.7_127, 1.8_848]				)
      elif model_name == "timesformer-hr-finetuned-ssv2":
            __A      =       torch.Size([1, 1_74]				)
            __A      =       torch.tensor([-3.6_756, -0.7_513, 0.7_180]				)
      else:
            raise ValueError(F"""Model name not supported. Should be one of {model_names}"""				)

      # verify logits
      assert logits.shape == expected_shape
      assert torch.allclose(logits[0, :3]      ,       lowerCamelCase      ,       atol=1e-4				)
      print('''Logits ok!'''				)

      if pytorch_dump_folder_path is not None:
            print(F"""Saving model and image processor to {pytorch_dump_folder_path}"""				)
            image_processor.save_pretrained(lowerCamelCase				)
            model.save_pretrained(lowerCamelCase				)

      if push_to_hub:
            print('''Pushing to the hub...'''				)
            model.push_to_hub(F"""fcakyon/{model_name}"""				)


if __name__ == "__main__":
       snake_case__      :							int									=  argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           '--checkpoint_url',
           default='https://drive.google.com/u/1/uc?id=17yvuYp9L4mn-HpIcK5Zo6K3UoOy1kA5l&export=download',
           type=str,
           help=(
               'URL of the original PyTorch checkpoint (on Google Drive) you\'d like to convert. Should be a direct'
               ' download link.'
           ),
       )
       parser.add_argument(
           '--pytorch_dump_folder_path',
           default='',
           type=str,
           help='Path to the output PyTorch model directory.',
       )
       parser.add_argument('--model_name', default='timesformer-base-finetuned-k400', type=str, help='Name of the model.')
       parser.add_argument(
           '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
       )

       snake_case__      :							str									=  parser.parse_args()
       convert_timesformer_checkpoint(
           args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
       )





