





'''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__   (	lowerCAmelCase						):





      """simple docstring"""


      _lowerCAmelCase							  =							TimesformerConfig()

      if "large" in model_name:
            _lowerCAmelCase							  =							96

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

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

      return config
def      UpperCamelCase__   (	lowerCAmelCase						):





      """simple docstring"""


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

      return name
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase						):





      """simple docstring"""


      for key in orig_state_dict.copy().keys():
            _lowerCAmelCase							  =							orig_state_dict.pop(lowerCAmelCase						)

            if key.startswith("""model."""						):
                  _lowerCAmelCase							  =							key.replace("""model."""							,			""""""						)

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

      return orig_state_dict
def      UpperCamelCase__   (	):





      """simple docstring"""


      _lowerCAmelCase							  =							hf_hub_download(
          repo_id="""hf-internal-testing/spaghetti-video"""							,			filename="""eating_spaghetti.npy"""							,			repo_type="""dataset"""						)
      _lowerCAmelCase							  =							np.load(lowerCAmelCase						)
      return list(lowerCAmelCase						)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase						):





      """simple docstring"""


      _lowerCAmelCase							  =							get_timesformer_config(lowerCAmelCase						)

      _lowerCAmelCase							  =							TimesformerForVideoClassification(lowerCAmelCase						)

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

      model.load_state_dict(lowerCAmelCase						)
      model.eval()

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

      _lowerCAmelCase							  =							model(**lowerCAmelCase						)
      _lowerCAmelCase							  =							outputs.logits

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

from .custom_configuration import CustomConfig, NoSuperInitConfig




class 				UpperCAmelCase  (    snake_case_     ):
      _lowercase:    str		  =     CustomConfig




      def __init__(    self    :				Optional[int]    ,    __snake_case    :				Tuple						)       ->   Optional[Any]:
            super().__init__(__snake_case						)
            _lowerCAmelCase							  =							torch.nn.Linear(config.hidden_size    ,    config.hidden_size						)




      def        lowercase__					(    self    :				Dict    ,    __snake_case    :				List[Any]						)       ->   Dict:
            return self.linear(__snake_case						)




      def        lowercase__					(    self    :				List[Any]    ,    __snake_case    :				int						)       ->   Any:
            pass




class 				UpperCAmelCase  (    snake_case_     ):
      _lowercase:    List[str]		  =     NoSuperInitConfig




      def __init__(    self    :				List[str]    ,    __snake_case    :				List[str]						)       ->   int:
            super().__init__(__snake_case						)
            _lowerCAmelCase							  =							torch.nn.Linear(config.attribute    ,    config.attribute						)




      def        lowercase__					(    self    :				List[str]    ,    __snake_case    :				Dict						)       ->   Any:
            return self.linear(__snake_case						)




      def        lowercase__					(    self    :				Optional[int]    ,    __snake_case    :				List[Any]						)       ->   Dict:
            pass






