


'''simple docstring'''



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       UpperCAmelCase__ (     UpperCAmelCase_   :    Union[str, Any]							)					->		Tuple:
     __lowerCamelCase       :     List[Any]        =		TimesformerConfig()

     if "large" in model_name:
          __lowerCamelCase       :     Optional[int]        =		96

     if "hr" in model_name:
          __lowerCamelCase       :     Optional[int]        =		16
          __lowerCamelCase       :     str        =		4_48

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

     return config



def       UpperCAmelCase__ (     UpperCAmelCase_   :    Tuple							)					->		Any:
     if "encoder." in name:
          __lowerCamelCase       :     List[Any]        =		name.replace('encoder.'				,   ''							)
     if "cls_token" in name:
          __lowerCamelCase       :     Dict        =		name.replace('cls_token'				,   'timesformer.embeddings.cls_token'							)
     if "pos_embed" in name:
          __lowerCamelCase       :     List[Any]        =		name.replace('pos_embed'				,   'timesformer.embeddings.position_embeddings'							)
     if "time_embed" in name:
          __lowerCamelCase       :     Any        =		name.replace('time_embed'				,   'timesformer.embeddings.time_embeddings'							)
     if "patch_embed.proj" in name:
          __lowerCamelCase       :     Tuple        =		name.replace('patch_embed.proj'				,   'timesformer.embeddings.patch_embeddings.projection'							)
     if "patch_embed.norm" in name:
          __lowerCamelCase       :     Dict        =		name.replace('patch_embed.norm'				,   'timesformer.embeddings.norm'							)
     if "blocks" in name:
          __lowerCamelCase       :     Any        =		name.replace('blocks'				,   'timesformer.encoder.layer'							)
     if "attn.proj" in name:
          __lowerCamelCase       :     List[str]        =		name.replace('attn.proj'				,   'attention.output.dense'							)
     if "attn" in name and "bias" not in name and "temporal" not in name:
          __lowerCamelCase       :     Union[str, Any]        =		name.replace('attn'				,   'attention.self'							)
     if "attn" in name and "temporal" not in name:
          __lowerCamelCase       :     List[Any]        =		name.replace('attn'				,   'attention.attention'							)
     if "temporal_norm1" in name:
          __lowerCamelCase       :     str        =		name.replace('temporal_norm1'				,   'temporal_layernorm'							)
     if "temporal_attn.proj" in name:
          __lowerCamelCase       :     int        =		name.replace('temporal_attn'				,   'temporal_attention.output.dense'							)
     if "temporal_fc" in name:
          __lowerCamelCase       :     Dict        =		name.replace('temporal_fc'				,   'temporal_dense'							)
     if "norm1" in name and "temporal" not in name:
          __lowerCamelCase       :     Tuple        =		name.replace('norm1'				,   'layernorm_before'							)
     if "norm2" in name:
          __lowerCamelCase       :     int        =		name.replace('norm2'				,   'layernorm_after'							)
     if "mlp.fc1" in name:
          __lowerCamelCase       :     Tuple        =		name.replace('mlp.fc1'				,   'intermediate.dense'							)
     if "mlp.fc2" in name:
          __lowerCamelCase       :     Optional[Any]        =		name.replace('mlp.fc2'				,   'output.dense'							)
     if "norm.weight" in name and "fc" not in name and "temporal" not in name:
          __lowerCamelCase       :     str        =		name.replace('norm.weight'				,   'timesformer.layernorm.weight'							)
     if "norm.bias" in name and "fc" not in name and "temporal" not in name:
          __lowerCamelCase       :     List[Any]        =		name.replace('norm.bias'				,   'timesformer.layernorm.bias'							)
     if "head" in name:
          __lowerCamelCase       :     List[str]        =		name.replace('head'				,   'classifier'							)

     return name



