



"""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_		):
						UpperCAmelCase    =	TimesformerConfig()

						if "large" in model_name:
												UpperCAmelCase    =	96

						if "hr" in model_name:
												UpperCAmelCase    =	16
												UpperCAmelCase    =	448

						UpperCAmelCase    =	'huggingface/label-files'
						if "k400" in model_name:
												UpperCAmelCase    =	400
												UpperCAmelCase    =	'kinetics400-id2label.json'
						elif "k600" in model_name:
												UpperCAmelCase    =	600
												UpperCAmelCase    =	'kinetics600-id2label.json'
						elif "ssv2" in model_name:
												UpperCAmelCase    =	174
												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(UpperCamelCase_	,			UpperCamelCase_	,			repo_type='dataset'		)	,			'r'		)		)
						UpperCAmelCase    =	{int(UpperCamelCase_		): v for k, v in idalabel.items()}
						UpperCAmelCase    =	idalabel
						UpperCAmelCase    =	{v: k for k, v in idalabel.items()}

						return config

def 				_lowerCAmelCase   (     lowercase_		):






						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 				_lowerCAmelCase   (     lowercase_	,			lowercase_		):






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

												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 				_lowerCAmelCase   (     ):
						UpperCAmelCase    =	hf_hub_download(
						    repo_id='hf-internal-testing/spaghetti-video'	,			filename='eating_spaghetti.npy'	,			repo_type='dataset'		)
						UpperCAmelCase    =	np.load(UpperCamelCase_		)
						return list(UpperCamelCase_		)

def 				_lowerCAmelCase   (     lowercase_	,			lowercase_	,			lowercase_	,			lowercase_		):
						UpperCAmelCase    =	get_timesformer_config(UpperCamelCase_		)

						UpperCAmelCase    =	TimesformerForVideoClassification(UpperCamelCase_		)

						# download original checkpoint, hosted on Google Drive
						UpperCAmelCase    =	'pytorch_model.bin'
						gdown.cached_download(UpperCamelCase_	,			UpperCamelCase_	,			quiet=UpperCamelCase_		)
						UpperCAmelCase    =	torch.load(UpperCamelCase_	,			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(UpperCamelCase_	,			UpperCamelCase_		)

						model.load_state_dict(UpperCamelCase_		)
						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(**UpperCamelCase_		)
						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, 400]		)
												UpperCAmelCase    =	torch.tensor([-0.3_0_1_6, -0.7_7_1_3, -0.4_2_0_5]		)
						elif model_name == "timesformer-base-finetuned-k600":
												UpperCAmelCase    =	torch.Size([1, 600]		)
												UpperCAmelCase    =	torch.tensor([-0.7_2_6_7, -0.7_4_6_6, 3.2_4_0_4]		)
						elif model_name == "timesformer-base-finetuned-ssv2":
												UpperCAmelCase    =	torch.Size([1, 174]		)
												UpperCAmelCase    =	torch.tensor([-0.9_0_5_9, 0.6_4_3_3, -3.1_4_5_7]		)
						elif model_name == "timesformer-large-finetuned-k400":
												UpperCAmelCase    =	torch.Size([1, 400]		)
												UpperCAmelCase    =	torch.tensor([0, 0, 0]		)
						elif model_name == "timesformer-large-finetuned-k600":
												UpperCAmelCase    =	torch.Size([1, 600]		)
												UpperCAmelCase    =	torch.tensor([0, 0, 0]		)
						elif model_name == "timesformer-large-finetuned-ssv2":
												UpperCAmelCase    =	torch.Size([1, 174]		)
												UpperCAmelCase    =	torch.tensor([0, 0, 0]		)
						elif model_name == "timesformer-hr-finetuned-k400":
												UpperCAmelCase    =	torch.Size([1, 400]		)
												UpperCAmelCase    =	torch.tensor([-0.9_6_1_7, -3.7_3_1_1, -3.7_7_0_8]		)
						elif model_name == "timesformer-hr-finetuned-k600":
												UpperCAmelCase    =	torch.Size([1, 600]		)
												UpperCAmelCase    =	torch.tensor([2.5_2_7_3, 0.7_1_2_7, 1.8_8_4_8]		)
						elif model_name == "timesformer-hr-finetuned-ssv2":
												UpperCAmelCase    =	torch.Size([1, 174]		)
												UpperCAmelCase    =	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__":
				snake_case_         = 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."""
				)

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










"""simple docstring"""


import warnings

from .generation import FlaxGenerationMixin







class    A_    (  SCREAMING_SNAKE_CASE_       ):



						"""simple docstring"""
						# warning at import time
						warnings.warn(
						    """Importing `FlaxGenerationMixin` from `src/transformers/generation_flax_utils.py` is deprecated and will """
						    """be removed in Transformers v5. Import as `from transformers import FlaxGenerationMixin` instead."""			,     SCREAMING_SNAKE_CASE_			,     )





