
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_		(			SCREAMING_SNAKE_CASE_  )       ->       Dict:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE       	=			TimesformerConfig()

   if "large" in model_name:
      _SCREAMING_SNAKE_CASE       	=			96

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

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

   return config
def 				lowerCAmelCase_		(			SCREAMING_SNAKE_CASE_  )       ->       Union[str, Any]:




   """simple docstring"""







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

   return name
def 				lowerCAmelCase_		(			SCREAMING_SNAKE_CASE_      , SCREAMING_SNAKE_CASE_  )       ->       Union[str, Any]:




   """simple docstring"""







   for key in orig_state_dict.copy().keys():
      _SCREAMING_SNAKE_CASE       	=			orig_state_dict.pop(SCREAMING_SNAKE_CASE_  )

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

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

   return orig_state_dict
def 				lowerCAmelCase_		(			)       ->       str:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE       	=			hf_hub_download(
       repo_id="""hf-internal-testing/spaghetti-video"""      , filename="""eating_spaghetti.npy"""      , repo_type="""dataset"""  )
   _SCREAMING_SNAKE_CASE       	=			np.load(SCREAMING_SNAKE_CASE_  )
   return list(SCREAMING_SNAKE_CASE_  )
def 				lowerCAmelCase_		(			SCREAMING_SNAKE_CASE_      , SCREAMING_SNAKE_CASE_      , SCREAMING_SNAKE_CASE_      , SCREAMING_SNAKE_CASE_  )       ->       int:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE       	=			get_timesformer_config(SCREAMING_SNAKE_CASE_  )

   _SCREAMING_SNAKE_CASE       	=			TimesformerForVideoClassification(SCREAMING_SNAKE_CASE_  )

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

   model.load_state_dict(SCREAMING_SNAKE_CASE_  )
   model.eval()

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

   _SCREAMING_SNAKE_CASE       	=			model(**SCREAMING_SNAKE_CASE_  )
   _SCREAMING_SNAKE_CASE       	=			outputs.logits

   _SCREAMING_SNAKE_CASE       	=			[
       # 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":
      _SCREAMING_SNAKE_CASE       	=			torch.Size([1, 4_00]  )
      _SCREAMING_SNAKE_CASE       	=			torch.tensor([-0.3016, -0.7713, -0.4205]  )
   elif model_name == "timesformer-base-finetuned-k600":
      _SCREAMING_SNAKE_CASE       	=			torch.Size([1, 6_00]  )
      _SCREAMING_SNAKE_CASE       	=			torch.tensor([-0.7267, -0.7466, 3.2404]  )
   elif model_name == "timesformer-base-finetuned-ssv2":
      _SCREAMING_SNAKE_CASE       	=			torch.Size([1, 1_74]  )
      _SCREAMING_SNAKE_CASE       	=			torch.tensor([-0.9059, 0.6433, -3.1457]  )
   elif model_name == "timesformer-large-finetuned-k400":
      _SCREAMING_SNAKE_CASE       	=			torch.Size([1, 4_00]  )
      _SCREAMING_SNAKE_CASE       	=			torch.tensor([0, 0, 0]  )
   elif model_name == "timesformer-large-finetuned-k600":
      _SCREAMING_SNAKE_CASE       	=			torch.Size([1, 6_00]  )
      _SCREAMING_SNAKE_CASE       	=			torch.tensor([0, 0, 0]  )
   elif model_name == "timesformer-large-finetuned-ssv2":
      _SCREAMING_SNAKE_CASE       	=			torch.Size([1, 1_74]  )
      _SCREAMING_SNAKE_CASE       	=			torch.tensor([0, 0, 0]  )
   elif model_name == "timesformer-hr-finetuned-k400":
      _SCREAMING_SNAKE_CASE       	=			torch.Size([1, 4_00]  )
      _SCREAMING_SNAKE_CASE       	=			torch.tensor([-0.9617, -3.7311, -3.7708]  )
   elif model_name == "timesformer-hr-finetuned-k600":
      _SCREAMING_SNAKE_CASE       	=			torch.Size([1, 6_00]  )
      _SCREAMING_SNAKE_CASE       	=			torch.tensor([2.5273, 0.7127, 1.8848]  )
   elif model_name == "timesformer-hr-finetuned-ssv2":
      _SCREAMING_SNAKE_CASE       	=			torch.Size([1, 1_74]  )
      _SCREAMING_SNAKE_CASE       	=			torch.tensor([-3.6756, -0.7513, 0.7180]  )
   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]      , SCREAMING_SNAKE_CASE_      , 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(SCREAMING_SNAKE_CASE_  )
      model.save_pretrained(SCREAMING_SNAKE_CASE_  )

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