def       UpperCAmelCase__ (     UpperCAmelCase_   :    List[str]				,   UpperCAmelCase_   :    List[str]							)					->		str:
     for key in orig_state_dict.copy().keys():
          __lowerCamelCase       :     Any        =		orig_state_dict.pop(UpperCAmelCase_							)

          if key.startswith('model.'							):
               __lowerCamelCase       :     Optional[int]        =		key.replace('model.'				,   ''							)

          if "qkv" in key:
               __lowerCamelCase       :     str        =		key.split('.'							)
               __lowerCamelCase       :     Dict        =		int(key_split[1]							)
               __lowerCamelCase       :     List[Any]        =		'timesformer.encoder.layer.'
               if "temporal" in key:
                    __lowerCamelCase       :     int        =		'.temporal_attention.attention.qkv.'
               else:
                    __lowerCamelCase       :     Tuple        =		'.attention.attention.qkv.'
               if "weight" in key:
                    __lowerCamelCase       :     Dict        =		val
               else:
                    __lowerCamelCase       :     Optional[Any]        =		val
          else:
               __lowerCamelCase       :     List[Any]        =		val

     return orig_state_dict



def       UpperCAmelCase__ (     )					->		Tuple:
     __lowerCamelCase       :     List[Any]        =		hf_hub_download(
         repo_id='hf-internal-testing/spaghetti-video'				,   filename='eating_spaghetti.npy'				,   repo_type='dataset'							)
     __lowerCamelCase       :     Any        =		np.load(UpperCAmelCase_							)
     return list(UpperCAmelCase_							)



def       UpperCAmelCase__ (     UpperCAmelCase_   :    Optional[int]				,   UpperCAmelCase_   :    Optional[int]				,   UpperCAmelCase_   :    Dict				,   UpperCAmelCase_   :    Any							)					->		Tuple:
     __lowerCamelCase       :     Optional[Any]        =		get_timesformer_config(UpperCAmelCase_							)

     __lowerCamelCase       :     Optional[Any]        =		TimesformerForVideoClassification(UpperCAmelCase_							)

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

     model.load_state_dict(UpperCAmelCase_							)
     model.eval()

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

     __lowerCamelCase       :     Tuple        =		model(**UpperCAmelCase_							)
     __lowerCamelCase       :     Optional[int]        =		outputs.logits

     __lowerCamelCase       :     str        =		[
         # 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":
          __lowerCamelCase       :     Tuple        =		torch.Size([1, 4_00]							)
          __lowerCamelCase       :     List[str]        =		torch.tensor([-0.3_016, -0.7_713, -0.4_205]							)
     elif model_name == "timesformer-base-finetuned-k600":
          __lowerCamelCase       :     List[str]        =		torch.Size([1, 6_00]							)
          __lowerCamelCase       :     List[str]        =		torch.tensor([-0.7_267, -0.7_466, 3.2_404]							)
     elif model_name == "timesformer-base-finetuned-ssv2":
          __lowerCamelCase       :     List[Any]        =		torch.Size([1, 1_74]							)
          __lowerCamelCase       :     Any        =		torch.tensor([-0.9_059, 0.6_433, -3.1_457]							)
     elif model_name == "timesformer-large-finetuned-k400":
          __lowerCamelCase       :     Optional[int]        =		torch.Size([1, 4_00]							)
          __lowerCamelCase       :     Union[str, Any]        =		torch.tensor([0, 0, 0]							)
     elif model_name == "timesformer-large-finetuned-k600":
          __lowerCamelCase       :     List[Any]        =		torch.Size([1, 6_00]							)
          __lowerCamelCase       :     List[str]        =		torch.tensor([0, 0, 0]							)
     elif model_name == "timesformer-large-finetuned-ssv2":
          __lowerCamelCase       :     List[Any]        =		torch.Size([1, 1_74]							)
          __lowerCamelCase       :     str        =		torch.tensor([0, 0, 0]							)
     elif model_name == "timesformer-hr-finetuned-k400":
          __lowerCamelCase       :     Dict        =		torch.Size([1, 4_00]							)
          __lowerCamelCase       :     Tuple        =		torch.tensor([-0.9_617, -3.7_311, -3.7_708]							)
     elif model_name == "timesformer-hr-finetuned-k600":
          __lowerCamelCase       :     str        =		torch.Size([1, 6_00]							)
          __lowerCamelCase       :     Optional[int]        =		torch.tensor([2.5_273, 0.7_127, 1.8_848]							)
     elif model_name == "timesformer-hr-finetuned-ssv2":
          __lowerCamelCase       :     Dict        =		torch.Size([1, 1_74]							)
          __lowerCamelCase       :     Any        =		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]				,   UpperCAmelCase_				,   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(UpperCAmelCase_							)
          model.save_pretrained(UpperCAmelCase_							)

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


