

import json
import os
import shutil
import sys
import tempfile
import unittest
from unittest import TestCase
from unittest.mock import patch

import faiss
import numpy as np
from datasets import Dataset

from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
from transformers.integrations import is_ray_available
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_ray


sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip

if is_torch_available():
    from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
else:
    lowerCamelCase_					  =   None

if is_ray_available():
    import ray  # noqa: E402 # isort:skip
    from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
    lowerCamelCase_					  =   None
    lowerCamelCase_					  =   None
    lowerCamelCase_					  =   None



def   UpperCamelCase(   lowercase_				)      ->							Optional[Any]:





  '''simple docstring'''





  if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
    snake_case_					= unittest.skip("""test requires Datasets, Faiss, psutil"""				)(lowercase_				)
  return test_case







@require_distributed_retrieval
class 					__lowerCamelCase							(      __snake_case      ):







      def       lowerCAmelCase_       (	self  )		->			Dict:
        snake_case_					= tempfile.mkdtemp()
        snake_case_					= 8

        # DPR tok
        snake_case_					= [
            """[UNK]""",
            """[CLS]""",
            """[SEP]""",
            """[PAD]""",
            """[MASK]""",
            """want""",
            """##want""",
            """##ed""",
            """wa""",
            """un""",
            """runn""",
            """##ing""",
            """,""",
            """low""",
            """lowest""",
        ]
        snake_case_					= os.path.join(self.tmpdirname				,    """dpr_tokenizer"""  )
        os.makedirs(lowerCamelCase				,    exist_ok=lowerCamelCase  )
        snake_case_					= os.path.join(lowerCamelCase				,    DPR_VOCAB_FILES_NAMES["""vocab_file"""]  )
        with open(self.vocab_file				,    """w"""				,    encoding="""utf-8"""  ) as vocab_writer:
          vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]  )  )

        # BART tok
        snake_case_					= [
            """l""",
            """o""",
            """w""",
            """e""",
            """r""",
            """s""",
            """t""",
            """i""",
            """d""",
            """n""",
            """\u0120""",
            """\u0120l""",
            """\u0120n""",
            """\u0120lo""",
            """\u0120low""",
            """er""",
            """\u0120lowest""",
            """\u0120newer""",
            """\u0120wider""",
            """<unk>""",
        ]
        snake_case_					= dict(zip(lowerCamelCase				,    range(len(lowerCamelCase  )  )  )  )
        snake_case_					= ["""#version: 0.2""", """\u0120 l""", """\u0120l o""", """\u0120lo w""", """e r""", """"""]
        snake_case_					= {"""unk_token""": """<unk>"""}

        snake_case_					= os.path.join(self.tmpdirname				,    """bart_tokenizer"""  )
        os.makedirs(lowerCamelCase				,    exist_ok=lowerCamelCase  )
        snake_case_					= os.path.join(lowerCamelCase				,    BART_VOCAB_FILES_NAMES["""vocab_file"""]  )
        snake_case_					= os.path.join(lowerCamelCase				,    BART_VOCAB_FILES_NAMES["""merges_file"""]  )
        with open(self.vocab_file				,    """w"""				,    encoding="""utf-8"""  ) as fp:
          fp.write(json.dumps(lowerCamelCase  ) + """\n"""  )
        with open(self.merges_file				,    """w"""				,    encoding="""utf-8"""  ) as fp:
          fp.write("""\n""".join(lowerCamelCase  )  )







      def       lowerCAmelCase_       (	self  )		->			DPRQuestionEncoderTokenizer:
        return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname				,    """dpr_tokenizer"""  )  )







      def       lowerCAmelCase_       (	self  )		->			BartTokenizer:
        return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname				,    """bart_tokenizer"""  )  )







      def       lowerCAmelCase_       (	self  )		->			List[str]:
        shutil.rmtree(self.tmpdirname  )







      def       lowerCAmelCase_       (	self  )		->			Dict:
        snake_case_					= Dataset.from_dict(
            {
                """id""": ["""0""", """1"""],
                """text""": ["""foo""", """bar"""],
                """title""": ["""Foo""", """Bar"""],
                """embeddings""": [np.ones(self.retrieval_vector_size  ), 2 * np.ones(self.retrieval_vector_size  )],
            }  )
        dataset.add_faiss_index("""embeddings"""				,    string_factory="""Flat"""				,    metric_type=faiss.METRIC_INNER_PRODUCT  )
        return dataset







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase=12345  )		->			RagPyTorchDistributedRetriever:
        snake_case_					= self.get_dummy_dataset()
        snake_case_					= RagConfig(
            retrieval_vector_size=self.retrieval_vector_size				,    question_encoder=DPRConfig().to_dict()				,    generator=BartConfig().to_dict()				,    )
        with patch("""transformers.models.rag.retrieval_rag.load_dataset"""  ) as mock_load_dataset:
          snake_case_					= dataset
          snake_case_					= RagPyTorchDistributedRetriever(
              lowerCamelCase				,    question_encoder_tokenizer=self.get_dpr_tokenizer()				,    generator_tokenizer=self.get_bart_tokenizer()				,    )
          if init_retrieval:
            retriever.init_retrieval(lowerCamelCase  )
        return retriever







      def       lowerCAmelCase_       (	self				,    lowerCamelCase  )		->			RagRayDistributedRetriever:
        # Have to run in local mode because sys.path modifications at top of
        # file are not propogated to remote workers.
        # https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
        ray.init(local_mode=lowerCamelCase  )
        snake_case_					= RagConfig(
            retrieval_vector_size=self.retrieval_vector_size				,    question_encoder=DPRConfig().to_dict()				,    generator=BartConfig().to_dict()				,    )
        snake_case_					= ray.remote(lowerCamelCase  )
        snake_case_					= [remote_cls.remote() for _ in range(1  )]
        with patch("""transformers.models.rag.retrieval_rag.load_dataset"""  ) as mock_load_dataset:
          snake_case_					= self.get_dummy_dataset()
          snake_case_					= RagRayDistributedRetriever(
              lowerCamelCase				,    question_encoder_tokenizer=self.get_dpr_tokenizer()				,    generator_tokenizer=self.get_bart_tokenizer()				,    retrieval_workers=lowerCamelCase				,    )
          if init_retrieval:
            retriever.init_retrieval()
        return retriever







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase=12345  )		->			List[Any]:
        snake_case_					= self.get_dummy_dataset()
        snake_case_					= RagConfig(
            retrieval_vector_size=self.retrieval_vector_size				,    question_encoder=DPRConfig().to_dict()				,    generator=BartConfig().to_dict()				,    index_name="""custom"""				,    )
        if from_disk:
          snake_case_					= os.path.join(self.tmpdirname				,    """dataset"""  )
          snake_case_					= os.path.join(self.tmpdirname				,    """index.faiss"""  )
          dataset.get_index("""embeddings"""  ).save(os.path.join(self.tmpdirname				,    """index.faiss"""  )  )
          dataset.drop_index("""embeddings"""  )
          dataset.save_to_disk(os.path.join(self.tmpdirname				,    """dataset"""  )  )
          del dataset
          snake_case_					= RagPyTorchDistributedRetriever(
              lowerCamelCase				,    question_encoder_tokenizer=self.get_dpr_tokenizer()				,    generator_tokenizer=self.get_bart_tokenizer()				,    )
        else:
          snake_case_					= RagPyTorchDistributedRetriever(
              lowerCamelCase				,    question_encoder_tokenizer=self.get_dpr_tokenizer()				,    generator_tokenizer=self.get_bart_tokenizer()				,    index=CustomHFIndex(config.retrieval_vector_size				,    lowerCamelCase  )				,    )
        if init_retrieval:
          retriever.init_retrieval(lowerCamelCase  )
        return retriever







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase  )		->			str:
        # Have to run in local mode because sys.path modifications at top of
        # file are not propogated to remote workers.
        # https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
        ray.init(local_mode=lowerCamelCase  )
        snake_case_					= self.get_dummy_dataset()
        snake_case_					= RagConfig(
            retrieval_vector_size=self.retrieval_vector_size				,    question_encoder=DPRConfig().to_dict()				,    generator=BartConfig().to_dict()				,    index_name="""custom"""				,    )
        snake_case_					= ray.remote(lowerCamelCase  )
        snake_case_					= [remote_cls.remote() for _ in range(1  )]
        if from_disk:
          snake_case_					= os.path.join(self.tmpdirname				,    """dataset"""  )
          snake_case_					= os.path.join(self.tmpdirname				,    """index.faiss"""  )
          dataset.get_index("""embeddings"""  ).save(os.path.join(self.tmpdirname				,    """index.faiss"""  )  )
          dataset.drop_index("""embeddings"""  )
          dataset.save_to_disk(os.path.join(self.tmpdirname				,    """dataset"""  )  )
          del dataset
          snake_case_					= RagRayDistributedRetriever(
              lowerCamelCase				,    question_encoder_tokenizer=self.get_dpr_tokenizer()				,    generator_tokenizer=self.get_bart_tokenizer()				,    retrieval_workers=lowerCamelCase				,    index=CustomHFIndex.load_from_disk(
                  vector_size=config.retrieval_vector_size				,    dataset_path=config.passages_path				,    index_path=config.index_path				,    )				,    )
        else:
          snake_case_					= RagRayDistributedRetriever(
              lowerCamelCase				,    question_encoder_tokenizer=self.get_dpr_tokenizer()				,    generator_tokenizer=self.get_bart_tokenizer()				,    retrieval_workers=lowerCamelCase				,    index=CustomHFIndex(config.retrieval_vector_size				,    lowerCamelCase  )				,    )
        if init_retrieval:
          retriever.init_retrieval()
        return retriever







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			None:
        snake_case_					= retriever.retrieve(lowerCamelCase				,    n_docs=lowerCamelCase  )
        self.assertEqual(retrieved_doc_embeds.shape				,    (2, n_docs, self.retrieval_vector_size)  )
        self.assertEqual(len(lowerCamelCase  )				,    2  )
        self.assertEqual(sorted(doc_dicts[0]  )				,    ["""embeddings""", """id""", """text""", """title"""]  )
        self.assertEqual(len(doc_dicts[0]["""id"""]  )				,    lowerCamelCase  )
        self.assertEqual(doc_dicts[0]["""id"""][0]				,    """1"""  )  # max inner product is reached with second doc
        self.assertEqual(doc_dicts[1]["""id"""][0]				,    """0"""  )  # max inner product is reached with first doc
        self.assertListEqual(doc_ids.tolist()				,    [[1], [0]]  )







      def       lowerCAmelCase_       (	self  )		->			Optional[Any]:
        snake_case_					= 1
        snake_case_					= np.array(
            [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]				,    dtype=np.floataa  )

        self.distributed_retriever_check(
            self.get_dummy_pytorch_distributed_retriever(init_retrieval=lowerCamelCase  )				,    lowerCamelCase				,    lowerCamelCase  )







      def       lowerCAmelCase_       (	self  )		->			Optional[Any]:
        snake_case_					= 1
        snake_case_					= np.array(
            [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]				,    dtype=np.floataa  )

        self.distributed_retriever_check(
            self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=lowerCamelCase				,    from_disk=lowerCamelCase  )				,    lowerCamelCase				,    lowerCamelCase				,    )







      def       lowerCAmelCase_       (	self  )		->			Any:
        snake_case_					= 1
        snake_case_					= np.array(
            [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]				,    dtype=np.floataa  )

        self.distributed_retriever_check(
            self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=lowerCamelCase				,    from_disk=lowerCamelCase  )				,    lowerCamelCase				,    lowerCamelCase				,    )







      @require_ray
      def       lowerCAmelCase_       (	self  )		->			Any:
        snake_case_					= 1
        snake_case_					= np.array(
            [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]				,    dtype=np.floataa  )

        self.distributed_retriever_check(
            self.get_dummy_ray_distributed_retriever(init_retrieval=lowerCamelCase  )				,    lowerCamelCase				,    lowerCamelCase  )
        ray.shutdown()







      @require_ray
      def       lowerCAmelCase_       (	self  )		->			List[Any]:
        snake_case_					= 1
        snake_case_					= np.array(
            [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]				,    dtype=np.floataa  )
        with self.assertRaises(lowerCamelCase  ):
          self.distributed_retriever_check(
              self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=lowerCamelCase				,    from_disk=lowerCamelCase  )				,    lowerCamelCase				,    lowerCamelCase				,    )
        ray.shutdown()







      @require_ray
      def       lowerCAmelCase_       (	self  )		->			str:
        snake_case_					= 1
        snake_case_					= np.array(
            [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]				,    dtype=np.floataa  )

        self.distributed_retriever_check(
            self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=lowerCamelCase				,    from_disk=lowerCamelCase  )				,    lowerCamelCase				,    lowerCamelCase  )
        ray.shutdown()



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






       """simple docstring"""





       snake_case__					:	Dict          =					TimesformerConfig()

       if "large" in model_name:
              snake_case__					:	int          =					96

       if "hr" in model_name:
              snake_case__					:	int          =					16
              snake_case__					:	int          =					448

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

       return config



