


import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)





def   _lowerCAmelCase				(   UpperCamelCase__:     Union[str, Any]						) ->   str:





	"""simple docstring"""







	A								=				VideoMAEConfig()

	set_architecture_configs(UpperCamelCase__							,					UpperCamelCase__						)

	if "finetuned" not in model_name:
		A								=				False

	if "finetuned" in model_name:
		A								=				"""huggingface/label-files"""
		if "kinetics" in model_name:
			A								=				4_00
			A								=				"""kinetics400-id2label.json"""
		elif "ssv2" in model_name:
			A								=				1_74
			A								=				"""something-something-v2-id2label.json"""
		else:
			raise ValueError("""Model name should either contain 'kinetics' or 'ssv2' in case it's fine-tuned."""						)
		A								=				json.load(open(hf_hub_download(UpperCamelCase__							,					UpperCamelCase__							,					repo_type="""dataset"""						)							,					"""r"""						)						)
		A								=				{int(UpperCamelCase__						): v for k, v in idalabel.items()}
		A								=				idalabel
		A								=				{v: k for k, v in idalabel.items()}

	return config





def   _lowerCAmelCase				(   UpperCamelCase__:     Tuple							,					UpperCamelCase__:     List[str]						) ->   Dict:





	"""simple docstring"""







	if "small" in model_name:
		A								=				3_84
		A								=				15_36
		A								=				12
		A								=				16
		A								=				12
		A								=				3
		A								=				1_92
		A								=				7_68
	elif "large" in model_name:
		A								=				10_24
		A								=				40_96
		A								=				24
		A								=				16
		A								=				12
		A								=				8
		A								=				5_12
		A								=				20_48
	elif "huge" in model_name:
		A								=				12_80
		A								=				51_20
		A								=				32
		A								=				16
		A								=				12
		A								=				8
		A								=				6_40
		A								=				25_60
	elif "base" not in model_name:
		raise ValueError("""Model name should include either \"small\", \"base\", \"large\", or \"huge\""""						)





def   _lowerCAmelCase				(   UpperCamelCase__:     int						) ->   List[Any]:





	"""simple docstring"""







	if "encoder." in name:
		A								=				name.replace("""encoder."""							,					""""""						)
	if "cls_token" in name:
		A								=				name.replace("""cls_token"""							,					"""videomae.embeddings.cls_token"""						)
	if "decoder_pos_embed" in name:
		A								=				name.replace("""decoder_pos_embed"""							,					"""decoder.decoder_pos_embed"""						)
	if "pos_embed" in name and "decoder" not in name:
		A								=				name.replace("""pos_embed"""							,					"""videomae.embeddings.position_embeddings"""						)
	if "patch_embed.proj" in name:
		A								=				name.replace("""patch_embed.proj"""							,					"""videomae.embeddings.patch_embeddings.projection"""						)
	if "patch_embed.norm" in name:
		A								=				name.replace("""patch_embed.norm"""							,					"""videomae.embeddings.norm"""						)
	if "decoder.blocks" in name:
		A								=				name.replace("""decoder.blocks"""							,					"""decoder.decoder_layers"""						)
	if "blocks" in name:
		A								=				name.replace("""blocks"""							,					"""videomae.encoder.layer"""						)
	if "attn.proj" in name:
		A								=				name.replace("""attn.proj"""							,					"""attention.output.dense"""						)
	if "attn" in name and "bias" not in name:
		A								=				name.replace("""attn"""							,					"""attention.self"""						)
	if "attn" in name:
		A								=				name.replace("""attn"""							,					"""attention.attention"""						)
	if "norm1" in name:
		A								=				name.replace("""norm1"""							,					"""layernorm_before"""						)
	if "norm2" in name:
		A								=				name.replace("""norm2"""							,					"""layernorm_after"""						)
	if "mlp.fc1" in name:
		A								=				name.replace("""mlp.fc1"""							,					"""intermediate.dense"""						)
	if "mlp.fc2" in name:
		A								=				name.replace("""mlp.fc2"""							,					"""output.dense"""						)
	if "decoder_embed" in name:
		A								=				name.replace("""decoder_embed"""							,					"""decoder.decoder_embed"""						)
	if "decoder_norm" in name:
		A								=				name.replace("""decoder_norm"""							,					"""decoder.decoder_norm"""						)
	if "decoder_pred" in name:
		A								=				name.replace("""decoder_pred"""							,					"""decoder.decoder_pred"""						)
	if "norm.weight" in name and "decoder" not in name and "fc" not in name:
		A								=				name.replace("""norm.weight"""							,					"""videomae.layernorm.weight"""						)
	if "norm.bias" in name and "decoder" not in name and "fc" not in name:
		A								=				name.replace("""norm.bias"""							,					"""videomae.layernorm.bias"""						)
	if "head" in name and "decoder" not in name:
		A								=				name.replace("""head"""							,					"""classifier"""						)

	return name