if __name__ == "__main__":
     UpperCamelCase__ :						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."
     )

     UpperCamelCase__ :						Tuple    							=					parser.parse_args()
     convert_timesformer_checkpoint(
         args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
     )








'''simple docstring'''



import argparse
import time

import torch
import transformers
from measures_util import end_measure, log_measures, start_measure
from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeqaSeqLM, AutoTokenizer

from accelerate.utils import compute_module_sizes


UpperCamelCase__ :						Tuple    							=					{
    "gpt-j-6b": {"is_causal": True, "model": "sgugger/sharded-gpt-j-6B", "tokenizer": "EleutherAI/gpt-j-6B"},
    "gpt-neox": {"is_causal": True, "model": "EleutherAI/gpt-neox-20b"},
    "opt": {"is_causal": True, "model": "facebook/opt-30b"},
    "T0pp": {"is_causal": False, "model": "bigscience/T0pp", "model_revision": "sharded"},
}

UpperCamelCase__ :						int    							=					[
    "Hello, my name is",
    "Are unicorns real? Unicorns are",
    "For the first time in several years,",
    "My name is Julien and I am",
    "The goal of life is",
    "Whenever I'm sad, I like to",
]
def 				lowerCAmelCase_		(			)       ->       str:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE       	=			argparse.ArgumentParser(description="""Run and time generations on a big model using Accelerate."""  )
   parser.add_argument("""model_name"""      , type=SCREAMING_SNAKE_CASE_      , default=SCREAMING_SNAKE_CASE_      , help="""The name of the model to try."""  )
   parser.add_argument(
       """--tokenizer_name"""      , type=SCREAMING_SNAKE_CASE_      , default=SCREAMING_SNAKE_CASE_      , help="""The name of the tokenizer (if different from the model."""  )
   parser.add_argument("""--is_causal"""      , type=SCREAMING_SNAKE_CASE_      , default=SCREAMING_SNAKE_CASE_      , help="""Whether or not the model is causal."""  )
   parser.add_argument(
       """--model_revision"""      , type=SCREAMING_SNAKE_CASE_      , default=SCREAMING_SNAKE_CASE_      , help="""The revision to use for the model checkpoint."""  )
   parser.add_argument("""--torch_dtype"""      , type=SCREAMING_SNAKE_CASE_      , default=SCREAMING_SNAKE_CASE_      , help="""The dtype for the model."""  )
   parser.add_argument("""--disk_offload"""      , action="""store_true"""  )

   _SCREAMING_SNAKE_CASE       	=			parser.parse_args()

   # Sanitize args
   if args.model_name in DEFAULT_MODELS:
      _SCREAMING_SNAKE_CASE       	=			DEFAULT_MODELS[args.model_name]
      _SCREAMING_SNAKE_CASE       	=			defaults["""model"""]
      if args.tokenizer_name is None:
         _SCREAMING_SNAKE_CASE       	=			defaults.get("""tokenizer"""      , args.model_name  )
      if args.is_causal is None:
         _SCREAMING_SNAKE_CASE       	=			defaults["""is_causal"""]
      if args.model_revision is None:
         _SCREAMING_SNAKE_CASE       	=			defaults.get("""model_revision"""      , """main"""  )

   if args.is_causal is None:
      raise ValueError("""Could not infer the default for `--is_causal`, pass either True or False for it."""  )
   if args.tokenizer_name is None:
      _SCREAMING_SNAKE_CASE       	=			args.model_name
   if args.model_revision is None:
      _SCREAMING_SNAKE_CASE       	=			"""main"""

   return args
