



'''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    (   _lowercase						) ->     str:
							UpperCAmelCase	:						Any					=						TimesformerConfig()

							if "large" in model_name:
														UpperCAmelCase	:						Optional[int]					=						9_6

							if "hr" in model_name:
														UpperCAmelCase	:						str					=						1_6
														UpperCAmelCase	:						Tuple					=						4_4_8

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

							return config



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

							return name



def  __lowerCamelCase    (   _lowercase      ,							_lowercase						) ->     Any:
							for key in orig_state_dict.copy().keys():
														UpperCAmelCase	:						str					=						orig_state_dict.pop(_lowercase						)

														if key.startswith("""model."""						):
																					UpperCAmelCase	:						Optional[int]					=						key.replace("""model."""      ,							""""""						)

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

							return orig_state_dict



def  __lowerCamelCase    (   ) ->     int:
							UpperCAmelCase	:						List[Any]					=						hf_hub_download(
							    repo_id="""hf-internal-testing/spaghetti-video"""      ,							filename="""eating_spaghetti.npy"""      ,							repo_type="""dataset"""						)
							UpperCAmelCase	:						str					=						np.load(_lowercase						)
							return list(_lowercase						)



def  __lowerCamelCase    (   _lowercase      ,							_lowercase      ,							_lowercase      ,							_lowercase						) ->     Optional[Any]:
							UpperCAmelCase	:						str					=						get_timesformer_config(_lowercase						)

							UpperCAmelCase	:						str					=						TimesformerForVideoClassification(_lowercase						)

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

							model.load_state_dict(_lowercase						)
							model.eval()

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

							UpperCAmelCase	:						List[str]					=						model(**_lowercase						)
							UpperCAmelCase	:						str					=						outputs.logits

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

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


if __name__ == "__main__":
							a  :     Union[str, Any]       =       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."""
							)

							a  :     List[Any]       =       parser.parse_args()
							convert_timesformer_checkpoint(
							    args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
							)






'''simple docstring'''




import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils_fast import AddedToken, PreTrainedTokenizerFast
from ...utils import logging


a  :     Any       =       logging.get_logger(__name__)

a  :     int       =       {"""vocab_file""": """spiece.model""", """tokenizer_file""": """tokenizer.json"""}

a  :     List[Any]       =       {
    """vocab_file""": {
        """TsinghuaAI/CPM-Generate""": """https://huggingface.co/TsinghuaAI/CPM-Generate/resolve/main/spiece.model""",
    },
    """tokenizer_file""": {
        """TsinghuaAI/CPM-Generate""": """https://huggingface.co/TsinghuaAI/CPM-Generate/resolve/main/tokenizer.json""",
    },
}

class      UpperCamelCase_					(				__magic_name__     ):





	def __init__(    self						,						A=None						,						A=None						,						A=False						,						A=True						,						A=False						,						A="<s>"						,						A="</s>"						,						A="<unk>"						,						A="<sep>"						,						A="<pad>"						,						A="<cls>"						,						A="<mask>"						,						A=["<eop>", "<eod>"]						,						**A						,						)     ->     Tuple:
								UpperCAmelCase	:						str					=						AddedToken(A						,						lstrip=A						,						rstrip=A       ) if isinstance(A						,						A       ) else mask_token

								super().__init__(
								    vocab_file=A						,						tokenizer_file=A						,						do_lower_case=A						,						remove_space=A						,						keep_accents=A						,						bos_token=A						,						eos_token=A						,						unk_token=A						,						sep_token=A						,						pad_token=A						,						cls_token=A						,						mask_token=A						,						additional_special_tokens=A						,						**A						,						)

								UpperCAmelCase	:						List[Any]					=						3
								UpperCAmelCase	:						Tuple					=						do_lower_case
								UpperCAmelCase	:						Union[str, Any]					=						remove_space
								UpperCAmelCase	:						Tuple					=						keep_accents
								UpperCAmelCase	:						Any					=						vocab_file
								UpperCAmelCase	:						Tuple					=						False if not self.vocab_file else True

								try:
															import jieba
								except ModuleNotFoundError as error:
															raise error.__class__(
															    """You need to install jieba to use CpmTokenizer or CpmTokenizerFast. """
															    """See https://pypi.org/project/jieba/ for installation."""       )
								UpperCAmelCase	:						int					=						jieba
								UpperCAmelCase	:						Optional[int]					=						str.maketrans(""" \n"""						,						"""\u2582\u2583"""       )





	def 					_lowercase(    self						,						A						,						A = None       )     ->     List[int]:
								UpperCAmelCase	:						Optional[int]					=						[self.sep_token_id]
								UpperCAmelCase	:						str					=						[self.cls_token_id]
								if token_ids_a is None:
															return token_ids_a + sep + cls
								return token_ids_a + sep + token_ids_a + sep + cls





	def 					_lowercase(    self						,						A						,						A = None       )     ->     List[int]:
								UpperCAmelCase	:						Union[str, Any]					=						[self.sep_token_id]
								UpperCAmelCase	:						str					=						[2]

								if token_ids_a is None:
															return len(token_ids_a + sep       ) * [0] + cls_segment_id
								return len(token_ids_a + sep       ) * [0] + len(token_ids_a + sep       ) * [1] + cls_segment_id





	def 					_lowercase(    self						,						A						,						A = None       )     ->     Tuple[str]:
								if not self.can_save_slow_tokenizer:
															raise ValueError(
															    """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
															    """tokenizer."""       )

								if not os.path.isdir(A       ):
															logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''       )
															return
								UpperCAmelCase	:						Tuple					=						os.path.join(
								    A						,						(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]       )

								if os.path.abspath(self.vocab_file       ) != os.path.abspath(A       ):
															copyfile(self.vocab_file						,						A       )

								return (out_vocab_file,)





	def 					_lowercase(    self						,						A						,						*A						,						**A       )     ->     List[Any]:
								UpperCAmelCase	:						Optional[Any]					=						[
								    """ """.join([x.translate(self.translator       ) for x in self.jieba.cut(A						,						cut_all=A       )]       )
								    for text in batch_text_or_text_pairs
								]
								return super()._batch_encode_plus(A						,						*A						,						**A       )






	def 					_lowercase(    self						,						*A						,						**A       )     ->     str:
								UpperCAmelCase	:						Optional[int]					=						super()._decode(*A						,						**A       )
								UpperCAmelCase	:						int					=						text.replace(""" """						,						""""""       ).replace("""\u2582"""						,						""" """       ).replace("""\u2583"""						,						"""\n"""       )
								return text