def   _lowerCAmelCase				(   UpperCamelCase__:     List[Any]							,					UpperCamelCase__:     int						) ->   Optional[int]:





	"""simple docstring"""







	for key in orig_state_dict.copy().keys():
		A								=				orig_state_dict.pop(UpperCamelCase__						)

		if key.startswith("""encoder."""						):
			A								=				key.replace("""encoder."""							,					""""""						)

		if "qkv" in key:
			A								=				key.split("""."""						)
			if key.startswith("""decoder.blocks"""						):
				A								=				config.decoder_hidden_size
				A								=				int(key_split[2]						)
				A								=				"""decoder.decoder_layers."""
				if "weight" in key:
					A								=				val[:dim, :]
					A								=				val[dim : dim * 2, :]
					A								=				val[-dim:, :]
			else:
				A								=				config.hidden_size
				A								=				int(key_split[1]						)
				A								=				"""videomae.encoder.layer."""
				if "weight" in key:
					A								=				val[:dim, :]
					A								=				val[dim : dim * 2, :]
					A								=				val[-dim:, :]
		else:
			A								=				val

	return orig_state_dict





def   _lowerCAmelCase				(   ) ->   Dict:





	"""simple docstring"""







	A								=				hf_hub_download(
	    repo_id="""hf-internal-testing/spaghetti-video"""							,					filename="""eating_spaghetti.npy"""							,					repo_type="""dataset"""						)
	A								=				np.load(UpperCamelCase__						)
	return list(UpperCamelCase__						)





