


'''simple docstring'''
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 				lowerCAmelCase__	(       UpperCAmelCase     ):






       """simple docstring"""





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







@require_distributed_retrieval
class 					_A						(						UpperCamelCase							):



   '''simple docstring'''






   def   __lowerCAmelCase						(	self	:					Any						)->	Dict:
          snake_case__					:	Tuple          =					tempfile.mkdtemp()
          snake_case__					:	Tuple          =					8

          # DPR tok
          snake_case__					:	Optional[int]          =					[
              """[UNK]""",
              """[CLS]""",
              """[SEP]""",
              """[PAD]""",
              """[MASK]""",
              """want""",
              """##want""",
              """##ed""",
              """wa""",
              """un""",
              """runn""",
              """##ing""",
              """,""",
              """low""",
              """lowest""",
          ]
          snake_case__					:	List[str]          =					os.path.join(self.tmpdirname					, """dpr_tokenizer"""						)
          os.makedirs(lowerCamelCase					, exist_ok=lowerCamelCase						)
          snake_case__					:	Any          =					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__					:	Optional[Any]          =					[
              """l""",
              """o""",
              """w""",
              """e""",
              """r""",
              """s""",
              """t""",
              """i""",
              """d""",
              """n""",
              """\u0120""",
              """\u0120l""",
              """\u0120n""",
              """\u0120lo""",
              """\u0120low""",
              """er""",
              """\u0120lowest""",
              """\u0120newer""",
              """\u0120wider""",
              """<unk>""",
          ]
          snake_case__					:	int          =					dict(zip(lowerCamelCase					, range(len(lowerCamelCase						)						)						)						)
          snake_case__					:	str          =					["""#version: 0.2""", """\u0120 l""", """\u0120l o""", """\u0120lo w""", """e r""", """"""]
          snake_case__					:	Dict          =					{"""unk_token""": """<unk>"""}

          snake_case__					:	Optional[Any]          =					os.path.join(self.tmpdirname					, """bart_tokenizer"""						)
          os.makedirs(lowerCamelCase					, exist_ok=lowerCamelCase						)
          snake_case__					:	str          =					os.path.join(lowerCamelCase					, BART_VOCAB_FILES_NAMES["""vocab_file"""]						)
          snake_case__					:	Any          =					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	:					Optional[Any]						)->	DPRQuestionEncoderTokenizer:
          return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					, """dpr_tokenizer"""						)						)




   def   __lowerCAmelCase						(	self	:					List[str]						)->	BartTokenizer:
          return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname					, """bart_tokenizer"""						)						)




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




   def   __lowerCAmelCase						(	self	:					Optional[int]						)->	Dict:
          snake_case__					:	Optional[int]          =					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	:					Optional[int]					, lowerCamelCase	:					bool					, lowerCamelCase	:					Union[str, Any]=12_345						)->	RagPyTorchDistributedRetriever:
          snake_case__					:	Dict          =					self.get_dummy_dataset()
          snake_case__					:	Tuple          =					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__					:	List[Any]          =					dataset
                 snake_case__					:	int          =					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	:					Any					, lowerCamelCase	:					bool						)->	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__					:	Any          =					RagConfig(
              retrieval_vector_size=self.retrieval_vector_size					, question_encoder=DPRConfig().to_dict()					, generator=BartConfig().to_dict()					, )
          snake_case__					:	Dict          =					ray.remote(lowerCamelCase						)
          snake_case__					:	Optional[int]          =					[remote_cls.remote() for _ in range(1						)]
          with patch("""transformers.models.rag.retrieval_rag.load_dataset"""						) as mock_load_dataset:
                 snake_case__					:	str          =					self.get_dummy_dataset()
                 snake_case__					:	Dict          =					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	:					List[Any]					, lowerCamelCase	:					bool					, lowerCamelCase	:					bool					, lowerCamelCase	:					Dict=12_345						)->	List[Any]:
          snake_case__					:	List[str]          =					self.get_dummy_dataset()
          snake_case__					:	List[Any]          =					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__					:	List[str]          =					os.path.join(self.tmpdirname					, """dataset"""						)
                 snake_case__					:	List[str]          =					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__					:	Tuple          =					RagPyTorchDistributedRetriever(
                     lowerCamelCase					, question_encoder_tokenizer=self.get_dpr_tokenizer()					, generator_tokenizer=self.get_bart_tokenizer()					, )
          else:
                 snake_case__					:	List[str]          =					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	:					Tuple					, lowerCamelCase	:					bool					, lowerCamelCase	:					bool						)->	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__					:	Optional[int]          =					self.get_dummy_dataset()
          snake_case__					:	int          =					RagConfig(
              retrieval_vector_size=self.retrieval_vector_size					, question_encoder=DPRConfig().to_dict()					, generator=BartConfig().to_dict()					, index_name="""custom"""					, )
          snake_case__					:	Union[str, Any]          =					ray.remote(lowerCamelCase						)
          snake_case__					:	Optional[int]          =					[remote_cls.remote() for _ in range(1						)]
          if from_disk:
                 snake_case__					:	int          =					os.path.join(self.tmpdirname					, """dataset"""						)
                 snake_case__					:	Union[str, Any]          =					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__					:	str          =					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__					:	Any          =					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	:					Union[str, Any]					, lowerCamelCase	:					RagRetriever					, lowerCamelCase	:					np.array					, lowerCamelCase	:					int						)->	None:
          snake_case__       ,   snake_case__       ,   snake_case__					:	Optional[Any]          =					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	:					str						)->	Optional[Any]:
          snake_case__					:	List[str]          =					1
          snake_case__					:	Optional[Any]          =					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	:					Dict						)->	Optional[Any]:
          snake_case__					:	Dict          =					1
          snake_case__					:	str          =					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	:					List[str]						)->	Any:
          snake_case__					:	Dict          =					1
          snake_case__					:	List[Any]          =					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	:					Tuple						)->	Any:
          snake_case__					:	Any          =					1
          snake_case__					:	Dict          =					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	:					int						)->	List[Any]:
          snake_case__					:	Union[str, Any]          =					1
          snake_case__					:	int          =					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						)->	str:
          snake_case__					:	Union[str, Any]          =					1
          snake_case__					:	List[str]          =					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
     )




