





"""simple docstring"""
import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing the checkpointing capability,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################

lowerCamelCase_          =	1_6
lowerCamelCase_          =	3_2



def   __lowerCamelCase     (							a_  :	List[Any]				,     a_  :	int = 16							)    ->   Dict:
      __SCREAMING_SNAKE_CASE     :Tuple	       =						AutoTokenizer.from_pretrained('''bert-base-cased'''							)
      __SCREAMING_SNAKE_CASE     :Dict	       =						load_dataset('''glue'''				,     '''mrpc'''							)

      def tokenize_function(a_  :	List[Any]							):
            # max_length=None => use the model max length (it's actually the default)
            __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						tokenizer(examples['''sentence1''']				,     examples['''sentence2''']				,     truncation=__A				,     max_length=__A							)
            return outputs

      # Apply the method we just defined to all the examples in all the splits of the dataset
      # starting with the main process first:
      with accelerator.main_process_first():
            __SCREAMING_SNAKE_CASE     :int	       =						datasets.map(
                __A				,     batched=__A				,     remove_columns=['''idx''', '''sentence1''', '''sentence2''']				,     )

      # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
      # transformers library
      __SCREAMING_SNAKE_CASE     :List[Any]	       =						tokenized_datasets.rename_column('''label'''				,     '''labels'''							)

      def collate_fn(a_  :	Optional[Any]							):
            # On TPU it's best to pad everything to the same length or training will be very slow.
            __SCREAMING_SNAKE_CASE     :str	       =						1_28 if accelerator.distributed_type == DistributedType.TPU else None
            # When using mixed precision we want round multiples of 8/16
            if accelerator.mixed_precision == "fp8":
                  __SCREAMING_SNAKE_CASE     :int	       =						16
            elif accelerator.mixed_precision != "no":
                  __SCREAMING_SNAKE_CASE     :Any	       =						8
            else:
                  __SCREAMING_SNAKE_CASE     :Tuple	       =						None

            return tokenizer.pad(
                __A				,     padding='''longest'''				,     max_length=__A				,     pad_to_multiple_of=__A				,     return_tensors='''pt'''				,     )

      # Instantiate dataloaders.
      __SCREAMING_SNAKE_CASE     :Dict	       =						DataLoader(
          tokenized_datasets['''train''']				,     shuffle=__A				,     collate_fn=__A				,     batch_size=__A							)
      __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						DataLoader(
          tokenized_datasets['''validation''']				,     shuffle=__A				,     collate_fn=__A				,     batch_size=__A							)

      return train_dataloader, eval_dataloader


# For testing only
if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
  from accelerate.test_utils.training import mocked_dataloaders

  lowerCamelCase_          =	mocked_dataloaders  # noqa: F811



