



'''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     snake_case__			(lowerCAmelCase_						):
    UpperCamelCase__	     =      TimesformerConfig()

    if "large" in model_name:
        UpperCamelCase__	     =      96

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

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

    return config





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

    return name





def     snake_case__			(lowerCAmelCase_       ,	lowerCAmelCase_						):
    for key in orig_state_dict.copy().keys():
        UpperCamelCase__	     =      orig_state_dict.pop(lowerCAmelCase_						)

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

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

    return orig_state_dict





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





def     snake_case__			(lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_						):
    UpperCamelCase__	     =      get_timesformer_config(lowerCAmelCase_						)

    UpperCamelCase__	     =      TimesformerForVideoClassification(lowerCAmelCase_						)

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

    model.load_state_dict(lowerCAmelCase_						)
    model.eval()

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

    UpperCamelCase__	     =      model(**lowerCAmelCase_						)
    UpperCamelCase__	     =      outputs.logits

    UpperCamelCase__	     =      [
        # 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__	     =      torch.Size([1, 4_00]						)
        UpperCamelCase__	     =      torch.tensor([-0.30_16, -0.77_13, -0.42_05]						)
    elif model_name == "timesformer-base-finetuned-k600":
        UpperCamelCase__	     =      torch.Size([1, 6_00]						)
        UpperCamelCase__	     =      torch.tensor([-0.72_67, -0.74_66, 3.24_04]						)
    elif model_name == "timesformer-base-finetuned-ssv2":
        UpperCamelCase__	     =      torch.Size([1, 1_74]						)
        UpperCamelCase__	     =      torch.tensor([-0.90_59, 0.64_33, -3.14_57]						)
    elif model_name == "timesformer-large-finetuned-k400":
        UpperCamelCase__	     =      torch.Size([1, 4_00]						)
        UpperCamelCase__	     =      torch.tensor([0, 0, 0]						)
    elif model_name == "timesformer-large-finetuned-k600":
        UpperCamelCase__	     =      torch.Size([1, 6_00]						)
        UpperCamelCase__	     =      torch.tensor([0, 0, 0]						)
    elif model_name == "timesformer-large-finetuned-ssv2":
        UpperCamelCase__	     =      torch.Size([1, 1_74]						)
        UpperCamelCase__	     =      torch.tensor([0, 0, 0]						)
    elif model_name == "timesformer-hr-finetuned-k400":
        UpperCamelCase__	     =      torch.Size([1, 4_00]						)
        UpperCamelCase__	     =      torch.tensor([-0.96_17, -3.73_11, -3.77_08]						)
    elif model_name == "timesformer-hr-finetuned-k600":
        UpperCamelCase__	     =      torch.Size([1, 6_00]						)
        UpperCamelCase__	     =      torch.tensor([2.52_73, 0.71_27, 1.88_48]						)
    elif model_name == "timesformer-hr-finetuned-ssv2":
        UpperCamelCase__	     =      torch.Size([1, 1_74]						)
        UpperCamelCase__	     =      torch.tensor([-3.67_56, -0.75_13, 0.71_80]						)
    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__":
      lowercase__  :	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.'
      )

      lowercase__  :	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 unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
      from transformers.models.electra.modeling_flax_electra import (
          FlaxElectraForCausalLM,
          FlaxElectraForMaskedLM,
          FlaxElectraForMultipleChoice,
          FlaxElectraForPreTraining,
          FlaxElectraForQuestionAnswering,
          FlaxElectraForSequenceClassification,
          FlaxElectraForTokenClassification,
          FlaxElectraModel,
      )

class    lowerCamelCase_				(						unittest.TestCase	):






   '''simple docstring'''

   def __init__(	self  ,    __UpperCamelCase  ,    __UpperCamelCase=13  ,    __UpperCamelCase=7  ,    __UpperCamelCase=True  ,    __UpperCamelCase=True  ,    __UpperCamelCase=True  ,    __UpperCamelCase=True  ,    __UpperCamelCase=99  ,    __UpperCamelCase=24  ,    __UpperCamelCase=32  ,    __UpperCamelCase=5  ,    __UpperCamelCase=4  ,    __UpperCamelCase=37  ,    __UpperCamelCase="gelu"  ,    __UpperCamelCase=0.1  ,    __UpperCamelCase=0.1  ,    __UpperCamelCase=5_12  ,    __UpperCamelCase=16  ,    __UpperCamelCase=2  ,    __UpperCamelCase=0.02  ,    __UpperCamelCase=4  ,    ):
       UpperCamelCase__	     =      parent
       UpperCamelCase__	     =      batch_size
       UpperCamelCase__	     =      seq_length
       UpperCamelCase__	     =      is_training
       UpperCamelCase__	     =      use_attention_mask
       UpperCamelCase__	     =      use_token_type_ids
       UpperCamelCase__	     =      use_labels
       UpperCamelCase__	     =      vocab_size
       UpperCamelCase__	     =      hidden_size
       UpperCamelCase__	     =      embedding_size
       UpperCamelCase__	     =      num_hidden_layers
       UpperCamelCase__	     =      num_attention_heads
       UpperCamelCase__	     =      intermediate_size
       UpperCamelCase__	     =      hidden_act
       UpperCamelCase__	     =      hidden_dropout_prob
       UpperCamelCase__	     =      attention_probs_dropout_prob
       UpperCamelCase__	     =      max_position_embeddings
       UpperCamelCase__	     =      type_vocab_size
       UpperCamelCase__	     =      type_sequence_label_size
       UpperCamelCase__	     =      initializer_range
       UpperCamelCase__	     =      num_choices


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      ids_tensor([self.batch_size, self.seq_length]  ,    self.vocab_size			)

       UpperCamelCase__	     =      None
       if self.use_attention_mask:
           UpperCamelCase__	     =      random_attention_mask([self.batch_size, self.seq_length]			)

       UpperCamelCase__	     =      None
       if self.use_token_type_ids:
           UpperCamelCase__	     =      ids_tensor([self.batch_size, self.seq_length]  ,    self.type_vocab_size			)

       UpperCamelCase__	     =      ElectraConfig(
           vocab_size=self.vocab_size  ,    hidden_size=self.hidden_size  ,    embedding_size=self.embedding_size  ,    num_hidden_layers=self.num_hidden_layers  ,    num_attention_heads=self.num_attention_heads  ,    intermediate_size=self.intermediate_size  ,    hidden_act=self.hidden_act  ,    hidden_dropout_prob=self.hidden_dropout_prob  ,    attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,    max_position_embeddings=self.max_position_embeddings  ,    type_vocab_size=self.type_vocab_size  ,    initializer_range=self.initializer_range  ,    )

       return config, input_ids, token_type_ids, attention_mask







   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      self.prepare_config_and_inputs()
       UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__	     =      config_and_inputs
       UpperCamelCase__	     =      {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
       return config, inputs_dict

@require_flax
class    lowerCamelCase_				(						lowerCAmelCase_       ,     unittest.TestCase	):






   '''simple docstring'''

   A			:  Dict					=						True

   A			:  Optional[Any]					=						(
       (
           FlaxElectraModel,
           FlaxElectraForCausalLM,
           FlaxElectraForMaskedLM,
           FlaxElectraForPreTraining,
           FlaxElectraForTokenClassification,
           FlaxElectraForQuestionAnswering,
           FlaxElectraForMultipleChoice,
           FlaxElectraForSequenceClassification,
       )
       if is_flax_available()
       else ()
   )


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      FlaxElectraModelTester(self			)







   @slow
   def 						_UpperCamelCase							(	self			):
       for model_class_name in self.all_model_classes:
           if model_class_name == FlaxElectraForMaskedLM:
               UpperCamelCase__	     =      model_class_name.from_pretrained("google/electra-small-generator"			)
           else:
               UpperCamelCase__	     =      model_class_name.from_pretrained("google/electra-small-discriminator"			)
           UpperCamelCase__	     =      model(np.ones((1, 1)			)			)
           self.assertIsNotNone(__UpperCamelCase			)