if __name__ == "__main__":
      A__		: List[str]												=  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."""
      )

      A__		: int												=  parser.parse_args()
      convert_timesformer_checkpoint(
          args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
      )










'''simple docstring'''



import torch

from transformers import AutoTokenizer





class        UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self      ,					SCREAMING_SNAKE_CASE_     )    ->					int:
         __lowerCamelCase       :     Dict        =		AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE_     )







    def 				lowercase_    (							self      ,					SCREAMING_SNAKE_CASE_     )    ->					int:

         if isinstance(SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_     ) and all(isinstance(_x      ,					SCREAMING_SNAKE_CASE_     ) for _x in x     ):
              __lowerCamelCase       :     Union[str, Any]        =		None
              for l in x:
                   __lowerCamelCase       :     Tuple        =		self.tokenizer(
                       SCREAMING_SNAKE_CASE_      ,					padding='max_length'      ,					max_length=3_84      ,					truncation=SCREAMING_SNAKE_CASE_      ,					return_tensors='pt'      ,					)
                   __lowerCamelCase       :     str        =		torch.tensor([len(SCREAMING_SNAKE_CASE_     )]     )
                   if d is not None:
                        for k in d.keys():
                             __lowerCamelCase       :     List[Any]        =		torch.cat((d[k], t[k])      ,					0     )
                   else:
                        __lowerCamelCase       :     str        =		t

              __lowerCamelCase       :     List[Any]        =		torch.tensor(self.tokenizer.convert_tokens_to_ids('[E]'     )     )
              __lowerCamelCase       :     Tuple        =		torch.tensor(self.tokenizer.convert_tokens_to_ids('[/E]'     )     )

         elif isinstance(SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_     ) and all(isinstance(_x      ,					SCREAMING_SNAKE_CASE_     ) for _x in x     ):
              __lowerCamelCase       :     Optional[Any]        =		self.tokenizer(
                  SCREAMING_SNAKE_CASE_      ,					padding='max_length'      ,					max_length=3_84      ,					truncation=SCREAMING_SNAKE_CASE_      ,					return_tensors='pt'      ,					)

         else:
              raise Exception(
                  'Type of parameter x was not recognized! Only `list of strings` for query or `list of lists of'
                  ' strings` for supports are supported.'     )

         return d







    def 				lowercase_    (							self      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_=0.7_0     )    ->					str:
         __lowerCamelCase       :     str        =		[]
         for idx in range(len(W_query['input_ids']     )     ):
              __lowerCamelCase       :     Optional[Any]        =		range(p_start.shape[1]     )

              __lowerCamelCase       :     List[Any]        =		[]
              for start_id in start_indexes:
                   for end_id in end_indexes:
                        if start_id < end_id:
                             output.append(
                                 (
                                     start_id,
                                     end_id,
                                     p_start[idx][start_id].item(),
                                     p_end[idx][end_id].item(),
                                 )     )

              output.sort(key=lambda SCREAMING_SNAKE_CASE_     : (tup[2] * tup[3])      ,					reverse=SCREAMING_SNAKE_CASE_     )
              __lowerCamelCase       :     str        =		[]
              for k in range(len(SCREAMING_SNAKE_CASE_     )     ):
                   if output[k][2] * output[k][3] >= thresh:
                        __lowerCamelCase    ,	__lowerCamelCase       :     Union[str, Any]        =		output[k][0], output[k][1]
                        __lowerCamelCase       :     List[str]        =		self.tokenizer.decode(W_query['input_ids'][idx][c_start_pos:c_end_pos]     )
                        temp.append((decoded, output[k][2] * output[k][3])     )

              final_outputs.append(SCREAMING_SNAKE_CASE_     )

         return final_outputs