def   __lowerCamelCase     (							a_  :	Dict				,     a_  :	Dict							)    ->   Optional[int]:



      if os.environ.get('''TESTING_MOCKED_DATALOADERS'''				,     __A							) == "1":
            __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						2
      # Initialize accelerator
      __SCREAMING_SNAKE_CASE     :List[str]	       =						Accelerator(cpu=args.cpu				,     mixed_precision=args.mixed_precision							)
      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      __SCREAMING_SNAKE_CASE     :Any	       =						config['''lr''']
      __SCREAMING_SNAKE_CASE     :str	       =						int(config['''num_epochs''']							)
      __SCREAMING_SNAKE_CASE     :int	       =						int(config['''seed''']							)
      __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						int(config['''batch_size''']							)

      # New Code #
      # Parse out whether we are saving every epoch or after a certain number of batches
      if hasattr(args.checkpointing_steps				,     '''isdigit'''							):
            if args.checkpointing_steps == "epoch":
                  __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						args.checkpointing_steps
            elif args.checkpointing_steps.isdigit():
                  __SCREAMING_SNAKE_CASE     :Tuple	       =						int(args.checkpointing_steps							)
            else:
                  raise ValueError(
                      f'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.'''							)
      else:
            __SCREAMING_SNAKE_CASE     :Any	       =						None

      set_seed(__A							)

      __SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE     :Tuple	       =						get_dataloaders(__A				,     __A							)
      __SCREAMING_SNAKE_CASE     :List[Any]	       =						evaluate.load('''glue'''				,     '''mrpc'''							)

      # If the batch size is too big we use gradient accumulation
      __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						1
      if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
            __SCREAMING_SNAKE_CASE     :List[str]	       =						batch_size // MAX_GPU_BATCH_SIZE
            __SCREAMING_SNAKE_CASE     :Optional[int]	       =						MAX_GPU_BATCH_SIZE

      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      __SCREAMING_SNAKE_CASE     :Dict	       =						AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''				,     return_dict=__A							)

      # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
      # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
      # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
      __SCREAMING_SNAKE_CASE     :Tuple	       =						model.to(accelerator.device							)

      # Instantiate optimizer
      __SCREAMING_SNAKE_CASE     :Tuple	       =						AdamW(params=model.parameters()				,     lr=__A							)

      # Instantiate scheduler
      __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						get_linear_schedule_with_warmup(
          optimizer=__A				,     num_warmup_steps=1_00				,     num_training_steps=(len(__A							) * num_epochs) // gradient_accumulation_steps				,     )

      # Prepare everything
      # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
      # prepare method.
      __SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE     :Tuple	       =						accelerator.prepare(
          __A				,     __A				,     __A				,     __A				,     __A							)

      # New Code #
      # We need to keep track of how many total steps we have iterated over
      __SCREAMING_SNAKE_CASE     :Dict	       =						0
      # We also need to keep track of the stating epoch so files are named properly
      __SCREAMING_SNAKE_CASE     :Tuple	       =						0

      # We need to load the checkpoint back in before training here with `load_state`
      # The total number of epochs is adjusted based on where the state is being loaded from,
      # as we assume continuation of the same training script
      if args.resume_from_checkpoint:
            if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
                  accelerator.print(f'''Resumed from checkpoint: {args.resume_from_checkpoint}'''							)
                  accelerator.load_state(args.resume_from_checkpoint							)
                  __SCREAMING_SNAKE_CASE     :Dict	       =						os.path.basename(args.resume_from_checkpoint							)
            else:
                  # Get the most recent checkpoint
                  __SCREAMING_SNAKE_CASE     :Tuple	       =						[f.name for f in os.scandir(os.getcwd()							) if f.is_dir()]
                  dirs.sort(key=os.path.getctime							)
                  __SCREAMING_SNAKE_CASE     :List[str]	       =						dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
            # Extract `epoch_{i}` or `step_{i}`
            __SCREAMING_SNAKE_CASE     :Tuple	       =						os.path.splitext(__A							)[0]

            if "epoch" in training_difference:
                  __SCREAMING_SNAKE_CASE     :Tuple	       =						int(training_difference.replace('''epoch_'''				,     ''''''							)							) + 1
                  __SCREAMING_SNAKE_CASE     :Optional[int]	       =						None
            else:
                  __SCREAMING_SNAKE_CASE     :Optional[int]	       =						int(training_difference.replace('''step_'''				,     ''''''							)							)
                  __SCREAMING_SNAKE_CASE     :Dict	       =						resume_step // len(__A							)
                  resume_step -= starting_epoch * len(__A							)

                   # Now we train the model
      for epoch in range(__A				,     __A							):
            model.train()
            # New Code #
            if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
                  # We need to skip steps until we reach the resumed step
                  __SCREAMING_SNAKE_CASE     :Tuple	       =						accelerator.skip_first_batches(__A				,     __A							)
                  overall_step += resume_step
            else:
                  # After the first iteration though, we need to go back to the original dataloader
                  __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						train_dataloader
            for step, batch in enumerate(__A							):
                  # We could avoid this line since we set the accelerator with `device_placement=True`.
                  batch.to(accelerator.device							)
                  __SCREAMING_SNAKE_CASE     :str	       =						model(**__A							)
                  __SCREAMING_SNAKE_CASE     :Optional[int]	       =						outputs.loss
                  __SCREAMING_SNAKE_CASE     :List[Any]	       =						loss / gradient_accumulation_steps
                  accelerator.backward(__A							)
                  if step % gradient_accumulation_steps == 0:
                        optimizer.step()
                        lr_scheduler.step()
                        optimizer.zero_grad()
                  # New Code #
                  overall_step += 1

                  # New Code #
                  # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
                  # These are saved to folders named `step_{overall_step}`
                  # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
                  # If mixed precision was used, will also save a "scalar.bin" file
                  if isinstance(__A				,     __A							):
                        __SCREAMING_SNAKE_CASE     :List[Any]	       =						f'''step_{overall_step}'''
                        if overall_step % checkpointing_steps == 0:
                              if args.output_dir is not None:
                                    __SCREAMING_SNAKE_CASE     :List[str]	       =						os.path.join(args.output_dir				,     __A							)
                              accelerator.save_state(__A							)

            model.eval()
            for step, batch in enumerate(__A							):
                  # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
                  batch.to(accelerator.device							)
                  with torch.no_grad():
                        __SCREAMING_SNAKE_CASE     :Optional[int]	       =						model(**__A							)
                  __SCREAMING_SNAKE_CASE     :Tuple	       =						outputs.logits.argmax(dim=-1							)
                  __SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE     :Optional[Any]	       =						accelerator.gather_for_metrics((predictions, batch['''labels'''])							)
                  metric.add_batch(
                      predictions=__A				,     references=__A				,     )

            __SCREAMING_SNAKE_CASE     :List[str]	       =						metric.compute()
            # Use accelerator.print to print only on the main process.
            accelerator.print(f'''epoch {epoch}:'''				,     __A							)

            # New Code #
            # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
            # These are saved to folders named `epoch_{epoch}`
            # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
            # If mixed precision was used, will also save a "scalar.bin" file
            if checkpointing_steps == "epoch":
                  __SCREAMING_SNAKE_CASE     :str	       =						f'''epoch_{epoch}'''
                  if args.output_dir is not None:
                        __SCREAMING_SNAKE_CASE     :List[Any]	       =						os.path.join(args.output_dir				,     __A							)
                  accelerator.save_state(__A							)



def   __lowerCamelCase     (							)    ->   List[Any]:
      __SCREAMING_SNAKE_CASE     :Optional[int]	       =						argparse.ArgumentParser(description='''Simple example of training script.'''							)
      parser.add_argument(
          '''--mixed_precision'''				,     type=__A				,     default=__A				,     choices=['''no''', '''fp16''', '''bf16''', '''fp8''']				,     help='''Whether to use mixed precision. Choose'''
          '''between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'''
          '''and an Nvidia Ampere GPU.'''				,     )
      parser.add_argument('''--cpu'''				,     action='''store_true'''				,     help='''If passed, will train on the CPU.'''							)
      parser.add_argument(
          '''--checkpointing_steps'''				,     type=__A				,     default=__A				,     help='''Whether the various states should be saved at the end of every n steps, or \'epoch\' for each epoch.'''				,     )
      parser.add_argument(
          '''--output_dir'''				,     type=__A				,     default='''.'''				,     help='''Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'''				,     )
      parser.add_argument(
          '''--resume_from_checkpoint'''				,     type=__A				,     default=__A				,     help='''If the training should continue from a checkpoint folder.'''				,     )
      __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						parser.parse_args()
      __SCREAMING_SNAKE_CASE     :Dict	       =						{'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
      training_function(__A				,     __A							)


if __name__ == "__main__":
  main()






"""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   __lowerCamelCase     (							a_  :	Union[str, Any]							)    ->   Tuple:
      __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						TimesformerConfig()

      if "large" in model_name:
            __SCREAMING_SNAKE_CASE     :Dict	       =						96

      if "hr" in model_name:
            __SCREAMING_SNAKE_CASE     :List[Any]	       =						16
            __SCREAMING_SNAKE_CASE     :int	       =						4_48

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

      return config



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

      return name



def   __lowerCamelCase     (							a_  :	Optional[Any]				,     a_  :	Tuple							)    ->   Optional[int]:
      for key in orig_state_dict.copy().keys():
            __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						orig_state_dict.pop(a_							)

            if key.startswith('''model.'''							):
                  __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						key.replace('''model.'''				,     ''''''							)

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

      return orig_state_dict



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



def   __lowerCamelCase     (							a_  :	List[Any]				,     a_  :	int				,     a_  :	Optional[int]				,     a_  :	Optional[Any]							)    ->   Union[str, Any]:
      __SCREAMING_SNAKE_CASE     :List[Any]	       =						get_timesformer_config(a_							)

      __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						TimesformerForVideoClassification(a_							)

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

      model.load_state_dict(a_							)
      model.eval()

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

      __SCREAMING_SNAKE_CASE     :Tuple	       =						model(**a_							)
      __SCREAMING_SNAKE_CASE     :List[str]	       =						outputs.logits

      __SCREAMING_SNAKE_CASE     :Dict	       =						[
          # 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     :Union[str, Any]	       =						torch.Size([1, 4_00]							)
            __SCREAMING_SNAKE_CASE     :int	       =						torch.tensor([-0.3016, -0.7713, -0.4205]							)
      elif model_name == "timesformer-base-finetuned-k600":
            __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						torch.Size([1, 6_00]							)
            __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						torch.tensor([-0.7267, -0.7466, 3.2404]							)
      elif model_name == "timesformer-base-finetuned-ssv2":
            __SCREAMING_SNAKE_CASE     :List[str]	       =						torch.Size([1, 1_74]							)
            __SCREAMING_SNAKE_CASE     :Tuple	       =						torch.tensor([-0.9059, 0.6433, -3.1457]							)
      elif model_name == "timesformer-large-finetuned-k400":
            __SCREAMING_SNAKE_CASE     :Tuple	       =						torch.Size([1, 4_00]							)
            __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						torch.tensor([0, 0, 0]							)
      elif model_name == "timesformer-large-finetuned-k600":
            __SCREAMING_SNAKE_CASE     :List[Any]	       =						torch.Size([1, 6_00]							)
            __SCREAMING_SNAKE_CASE     :str	       =						torch.tensor([0, 0, 0]							)
      elif model_name == "timesformer-large-finetuned-ssv2":
            __SCREAMING_SNAKE_CASE     :int	       =						torch.Size([1, 1_74]							)
            __SCREAMING_SNAKE_CASE     :Any	       =						torch.tensor([0, 0, 0]							)
      elif model_name == "timesformer-hr-finetuned-k400":
            __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						torch.Size([1, 4_00]							)
            __SCREAMING_SNAKE_CASE     :Dict	       =						torch.tensor([-0.9617, -3.7311, -3.7708]							)
      elif model_name == "timesformer-hr-finetuned-k600":
            __SCREAMING_SNAKE_CASE     :int	       =						torch.Size([1, 6_00]							)
            __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						torch.tensor([2.5273, 0.7127, 1.8848]							)
      elif model_name == "timesformer-hr-finetuned-ssv2":
            __SCREAMING_SNAKE_CASE     :List[Any]	       =						torch.Size([1, 1_74]							)
            __SCREAMING_SNAKE_CASE     :List[Any]	       =						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]				,     a_				,     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(a_							)
            model.save_pretrained(a_							)

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


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

  lowerCamelCase_          =	parser.parse_args()
  convert_timesformer_checkpoint(
      args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
  )