def   _lowerCAmelCase				(   UpperCamelCase__:     str							,					UpperCamelCase__:     Any							,					UpperCamelCase__:     Union[str, Any]							,					UpperCamelCase__:     str						) ->   Optional[int]:





	"""simple docstring"""







	A								=				get_videomae_config(UpperCamelCase__						)

	if "finetuned" in model_name:
		A								=				VideoMAEForVideoClassification(UpperCamelCase__						)
	else:
		A								=				VideoMAEForPreTraining(UpperCamelCase__						)

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

	model.load_state_dict(UpperCamelCase__						)
	model.eval()

	# verify model on basic input
	A								=				VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]							,					image_std=[0.5, 0.5, 0.5]						)
	A								=				prepare_video()
	A								=				image_processor(UpperCamelCase__							,					return_tensors="""pt"""						)

	if "finetuned" not in model_name:
		A								=				hf_hub_download(repo_id="""hf-internal-testing/bool-masked-pos"""							,					filename="""bool_masked_pos.pt"""						)
		A								=				torch.load(UpperCamelCase__						)

	A								=				model(**UpperCamelCase__						)
	A								=				outputs.logits

	A								=				[
	    """videomae-small-finetuned-kinetics""",
	    """videomae-small-finetuned-ssv2""",
	    # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
	    """videomae-base-short""",
	    """videomae-base-short-finetuned-kinetics""",
	    """videomae-base""",
	    """videomae-base-finetuned-kinetics""",
	    """videomae-large""",
	    """videomae-large-finetuned-kinetics""",
	    """videomae-huge-finetuned-kinetics""",
	    # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
	    """videomae-base-short-ssv2""",
	    """videomae-base-short-finetuned-ssv2""",
	    """videomae-base-ssv2""",
	    """videomae-base-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 == "videomae-small-finetuned-kinetics":
		A								=				torch.Size([1, 4_00]						)
		A								=				torch.tensor([-0.92_91, -0.40_61, -0.93_07]						)
	elif model_name == "videomae-small-finetuned-ssv2":
		A								=				torch.Size([1, 1_74]						)
		A								=				torch.tensor([0.26_71, -0.46_89, -0.82_35]						)
	elif model_name == "videomae-base":
		A								=				torch.Size([1, 14_08, 15_36]						)
		A								=				torch.tensor([[0.77_39, 0.79_68, 0.70_89], [0.67_01, 0.74_87, 0.62_09], [0.42_87, 0.51_58, 0.47_73]]						)
	elif model_name == "videomae-base-short":
		A								=				torch.Size([1, 14_08, 15_36]						)
		A								=				torch.tensor([[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]]						)
		# we verified the loss both for normalized and unnormalized targets for this one
		A								=				torch.tensor([0.51_42]						) if config.norm_pix_loss else torch.tensor([0.64_69]						)
	elif model_name == "videomae-large":
		A								=				torch.Size([1, 14_08, 15_36]						)
		A								=				torch.tensor([[0.71_49, 0.79_97, 0.69_66], [0.67_68, 0.78_69, 0.69_48], [0.51_39, 0.62_21, 0.56_05]]						)
	elif model_name == "videomae-large-finetuned-kinetics":
		A								=				torch.Size([1, 4_00]						)
		A								=				torch.tensor([0.07_71, 0.00_11, -0.36_25]						)
	elif model_name == "videomae-huge-finetuned-kinetics":
		A								=				torch.Size([1, 4_00]						)
		A								=				torch.tensor([0.24_33, 0.16_32, -0.48_94]						)
	elif model_name == "videomae-base-short-finetuned-kinetics":
		A								=				torch.Size([1, 4_00]						)
		A								=				torch.tensor([0.65_88, 0.09_90, -0.24_93]						)
	elif model_name == "videomae-base-finetuned-kinetics":
		A								=				torch.Size([1, 4_00]						)
		A								=				torch.tensor([0.36_69, -0.06_88, -0.24_21]						)
	elif model_name == "videomae-base-short-ssv2":
		A								=				torch.Size([1, 14_08, 15_36]						)
		A								=				torch.tensor([[0.47_12, 0.52_96, 0.57_86], [0.22_78, 0.27_29, 0.40_26], [0.03_52, 0.07_30, 0.25_06]]						)
	elif model_name == "videomae-base-short-finetuned-ssv2":
		A								=				torch.Size([1, 1_74]						)
		A								=				torch.tensor([-0.05_37, -0.15_39, -0.32_66]						)
	elif model_name == "videomae-base-ssv2":
		A								=				torch.Size([1, 14_08, 15_36]						)
		A								=				torch.tensor([[0.81_31, 0.87_27, 0.85_46], [0.73_66, 0.93_77, 0.88_70], [0.59_35, 0.88_74, 0.85_64]]						)
	elif model_name == "videomae-base-finetuned-ssv2":
		A								=				torch.Size([1, 1_74]						)
		A								=				torch.tensor([0.19_61, -0.83_37, -0.63_89]						)
	else:
		raise ValueError(f'Model name not supported. Should be one of {model_names}'						)

	# verify logits
	assert logits.shape == expected_shape
	if "finetuned" in model_name:
		assert torch.allclose(logits[0, :3]							,					UpperCamelCase__							,					atol=1e-4						)
	else:
		print("""Logits:"""							,					logits[0, :3, :3]						)
		assert torch.allclose(logits[0, :3, :3]							,					UpperCamelCase__							,					atol=1e-4						)
	print("""Logits ok!"""						)

	# verify loss, if applicable
	if model_name == "videomae-base-short":
		A								=				outputs.loss
		assert torch.allclose(UpperCamelCase__							,					UpperCamelCase__							,					atol=1e-4						)
		print("""Loss 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(UpperCamelCase__							,					organization="""nielsr"""						)


if __name__ == "__main__":
			_lowercase       :      List[Any]	=	argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    "--checkpoint_url",
			    default="https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4",
			    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="/Users/nielsrogge/Documents/VideoMAE/Test",
			    type=str,
			    help="Path to the output PyTorch model directory.",
			)
			parser.add_argument("--model_name", default="videomae-base", 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."
			)

			_lowercase       :      str	=	parser.parse_args()
			convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)







import sys
from collections import defaultdict






