

"""simple docstring"""





import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
_SCREAMING_SNAKE_CASE							:	Union[str, Any]										=			logging.get_logger(__name__)



def 							_lowerCAmelCase   (    UpperCAmelCase  :				Any					):


					'''simple docstring'''






					UpperCamelCase__		:     str								=ConvNextConfig()

					if "tiny" in checkpoint_url:
										UpperCamelCase__		:     int								=[3, 3, 9, 3]
										UpperCamelCase__		:     List[Any]								=[96, 192, 384, 768]
					if "small" in checkpoint_url:
										UpperCamelCase__		:     Any								=[3, 3, 27, 3]
										UpperCamelCase__		:     List[Any]								=[96, 192, 384, 768]
					if "base" in checkpoint_url:
										UpperCamelCase__		:     Tuple								=[3, 3, 27, 3]
										UpperCamelCase__		:     Any								=[128, 256, 512, 1_024]
					if "large" in checkpoint_url:
										UpperCamelCase__		:     List[str]								=[3, 3, 27, 3]
										UpperCamelCase__		:     Optional[Any]								=[192, 384, 768, 1_536]
					if "xlarge" in checkpoint_url:
										UpperCamelCase__		:     int								=[3, 3, 27, 3]
										UpperCamelCase__		:     str								=[256, 512, 1_024, 2_048]

					if "1k" in checkpoint_url:
										UpperCamelCase__		:     Optional[Any]								=1_000
										UpperCamelCase__		:     List[Any]								="imagenet-1k-id2label.json"
										UpperCamelCase__		:     Dict								=(1, 1_000)
					else:
										UpperCamelCase__		:     int								=21_841
										UpperCamelCase__		:     int								="imagenet-22k-id2label.json"
										UpperCamelCase__		:     Optional[int]								=(1, 21_841)

					UpperCamelCase__		:     Dict								="huggingface/label-files"
					UpperCamelCase__		:     str								=num_labels
					UpperCamelCase__		:     Any								=json.load(open(hf_hub_download(A_   ,				A_   ,				repo_type='''dataset'''					)   ,				'''r'''					)					)
					UpperCamelCase__		:     Any								={int(A_					): v for k, v in idalabel.items()}
					if "1k" not in checkpoint_url:
										# this dataset contains 21843 labels but the model only has 21841
										# we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
										del idalabel[9_205]
										del idalabel[15_027]
					UpperCamelCase__		:     Tuple								=idalabel
					UpperCamelCase__		:     Optional[int]								={v: k for k, v in idalabel.items()}
					UpperCamelCase__		:     Union[str, Any]								=hidden_sizes
					UpperCamelCase__		:     Tuple								=depths

					return config, expected_shape



def 							_lowerCAmelCase   (    UpperCAmelCase  :				Optional[int]					):


					'''simple docstring'''






					if "downsample_layers.0.0" in name:
										UpperCamelCase__		:     int								=name.replace('''downsample_layers.0.0'''   ,				'''embeddings.patch_embeddings'''					)
					if "downsample_layers.0.1" in name:
										UpperCamelCase__		:     List[Any]								=name.replace('''downsample_layers.0.1'''   ,				'''embeddings.norm'''					)  # we rename to layernorm later on
					if "downsample_layers.1.0" in name:
										UpperCamelCase__		:     Optional[int]								=name.replace('''downsample_layers.1.0'''   ,				'''stages.1.downsampling_layer.0'''					)
					if "downsample_layers.1.1" in name:
										UpperCamelCase__		:     Dict								=name.replace('''downsample_layers.1.1'''   ,				'''stages.1.downsampling_layer.1'''					)
					if "downsample_layers.2.0" in name:
										UpperCamelCase__		:     Any								=name.replace('''downsample_layers.2.0'''   ,				'''stages.2.downsampling_layer.0'''					)
					if "downsample_layers.2.1" in name:
										UpperCamelCase__		:     Any								=name.replace('''downsample_layers.2.1'''   ,				'''stages.2.downsampling_layer.1'''					)
					if "downsample_layers.3.0" in name:
										UpperCamelCase__		:     Any								=name.replace('''downsample_layers.3.0'''   ,				'''stages.3.downsampling_layer.0'''					)
					if "downsample_layers.3.1" in name:
										UpperCamelCase__		:     Union[str, Any]								=name.replace('''downsample_layers.3.1'''   ,				'''stages.3.downsampling_layer.1'''					)
					if "stages" in name and "downsampling_layer" not in name:
										# stages.0.0. for instance should be renamed to stages.0.layers.0.
										UpperCamelCase__		:     str								=name[: len('''stages.0'''					)] + ".layers" + name[len('''stages.0'''					) :]
					if "stages" in name:
										UpperCamelCase__		:     List[str]								=name.replace('''stages'''   ,				'''encoder.stages'''					)
					if "norm" in name:
										UpperCamelCase__		:     str								=name.replace('''norm'''   ,				'''layernorm'''					)
					if "gamma" in name:
										UpperCamelCase__		:     List[str]								=name.replace('''gamma'''   ,				'''layer_scale_parameter'''					)
					if "head" in name:
										UpperCamelCase__		:     str								=name.replace('''head'''   ,				'''classifier'''					)

					return name



