


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        lowerCamelCase__   (					__lowerCamelCase    :  List[Any]	):
  __UpperCAmelCase							:			int     =						TimesformerConfig()

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

  if "hr" in model_name:
    __UpperCAmelCase							:			Any     =						16
    __UpperCAmelCase							:			Optional[Any]     =						448

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

  return config






def        lowerCamelCase__   (					__lowerCamelCase    :  Dict	):



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

  return name






def        lowerCamelCase__   (					__lowerCamelCase    :  List[Any]						,	__lowerCamelCase    :  Any	):



  for key in orig_state_dict.copy().keys():
    __UpperCAmelCase							:			int     =						orig_state_dict.pop(__lowerCamelCase	)

    if key.startswith("""model."""	):
      __UpperCAmelCase							:			Optional[int]     =						key.replace("""model."""						,	""""""	)

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

  return orig_state_dict






def        lowerCamelCase__   (					):
  __UpperCAmelCase							:			Dict     =						hf_hub_download(
      repo_id="""hf-internal-testing/spaghetti-video"""						,	filename="""eating_spaghetti.npy"""						,	repo_type="""dataset"""	)
  __UpperCAmelCase							:			str     =						np.load(__lowerCamelCase	)
  return list(__lowerCamelCase	)






def        lowerCamelCase__   (					__lowerCamelCase    :  Optional[Any]						,	__lowerCamelCase    :  Dict						,	__lowerCamelCase    :  Tuple						,	__lowerCamelCase    :  Union[str, Any]	):
  __UpperCAmelCase							:			int     =						get_timesformer_config(__lowerCamelCase	)

  __UpperCAmelCase							:			Dict     =						TimesformerForVideoClassification(__lowerCamelCase	)

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

  model.load_state_dict(__lowerCamelCase	)
  model.eval()

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

  __UpperCAmelCase							:			int     =						model(**__lowerCamelCase	)
  __UpperCAmelCase							:			Any     =						outputs.logits

  __UpperCAmelCase							:			List[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":
    __UpperCAmelCase							:			str     =						torch.Size([1, 400]	)
    __UpperCAmelCase							:			Optional[Any]     =						torch.tensor([-0.3_0_1_6, -0.7_7_1_3, -0.4_2_0_5]	)
  elif model_name == "timesformer-base-finetuned-k600":
    __UpperCAmelCase							:			str     =						torch.Size([1, 600]	)
    __UpperCAmelCase							:			str     =						torch.tensor([-0.7_2_6_7, -0.7_4_6_6, 3.2_4_0_4]	)
  elif model_name == "timesformer-base-finetuned-ssv2":
    __UpperCAmelCase							:			int     =						torch.Size([1, 174]	)
    __UpperCAmelCase							:			Optional[Any]     =						torch.tensor([-0.9_0_5_9, 0.6_4_3_3, -3.1_4_5_7]	)
  elif model_name == "timesformer-large-finetuned-k400":
    __UpperCAmelCase							:			List[Any]     =						torch.Size([1, 400]	)
    __UpperCAmelCase							:			Optional[int]     =						torch.tensor([0, 0, 0]	)
  elif model_name == "timesformer-large-finetuned-k600":
    __UpperCAmelCase							:			List[str]     =						torch.Size([1, 600]	)
    __UpperCAmelCase							:			Optional[Any]     =						torch.tensor([0, 0, 0]	)
  elif model_name == "timesformer-large-finetuned-ssv2":
    __UpperCAmelCase							:			Union[str, Any]     =						torch.Size([1, 174]	)
    __UpperCAmelCase							:			int     =						torch.tensor([0, 0, 0]	)
  elif model_name == "timesformer-hr-finetuned-k400":
    __UpperCAmelCase							:			int     =						torch.Size([1, 400]	)
    __UpperCAmelCase							:			str     =						torch.tensor([-0.9_6_1_7, -3.7_3_1_1, -3.7_7_0_8]	)
  elif model_name == "timesformer-hr-finetuned-k600":
    __UpperCAmelCase							:			Optional[Any]     =						torch.Size([1, 600]	)
    __UpperCAmelCase							:			Optional[Any]     =						torch.tensor([2.5_2_7_3, 0.7_1_2_7, 1.8_8_4_8]	)
  elif model_name == "timesformer-hr-finetuned-ssv2":
    __UpperCAmelCase							:			Union[str, Any]     =						torch.Size([1, 174]	)
    __UpperCAmelCase							:			str     =						torch.tensor([-3.6_7_5_6, -0.7_5_1_3, 0.7_1_8_0]	)
  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	: Optional[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."
       )

       a	: Union[str, Any]      =				parser.parse_args()
       convert_timesformer_checkpoint(
           args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
       )










import enum
import inspect
import warnings
from functools import wraps
from typing import Callable, Optional

from .logging import get_logger


a	: Any      =				set()
a	: List[Any]      =				get_logger(__name__)






def        lowerCamelCase__   (					__lowerCamelCase    :  Optional[str] = None	):

  def decorator(__lowerCamelCase    :  Callable	):
    global _emitted_deprecation_warnings

    if inspect.isclass(__lowerCamelCase	):
      __UpperCAmelCase							:			Optional[Any]     =						deprecated_class_or_function.__init__
      __UpperCAmelCase							:			int     =						deprecated_class_or_function.__name__
    else:
      __UpperCAmelCase							:			Any     =						deprecated_class_or_function
      __UpperCAmelCase							:			Any     =						deprecated_function.__name__
      # Support deprecating __init__ class method: class name instead
      __UpperCAmelCase							:			List[str]     =						name if name != """__init__""" else deprecated_function.__qualname__.split("""."""	)[-2]

    __UpperCAmelCase							:			Union[str, Any]     =						(
        f"""{name} is deprecated and will be removed in the next major version of datasets.""" + f""" {help_message}"""
        if help_message
        else """"""
    )

    @wraps(__lowerCamelCase	)
    def wrapper(*__lowerCamelCase    :  List[str]						,	**__lowerCamelCase    :  Tuple	):
      __UpperCAmelCase							:			Optional[Any]     =						hash(__lowerCamelCase	)
      if func_hash not in _emitted_deprecation_warnings:
        warnings.warn(__lowerCamelCase						,	category=__lowerCamelCase						,	stacklevel=2	)
        _emitted_deprecation_warnings.add(__lowerCamelCase	)
      return deprecated_function(*__lowerCamelCase						,	**__lowerCamelCase	)

    __UpperCAmelCase							:			List[str]     =						"""deprecated"""

    if inspect.isclass(__lowerCamelCase	):
      __UpperCAmelCase							:			int     =						wrapper
      return deprecated_class_or_function
    else:
      return wrapper

  return decorator
class 					a					(				enum.EnumMeta     ):



      """simple docstring"""


      def __getattribute__(   cls      :    Optional[int]	,   __lowercase      :    List[Any]      )							->   str:
        __UpperCAmelCase							:			List[Any]     =						super().__getattribute__(__lowercase      )
        if isinstance(__lowercase	,   enum.Enum      ) and obj._on_access:
          obj._on_access()
        return obj
      def __getitem__(   cls      :    Union[str, Any]	,   __lowercase      :    Optional[int]      )							->   Optional[int]:
        __UpperCAmelCase							:			Union[str, Any]     =						super().__getitem__(__lowercase      )
        if member._on_access:
          member._on_access()
        return member




      def __call__(   cls      :    List[str]	,   __lowercase      :    Dict	,   __lowercase      :    List[Any]=None	,   *, __lowercase      :    Dict=None	,   __lowercase      :    Any=None	,   __lowercase      :    Optional[Any]=None	,   __lowercase      :    Dict=1      )							->   List[Any]:
        __UpperCAmelCase							:			Union[str, Any]     =						super().__call__(__lowercase	,   __lowercase	,   module=__lowercase	,   qualname=__lowercase	,   type=__lowercase	,   start=__lowercase      )
        if isinstance(__lowercase	,   enum.Enum      ) and obj._on_access:
          obj._on_access()
        return obj


class 					a					(				enum.Enum							, metaclass=lowercase__     ):



      """simple docstring"""


      def __new__(   cls      :    Dict	,   __lowercase      :    Tuple      )							->   Optional[Any]:
        __UpperCAmelCase							:			List[Any]     =						object.__new__(cls      )
        __UpperCAmelCase							:			Any     =						value
        __UpperCAmelCase							:			int     =						member.deprecate
        return member
      @property
      def 			UpperCAmelCase   (   self      :    Optional[Any]      )							->   Dict:
        return ""




      def 			UpperCAmelCase   (   self      :    str      )							->   Dict:
        __UpperCAmelCase							:			Optional[int]     =						f""" {self.help_message}""" if self.help_message else """"""
        warnings.warn(
            f"""'{self.__objclass__.__name__}' is deprecated and will be removed in the next major version of datasets."""
            + help_message	,   __lowercase	,   stacklevel=3	,   )