class       _UpperCamelCase    :

				"""simple docstring"""







				def __init__(						self   )	->				Any:
					A								=				[]

				def      _UpperCAmelCase					(						self				,   a__   )	->				List[str]:
					return self.node_position[vertex]

				def      _UpperCAmelCase					(						self				,   a__				,   a__   )	->				Optional[int]:
					A								=				pos

				def      _UpperCAmelCase					(						self				,   a__				,   a__				,   a__				,   a__   )	->				Any:
					if start > size // 2 - 1:
						return
					else:
						if 2 * start + 2 >= size:
							A								=				2 * start + 1
						else:
							if heap[2 * start + 1] < heap[2 * start + 2]:
								A								=				2 * start + 1
							else:
								A								=				2 * start + 2
						if heap[smallest_child] < heap[start]:
							A					,			A								=				heap[smallest_child], positions[smallest_child]
							A					,			A								=				(
							    heap[start],
							    positions[start],
							)
							A					,			A								=				temp, tempa

							A								=				self.get_position(positions[smallest_child]   )
							self.set_position(
							    positions[smallest_child]				,   self.get_position(positions[start]   )   )
							self.set_position(positions[start]				,   a__   )

							self.top_to_bottom(a__				,   a__				,   a__				,   a__   )

				def      _UpperCAmelCase					(						self				,   a__				,   a__				,   a__				,   a__   )	->				Dict:
					A								=				position[index]

					while index != 0:
						A								=				int((index - 2) / 2   ) if index % 2 == 0 else int((index - 1) / 2   )

						if val < heap[parent]:
							A								=				heap[parent]
							A								=				position[parent]
							self.set_position(position[parent]				,   a__   )
						else:
							A								=				val
							A								=				temp
							self.set_position(a__				,   a__   )
							break
						A								=				parent
					else:
						A								=				val
						A								=				temp
						self.set_position(a__				,   0   )

				def      _UpperCAmelCase					(						self				,   a__				,   a__   )	->				Optional[int]:
					A								=				len(a__   ) // 2 - 1
					for i in range(a__				,   -1				,   -1   ):
						self.top_to_bottom(a__				,   a__				,   len(a__   )				,   a__   )

				def      _UpperCAmelCase					(						self				,   a__				,   a__   )	->				Any:
					A								=				positions[0]
					A								=				sys.maxsize
					self.top_to_bottom(a__				,   0				,   len(a__   )				,   a__   )
					return temp





def   _lowerCAmelCase				(   UpperCamelCase__:     List[str]						) ->   List[str]:





	"""simple docstring"""







	A								=				Heap()

	A								=				[0] * len(UpperCamelCase__						)
	A								=				[-1] * len(UpperCamelCase__						)  # Neighboring Tree Vertex of selected vertex
	# Minimum Distance of explored vertex with neighboring vertex of partial tree
	# formed in graph
	A								=				[]  # Heap of Distance of vertices from their neighboring vertex
	A								=				[]

	for vertex in range(len(UpperCamelCase__						)						):
		distance_tv.append(sys.maxsize						)
		positions.append(UpperCamelCase__						)
		heap.node_position.append(UpperCamelCase__						)

	A								=				[]
	A								=				1
	A								=				sys.maxsize
	for neighbor, distance in adjacency_list[0]:
		A								=				0
		A								=				distance
	heap.heapify(UpperCamelCase__							,					UpperCamelCase__						)

	for _ in range(1							,					len(UpperCamelCase__						)						):
		A								=				heap.delete_minimum(UpperCamelCase__							,					UpperCamelCase__						)
		if visited[vertex] == 0:
			tree_edges.append((nbr_tv[vertex], vertex)						)
			A								=				1
			for neighbor, distance in adjacency_list[vertex]:
				if (
				    visited[neighbor] == 0
				    and distance < distance_tv[heap.get_position(UpperCamelCase__						)]
				):
					A								=				distance
					heap.bottom_to_top(
					    UpperCamelCase__							,					heap.get_position(UpperCamelCase__						)							,					UpperCamelCase__							,					UpperCamelCase__						)
					A								=				vertex
	return tree_edges


if __name__ == "__main__":  # pragma: no cover
				# < --------- Prims Algorithm --------- >
				_lowercase       :      Tuple	=	int(input("Enter number of edges: ").strip())
				_lowercase       :      int	=	defaultdict(list)
				for _ in range(edges_number):
								_lowercase       :      Optional[Any]	=	[int(x) for x in input().strip().split()]
								adjacency_list[edge[0]].append([edge[1], edge[2]])
								adjacency_list[edge[1]].append([edge[0], edge[2]])
				print(prisms_algorithm(adjacency_list))