def 							_lowerCAmelCase   (    ):


					'''simple docstring'''






					UpperCamelCase__		:     List[str]								="http://images.cocodataset.org/val2017/000000039769.jpg"
					UpperCamelCase__		:     Dict								=Image.open(requests.get(A_   ,				stream=A_					).raw					)
					return im



@torch.no_grad()
def 							_lowerCAmelCase   (    UpperCAmelCase  :				List[Any]   ,				UpperCAmelCase  :				Dict					):


					'''simple docstring'''






					UpperCamelCase__		:     List[Any]								=get_convnext_config(A_					)
					# load original state_dict from URL
					UpperCamelCase__		:     Union[str, Any]								=torch.hub.load_state_dict_from_url(A_					)["model"]
					# rename keys
					for key in state_dict.copy().keys():
										UpperCamelCase__		:     Dict								=state_dict.pop(A_					)
										UpperCamelCase__		:     int								=val
					# add prefix to all keys expect classifier head
					for key in state_dict.copy().keys():
										UpperCamelCase__		:     Dict								=state_dict.pop(A_					)
										if not key.startswith('''classifier'''					):
															UpperCamelCase__		:     Union[str, Any]								="convnext." + key
										UpperCamelCase__		:     List[str]								=val

					# load HuggingFace model
					UpperCamelCase__		:     Optional[Any]								=ConvNextForImageClassification(A_					)
					model.load_state_dict(A_					)
					model.eval()

					# Check outputs on an image, prepared by ConvNextImageProcessor
					UpperCamelCase__		:     Tuple								=224 if "224" in checkpoint_url else 384
					UpperCamelCase__		:     Optional[int]								=ConvNextImageProcessor(size=A_					)
					UpperCamelCase__		:     List[Any]								=image_processor(images=prepare_img()   ,				return_tensors='''pt'''					).pixel_values

					UpperCamelCase__		:     str								=model(A_					).logits

					# note: the logits below were obtained without center cropping
					if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
										UpperCamelCase__		:     int								=torch.tensor([-0.1210, -0.6605, 0.1918]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
										UpperCamelCase__		:     Optional[Any]								=torch.tensor([-0.4473, -0.1847, -0.6365]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
										UpperCamelCase__		:     Any								=torch.tensor([0.4525, 0.7539, 0.0308]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
										UpperCamelCase__		:     Dict								=torch.tensor([0.3561, 0.6350, -0.0384]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
										UpperCamelCase__		:     List[Any]								=torch.tensor([0.4174, -0.0989, 0.1489]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
										UpperCamelCase__		:     Any								=torch.tensor([0.2513, -0.1349, -0.1613]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
										UpperCamelCase__		:     int								=torch.tensor([1.2980, 0.3631, -0.1198]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
										UpperCamelCase__		:     Dict								=torch.tensor([1.2963, 0.1227, 0.1723]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
										UpperCamelCase__		:     Union[str, Any]								=torch.tensor([1.7956, 0.8390, 0.2820]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
										UpperCamelCase__		:     List[str]								=torch.tensor([-0.2822, -0.0502, -0.0878]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
										UpperCamelCase__		:     Optional[Any]								=torch.tensor([-0.5672, -0.0730, -0.4348]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
										UpperCamelCase__		:     str								=torch.tensor([0.2681, 0.2365, 0.6246]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
										UpperCamelCase__		:     Optional[Any]								=torch.tensor([-0.2642, 0.3931, 0.5116]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
										UpperCamelCase__		:     Tuple								=torch.tensor([-0.6677, -0.1873, -0.8379]					)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
										UpperCamelCase__		:     int								=torch.tensor([-0.7749, -0.2967, -0.6444]					)
					else:
										raise ValueError(F'''Unknown URL: {checkpoint_url}'''					)

					assert torch.allclose(logits[0, :3]   ,				A_   ,				atol=1E-3					)
					assert logits.shape == expected_shape

					Path(A_					).mkdir(exist_ok=A_					)
					print(F'''Saving model to {pytorch_dump_folder_path}'''					)
					model.save_pretrained(A_					)
					print(F'''Saving image processor to {pytorch_dump_folder_path}'''					)
					image_processor.save_pretrained(A_					)

					print('''Pushing model to the hub...'''					)
					UpperCamelCase__		:     str								="convnext"
					if "tiny" in checkpoint_url:
										model_name += "-tiny"
					elif "small" in checkpoint_url:
										model_name += "-small"
					elif "base" in checkpoint_url:
										model_name += "-base"
					elif "xlarge" in checkpoint_url:
										model_name += "-xlarge"
					elif "large" in checkpoint_url:
										model_name += "-large"
					if "224" in checkpoint_url:
										model_name += "-224"
					elif "384" in checkpoint_url:
										model_name += "-384"
					if "22k" in checkpoint_url and "1k" not in checkpoint_url:
										model_name += "-22k"
					if "22k" in checkpoint_url and "1k" in checkpoint_url:
										model_name += "-22k-1k"

					model.push_to_hub(
					    repo_path_or_name=Path(A_   ,				A_					)   ,				organization='''nielsr'''   ,				commit_message='''Add model'''   ,				)


if __name__ == "__main__":
				_SCREAMING_SNAKE_CASE							:	Optional[int]										=			argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    """--checkpoint_url""",
				    default="""https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth""",
				    type=str,
				    help="""URL of the original ConvNeXT checkpoint you'd like to convert.""",
				)
				parser.add_argument(
				    """--pytorch_dump_folder_path""",
				    default=None,
				    type=str,
				    required=True,
				    help="""Path to the output PyTorch model directory.""",
				)

				_SCREAMING_SNAKE_CASE							:	List[str]										=			parser.parse_args()
				convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






"""simple docstring"""







def 	lowercase     (							A_    ,  A_			)-> list[tuple[int, int, int]]:
  '''simple docstring'''







  a		:						Union[str, Any] 				=     sorted(A_    ,  key=lambda A_			: edge[2]			)

  a		:						List[Any] 				=     list(range(A_			)			)

  def find_parent(A_			):
    if i != parent[i]:
      a		:						Dict 				=     find_parent(parent[i]			)
    return parent[i]

  a		:						Optional[Any] 				=     0
  a		:						Optional[int] 				=     []

  for edge in edges:
    a		:						List[str] 				=     find_parent(edge[0]			)
    a		:						int 				=     find_parent(edge[1]			)
    if parent_a != parent_b:
      minimum_spanning_tree_cost += edge[2]
      minimum_spanning_tree.append(A_			)
      a		:						Tuple 				=     parent_b

  return minimum_spanning_tree


if __name__ == "__main__":  # pragma: no cover
 __lowercase   ,				__lowercase          = list(map(int, input().strip().split()))
 __lowercase          = []

 for _ in range(num_edges):
  __lowercase   ,				__lowercase   ,				__lowercase          = (int(x) for x in input().strip().split())
  edges.append((nodea, nodea, cost))

 kruskal(num_nodes, edges)