def 				lowerCAmelCase_		(			)       ->       Optional[int]:




   """simple docstring"""
   transformers.utils.logging.set_verbosity_error()
   _SCREAMING_SNAKE_CASE       	=			parse_args()

   if args.torch_dtype is None:
      _SCREAMING_SNAKE_CASE       	=			AutoConfig.from_pretrained(args.model_name  )
      _SCREAMING_SNAKE_CASE       	=			getattr(SCREAMING_SNAKE_CASE_      , """torch_dtype"""      , torch.floataa  )
   else:
      _SCREAMING_SNAKE_CASE       	=			getattr(SCREAMING_SNAKE_CASE_      , args.torch_dtype  )
   _SCREAMING_SNAKE_CASE       	=			AutoModelForCausalLM if args.is_causal else AutoModelForSeqaSeqLM
   _SCREAMING_SNAKE_CASE       	=			{
       """torch_dtype""": torch_dtype,
       """revision""": args.model_revision,
   }
   if args.disk_offload:
      _SCREAMING_SNAKE_CASE       	=			"""tmp_offload"""
      _SCREAMING_SNAKE_CASE       	=			True

   _SCREAMING_SNAKE_CASE       	=			start_measure()
   _SCREAMING_SNAKE_CASE       	=			model_cls.from_pretrained(args.model_name      , device_map="""auto"""      , **SCREAMING_SNAKE_CASE_  )
   _SCREAMING_SNAKE_CASE       	=			end_measure(SCREAMING_SNAKE_CASE_  )
   log_measures(SCREAMING_SNAKE_CASE_      , """Model loading"""  )

   _SCREAMING_SNAKE_CASE       	=			compute_module_sizes(SCREAMING_SNAKE_CASE_  )
   _SCREAMING_SNAKE_CASE       	=			{v: 0 for v in model.hf_device_map.values()}
   for module, device in model.hf_device_map.items():
      device_size[device] += module_sizes[module]
   _SCREAMING_SNAKE_CASE       	=			"""\n""".join([F"- {device}: {size // 2**20}MiB" for device, size in device_size.items()]  )
   print(F"\nTheoretical use:\n{message}"  )

   _SCREAMING_SNAKE_CASE       	=			AutoTokenizer.from_pretrained(args.tokenizer_name  )

   _SCREAMING_SNAKE_CASE       	=			start_measure()
   _SCREAMING_SNAKE_CASE       	=			[]
   _SCREAMING_SNAKE_CASE       	=			[]
   _SCREAMING_SNAKE_CASE       	=			[]
   for prompt in PROMPTS:
      _SCREAMING_SNAKE_CASE       	=			tokenizer(SCREAMING_SNAKE_CASE_      , return_tensors="""pt"""  ).to(0  )
      _SCREAMING_SNAKE_CASE       	=			inputs["""input_ids"""][0].tolist()
      _SCREAMING_SNAKE_CASE       	=			time.time()
      _SCREAMING_SNAKE_CASE       	=			model.generate(inputs["""input_ids"""]  )
      _SCREAMING_SNAKE_CASE       	=			time.time()
      _SCREAMING_SNAKE_CASE       	=			outputs[0].tolist()
      _SCREAMING_SNAKE_CASE       	=			len(SCREAMING_SNAKE_CASE_  ) if outputs[: len(SCREAMING_SNAKE_CASE_  )] != tokens else len(SCREAMING_SNAKE_CASE_  ) - len(SCREAMING_SNAKE_CASE_  )
      _SCREAMING_SNAKE_CASE       	=			after_generate - before_generate

      _SCREAMING_SNAKE_CASE       	=			tokenizer.decode(SCREAMING_SNAKE_CASE_      , skip_special_tokens=SCREAMING_SNAKE_CASE_  )
      texts_outs.append(SCREAMING_SNAKE_CASE_  )
      generation_times.append(SCREAMING_SNAKE_CASE_  )
      gen_tokens.append(SCREAMING_SNAKE_CASE_  )
      print(F"Prompt: {prompt}\nGeneration {text_out}\nIn {generation_time:.2f}s for {num_gen_tokens} tokens\n"  )

   _SCREAMING_SNAKE_CASE       	=			end_measure(SCREAMING_SNAKE_CASE_  )
   log_measures(SCREAMING_SNAKE_CASE_      , """Model generation"""  )

   _SCREAMING_SNAKE_CASE       	=			[gen / tok for gen, tok in zip(SCREAMING_SNAKE_CASE_      , SCREAMING_SNAKE_CASE_  )]
   _SCREAMING_SNAKE_CASE       	=			sum(SCREAMING_SNAKE_CASE_  ) / len(SCREAMING_SNAKE_CASE_  )
   print(F"Average time of generation per token: {avg_gen:.2f}s"  )
   print(F"First generation (avg time per token): {generation_times_per_token[0]:.2f}s"  )
   _SCREAMING_SNAKE_CASE       	=			sum(generation_times_per_token[1:]  ) / (len(SCREAMING_SNAKE_CASE_  ) - 1)
   print(F"Average time of generation per token (excluding the first): {avg_gen:.2f}s"  )


if __name__ == "__main__":
     main()






