





'''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_       )      ->Any:
	snake_case__  					=      TimesformerConfig()

	if "large" in model_name:
		snake_case__  					=      96

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

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

	return config




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

	return name




def 	__lowerCamelCase (  UpperCAmelCase_				,							UpperCAmelCase_       )      ->Tuple:
	for key in orig_state_dict.copy().keys():
		snake_case__  					=      orig_state_dict.pop(UpperCAmelCase_       )

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

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

	return orig_state_dict




def 	__lowerCamelCase (  )      ->List[str]:
	snake_case__  					=      hf_hub_download(
	    repo_id='hf-internal-testing/spaghetti-video'				,							filename='eating_spaghetti.npy'				,							repo_type='dataset'       )
	snake_case__  					=      np.load(UpperCAmelCase_       )
	return list(UpperCAmelCase_       )




def 	__lowerCamelCase (  UpperCAmelCase_				,							UpperCAmelCase_				,							UpperCAmelCase_				,							UpperCAmelCase_       )      ->Dict:
	snake_case__  					=      get_timesformer_config(UpperCAmelCase_       )

	snake_case__  					=      TimesformerForVideoClassification(UpperCAmelCase_       )

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

	model.load_state_dict(UpperCAmelCase_       )
	model.eval()

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

	snake_case__  					=      model(**UpperCAmelCase_       )
	snake_case__  					=      outputs.logits

	snake_case__  					=      [
	    # 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__  					=      torch.Size([1, 4_00]       )
		snake_case__  					=      torch.tensor([-0.3016, -0.7713, -0.4205]       )
	elif model_name == "timesformer-base-finetuned-k600":
		snake_case__  					=      torch.Size([1, 6_00]       )
		snake_case__  					=      torch.tensor([-0.7267, -0.7466, 3.2404]       )
	elif model_name == "timesformer-base-finetuned-ssv2":
		snake_case__  					=      torch.Size([1, 1_74]       )
		snake_case__  					=      torch.tensor([-0.9059, 0.6433, -3.1457]       )
	elif model_name == "timesformer-large-finetuned-k400":
		snake_case__  					=      torch.Size([1, 4_00]       )
		snake_case__  					=      torch.tensor([0, 0, 0]       )
	elif model_name == "timesformer-large-finetuned-k600":
		snake_case__  					=      torch.Size([1, 6_00]       )
		snake_case__  					=      torch.tensor([0, 0, 0]       )
	elif model_name == "timesformer-large-finetuned-ssv2":
		snake_case__  					=      torch.Size([1, 1_74]       )
		snake_case__  					=      torch.tensor([0, 0, 0]       )
	elif model_name == "timesformer-hr-finetuned-k400":
		snake_case__  					=      torch.Size([1, 4_00]       )
		snake_case__  					=      torch.tensor([-0.9617, -3.7311, -3.7708]       )
	elif model_name == "timesformer-hr-finetuned-k600":
		snake_case__  					=      torch.Size([1, 6_00]       )
		snake_case__  					=      torch.tensor([2.5273, 0.7127, 1.8848]       )
	elif model_name == "timesformer-hr-finetuned-ssv2":
		snake_case__  					=      torch.Size([1, 1_74]       )
		snake_case__  					=      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]				,							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__":
			a__				: int									=    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__				: Union[str, 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 shutil
import tempfile
import unittest

import pytest

from transformers import WhisperTokenizer, is_speech_available
from transformers.testing_utils import require_sentencepiece, require_torch, require_torchaudio

from .test_feature_extraction_whisper import floats_list


if is_speech_available():
			from transformers import WhisperFeatureExtractor, WhisperProcessor


a__				: Tuple									=    5_0358
a__				: str									=    5_0362






@require_torch
@require_torchaudio
@require_sentencepiece
class 							__snake_case    (				unittest.TestCase       ):





	def      _snake_case      (		self		)			->		Dict:
		snake_case__  					=      'openai/whisper-small.en'
		snake_case__  					=      tempfile.mkdtemp()





	def      _snake_case      (		self			,					**UpperCamelCase_		)			->		Optional[int]:
		return WhisperTokenizer.from_pretrained(self.checkpoint			,					**UpperCamelCase_		)





	def      _snake_case      (		self			,					**UpperCamelCase_		)			->		Any:
		return WhisperFeatureExtractor.from_pretrained(self.checkpoint			,					**UpperCamelCase_		)





	def      _snake_case      (		self		)			->		Tuple:
		shutil.rmtree(self.tmpdirname		)





	def      _snake_case      (		self		)			->		Any:
		snake_case__  					=      self.get_tokenizer()
		snake_case__  					=      self.get_feature_extractor()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)

		processor.save_pretrained(self.tmpdirname		)
		snake_case__  					=      WhisperProcessor.from_pretrained(self.tmpdirname		)

		self.assertEqual(processor.tokenizer.get_vocab()			,					tokenizer.get_vocab()		)
		self.assertIsInstance(processor.tokenizer			,					UpperCamelCase_		)

		self.assertEqual(processor.feature_extractor.to_json_string()			,					feature_extractor.to_json_string()		)
		self.assertIsInstance(processor.feature_extractor			,					UpperCamelCase_		)





	def      _snake_case      (		self		)			->		Dict:
		snake_case__  					=      WhisperProcessor(tokenizer=self.get_tokenizer()			,					feature_extractor=self.get_feature_extractor()		)
		processor.save_pretrained(self.tmpdirname		)

		snake_case__  					=      self.get_tokenizer(bos_token='(BOS)'			,					eos_token='(EOS)'		)
		snake_case__  					=      self.get_feature_extractor(do_normalize=UpperCamelCase_			,					padding_value=1.0		)

		snake_case__  					=      WhisperProcessor.from_pretrained(
		    self.tmpdirname			,					bos_token='(BOS)'			,					eos_token='(EOS)'			,					do_normalize=UpperCamelCase_			,					padding_value=1.0		)

		self.assertEqual(processor.tokenizer.get_vocab()			,					tokenizer_add_kwargs.get_vocab()		)
		self.assertIsInstance(processor.tokenizer			,					UpperCamelCase_		)

		self.assertEqual(processor.feature_extractor.to_json_string()			,					feature_extractor_add_kwargs.to_json_string()		)
		self.assertIsInstance(processor.feature_extractor			,					UpperCamelCase_		)





	def      _snake_case      (		self		)			->		List[str]:
		snake_case__  					=      self.get_feature_extractor()
		snake_case__  					=      self.get_tokenizer()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)

		snake_case__  					=      floats_list((3, 1000)		)

		snake_case__  					=      feature_extractor(UpperCamelCase_			,					return_tensors='np'		)
		snake_case__  					=      processor(UpperCamelCase_			,					return_tensors='np'		)

		for key in input_feat_extract.keys():
			self.assertAlmostEqual(input_feat_extract[key].sum()			,					input_processor[key].sum()			,					delta=1E-2		)





	def      _snake_case      (		self		)			->		Union[str, Any]:
		snake_case__  					=      self.get_feature_extractor()
		snake_case__  					=      self.get_tokenizer()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)

		snake_case__  					=      'This is a test string'

		snake_case__  					=      processor(text=UpperCamelCase_		)

		snake_case__  					=      tokenizer(UpperCamelCase_		)

		for key in encoded_tok.keys():
			self.assertListEqual(encoded_tok[key]			,					encoded_processor[key]		)





	def      _snake_case      (		self		)			->		Tuple:
		snake_case__  					=      self.get_feature_extractor()
		snake_case__  					=      self.get_tokenizer()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)

		snake_case__  					=      [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

		snake_case__  					=      processor.batch_decode(UpperCamelCase_		)
		snake_case__  					=      tokenizer.batch_decode(UpperCamelCase_		)

		self.assertListEqual(UpperCamelCase_			,					UpperCamelCase_		)





	def      _snake_case      (		self		)			->		List[str]:
		snake_case__  					=      self.get_feature_extractor()
		snake_case__  					=      self.get_tokenizer()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)

		self.assertListEqual(
		    processor.model_input_names			,					feature_extractor.model_input_names			,					msg='`processor` and `feature_extractor` model input names do not match'			,					)





	def      _snake_case      (		self		)			->		Dict:
		snake_case__  					=      self.get_feature_extractor()
		snake_case__  					=      self.get_tokenizer()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)
		snake_case__  					=      processor.get_decoder_prompt_ids(task='transcribe'			,					no_timestamps=UpperCamelCase_		)

		self.assertIsInstance(UpperCamelCase_			,					UpperCamelCase_		)
		for ids in forced_decoder_ids:
			self.assertIsInstance(UpperCamelCase_			,					(list, tuple)		)

		snake_case__  					=      [TRANSCRIBE, NOTIMESTAMPS]
		self.assertListEqual([ids[-1] for ids in forced_decoder_ids]			,					UpperCamelCase_		)





	def      _snake_case      (		self		)			->		List[str]:
		snake_case__  					=      WhisperProcessor(tokenizer=self.get_tokenizer()			,					feature_extractor=self.get_feature_extractor()		)
		snake_case__  					=      processor.get_prompt_ids('Mr. Quilter'		)
		snake_case__  					=      processor.tokenizer.decode(UpperCamelCase_		)

		self.assertListEqual(prompt_ids.tolist()			,					[5_0360, 1770, 13, 2264, 346, 353]		)
		self.assertEqual(UpperCamelCase_			,					'<|startofprev|> Mr. Quilter'		)





	def      _snake_case      (		self		)			->		Dict:
		snake_case__  					=      WhisperProcessor(tokenizer=self.get_tokenizer()			,					feature_extractor=self.get_feature_extractor()		)
		snake_case__  					=      processor.get_prompt_ids(''		)
		snake_case__  					=      processor.tokenizer.decode(UpperCamelCase_		)

		self.assertListEqual(prompt_ids.tolist()			,					[5_0360, 220]		)
		self.assertEqual(UpperCamelCase_			,					'<|startofprev|> '		)





	def      _snake_case      (		self		)			->		int:
		snake_case__  					=      WhisperProcessor(tokenizer=self.get_tokenizer()			,					feature_extractor=self.get_feature_extractor()		)

		def _test_prompt_error_raised_helper(UpperCamelCase_			,					UpperCamelCase_		):
			with pytest.raises(UpperCamelCase_		) as excinfo:
				processor.get_prompt_ids(UpperCamelCase_		)
			snake_case__  					=      F'''Encountered text in the prompt corresponding to disallowed special token: {special_token}.'''
			self.assertEqual(UpperCamelCase_			,					str(excinfo.value		)		)

		_test_prompt_error_raised_helper('<|startofprev|> test'			,					'<|startofprev|>'		)
		_test_prompt_error_raised_helper('test <|notimestamps|>'			,					'<|notimestamps|>'		)
		_test_prompt_error_raised_helper('test <|zh|> test <|transcribe|>'			,					'<|zh|>'		)