def 				lowerCAmelCase__	(       UpperCAmelCase     ):






       """simple docstring"""





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

       return name



def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase     ):






       """simple docstring"""





       for key in orig_state_dict.copy().keys():
              snake_case__					:	Any          =					orig_state_dict.pop(UpperCAmelCase     )

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

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

       return orig_state_dict



def 				lowerCAmelCase__	(       ):






       """simple docstring"""





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



def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase    ,       UpperCAmelCase    ,       UpperCAmelCase     ):






       """simple docstring"""





       snake_case__					:	Any          =					get_timesformer_config(UpperCAmelCase     )

       snake_case__					:	Optional[int]          =					TimesformerForVideoClassification(UpperCAmelCase     )

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

       model.load_state_dict(UpperCAmelCase     )
       model.eval()

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

       snake_case__					:	int          =					model(**UpperCAmelCase     )
       snake_case__					:	Tuple          =					outputs.logits

       snake_case__					:	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":
              snake_case__					:	str          =					torch.Size([1, 400]     )
              snake_case__					:	str          =					torch.tensor([-0.3_0_1_6, -0.7_7_1_3, -0.4_2_0_5]     )
       elif model_name == "timesformer-base-finetuned-k600":
              snake_case__					:	Dict          =					torch.Size([1, 600]     )
              snake_case__					:	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":
              snake_case__					:	List[Any]          =					torch.Size([1, 174]     )
              snake_case__					:	Union[str, 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":
              snake_case__					:	Optional[Any]          =					torch.Size([1, 400]     )
              snake_case__					:	Tuple          =					torch.tensor([0, 0, 0]     )
       elif model_name == "timesformer-large-finetuned-k600":
              snake_case__					:	Optional[int]          =					torch.Size([1, 600]     )
              snake_case__					:	int          =					torch.tensor([0, 0, 0]     )
       elif model_name == "timesformer-large-finetuned-ssv2":
              snake_case__					:	Union[str, Any]          =					torch.Size([1, 174]     )
              snake_case__					:	List[Any]          =					torch.tensor([0, 0, 0]     )
       elif model_name == "timesformer-hr-finetuned-k400":
              snake_case__					:	str          =					torch.Size([1, 400]     )
              snake_case__					:	List[Any]          =					torch.tensor([-0.9_6_1_7, -3.7_3_1_1, -3.7_7_0_8]     )
       elif model_name == "timesformer-hr-finetuned-k600":
              snake_case__					:	Optional[int]          =					torch.Size([1, 600]     )
              snake_case__					:	str          =					torch.tensor([2.5_2_7_3, 0.7_1_2_7, 1.8_8_4_8]     )
       elif model_name == "timesformer-hr-finetuned-ssv2":
              snake_case__					:	int          =					torch.Size([1, 174]     )
              snake_case__					:	Optional[Any]          =					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]    ,       UpperCAmelCase    ,       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(UpperCAmelCase     )
              model.save_pretrained(UpperCAmelCase     )

       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
     )




