
import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD






def 							UpperCAmelCase_		(							):
	lowercase       =    'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
	lowercase       =    Image.open(requests.get(__SCREAMING_SNAKE_CASE ,							stream=__SCREAMING_SNAKE_CASE							).raw							).convert('RGB'							)

	return image






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	lowercase       =    []
	# fmt: off

	# vision encoder
	rename_keys.append(('visual_encoder.cls_token', 'vision_model.embeddings.class_embedding')							)
	rename_keys.append(('visual_encoder.pos_embed', 'vision_model.embeddings.position_embedding')							)
	rename_keys.append(('visual_encoder.patch_embed.proj.weight', 'vision_model.embeddings.patch_embedding.weight')							)
	rename_keys.append(('visual_encoder.patch_embed.proj.bias', 'vision_model.embeddings.patch_embedding.bias')							)
	rename_keys.append(('ln_vision.weight', 'vision_model.post_layernorm.weight')							)
	rename_keys.append(('ln_vision.bias', 'vision_model.post_layernorm.bias')							)

	for i in range(config.vision_config.num_hidden_layers							):
		rename_keys.append((F'''visual_encoder.blocks.{i}.norm1.weight''', F'''vision_model.encoder.layers.{i}.layer_norm1.weight''')							)
		rename_keys.append((F'''visual_encoder.blocks.{i}.norm1.bias''', F'''vision_model.encoder.layers.{i}.layer_norm1.bias''')							)
		rename_keys.append((F'''visual_encoder.blocks.{i}.norm2.weight''', F'''vision_model.encoder.layers.{i}.layer_norm2.weight''')							)
		rename_keys.append((F'''visual_encoder.blocks.{i}.norm2.bias''', F'''vision_model.encoder.layers.{i}.layer_norm2.bias''')							)
		rename_keys.append((F'''visual_encoder.blocks.{i}.attn.qkv.weight''', F'''vision_model.encoder.layers.{i}.self_attn.qkv.weight''')							)
		rename_keys.append((F'''visual_encoder.blocks.{i}.attn.proj.weight''', F'''vision_model.encoder.layers.{i}.self_attn.projection.weight''',)							)
		rename_keys.append((F'''visual_encoder.blocks.{i}.attn.proj.bias''', F'''vision_model.encoder.layers.{i}.self_attn.projection.bias''')							)
		rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc1.weight''', F'''vision_model.encoder.layers.{i}.mlp.fc1.weight''')							)
		rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc1.bias''', F'''vision_model.encoder.layers.{i}.mlp.fc1.bias''')							)
		rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc2.weight''', F'''vision_model.encoder.layers.{i}.mlp.fc2.weight''')							)
		rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc2.bias''', F'''vision_model.encoder.layers.{i}.mlp.fc2.bias''')							)

	# QFormer
	rename_keys.append(('Qformer.bert.embeddings.LayerNorm.weight', 'qformer.layernorm.weight')							)
	rename_keys.append(('Qformer.bert.embeddings.LayerNorm.bias', 'qformer.layernorm.bias')							)

	# fmt: on
	return rename_keys






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	lowercase       =    dct.pop(__SCREAMING_SNAKE_CASE							)
	lowercase       =    val






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	for i in range(config.vision_config.num_hidden_layers							):
		# read in original q and v biases
		lowercase       =    state_dict.pop(F'''visual_encoder.blocks.{i}.attn.q_bias'''							)
		lowercase       =    state_dict.pop(F'''visual_encoder.blocks.{i}.attn.v_bias'''							)

		# next, set bias in the state dict
		lowercase       =    torch.cat((q_bias, torch.zeros_like(__SCREAMING_SNAKE_CASE ,							requires_grad=__SCREAMING_SNAKE_CASE							), v_bias)							)
		lowercase       =    qkv_bias






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	lowercase       =    364 if 'coco' in model_name else 224
	lowercase       =    BlipaVisionConfig(image_size=__SCREAMING_SNAKE_CASE							).to_dict()

	# make sure the models have proper bos_token_id and eos_token_id set (important for generation)
	# seems like flan-T5 models don't have bos_token_id properly set?
	if "opt-2.7b" in model_name:
		lowercase       =    OPTConfig.from_pretrained('facebook/opt-2.7b' ,							eos_token_id=__SCREAMING_SNAKE_CASE							).to_dict()
	elif "opt-6.7b" in model_name:
		lowercase       =    OPTConfig.from_pretrained('facebook/opt-6.7b' ,							eos_token_id=__SCREAMING_SNAKE_CASE							).to_dict()
	elif "t5-xl" in model_name:
		lowercase       =    TaConfig.from_pretrained('google/flan-t5-xl' ,							dense_act_fn='gelu' ,							bos_token_id=1							).to_dict()
	elif "t5-xxl" in model_name:
		lowercase       =    TaConfig.from_pretrained('google/flan-t5-xxl' ,							dense_act_fn='gelu' ,							bos_token_id=1							).to_dict()

	lowercase       =    BlipaConfig(vision_config=__SCREAMING_SNAKE_CASE ,							text_config=__SCREAMING_SNAKE_CASE							)

	return config, image_size






@torch.no_grad()
def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE=None ,							__SCREAMING_SNAKE_CASE=False							):
	lowercase       =    (
	    AutoTokenizer.from_pretrained('facebook/opt-2.7b'							)
	    if 'opt' in model_name
	    else AutoTokenizer.from_pretrained('google/flan-t5-xl'							)
	)
	lowercase       =    tokenizer('\n' ,							add_special_tokens=__SCREAMING_SNAKE_CASE							).input_ids[0]
	lowercase  ,					lowercase       =    get_blipa_config(__SCREAMING_SNAKE_CASE ,							eos_token_id=__SCREAMING_SNAKE_CASE							)

	lowercase       =    BlipaForConditionalGeneration(__SCREAMING_SNAKE_CASE							).eval()

	lowercase       =    {
	    'blip2-opt-2.7b': ('blip2_opt', 'pretrain_opt2.7b'),
	    'blip2-opt-6.7b': ('blip2_opt', 'pretrain_opt6.7b'),
	    'blip2-opt-2.7b-coco': ('blip2_opt', 'caption_coco_opt2.7b'),
	    'blip2-opt-6.7b-coco': ('blip2_opt', 'caption_coco_opt6.7b'),
	    'blip2-flan-t5-xl': ('blip2_t5', 'pretrain_flant5xl'),
	    'blip2-flan-t5-xl-coco': ('blip2_t5', 'caption_coco_flant5xl'),
	    'blip2-flan-t5-xxl': ('blip2_t5', 'pretrain_flant5xxl'),
	}

	lowercase  ,					lowercase       =    model_name_to_original[model_name]

	# load original model
	print('Loading original model...'							)
	lowercase       =    'cuda' if torch.cuda.is_available() else 'cpu'
	lowercase  ,					lowercase  ,					lowercase       =    load_model_and_preprocess(
	    name=__SCREAMING_SNAKE_CASE ,							model_type=__SCREAMING_SNAKE_CASE ,							is_eval=__SCREAMING_SNAKE_CASE ,							device=__SCREAMING_SNAKE_CASE							)
	original_model.eval()
	print('Done!'							)

	# update state dict keys
	lowercase       =    original_model.state_dict()
	lowercase       =    create_rename_keys(__SCREAMING_SNAKE_CASE							)
	for src, dest in rename_keys:
		rename_key(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)

	# some keys can be renamed efficiently
	for key, val in state_dict.copy().items():
		lowercase       =    state_dict.pop(__SCREAMING_SNAKE_CASE							)
		if key.startswith('Qformer.bert'							):
			lowercase       =    key.replace('Qformer.bert' ,							'qformer'							)
		if "attention.self" in key:
			lowercase       =    key.replace('self' ,							'attention'							)
		if "opt_proj" in key:
			lowercase       =    key.replace('opt_proj' ,							'language_projection'							)
		if "t5_proj" in key:
			lowercase       =    key.replace('t5_proj' ,							'language_projection'							)
		if key.startswith('opt'							):
			lowercase       =    key.replace('opt' ,							'language'							)
		if key.startswith('t5'							):
			lowercase       =    key.replace('t5' ,							'language'							)
		lowercase       =    val

	# read in qv biases
	read_in_q_v_bias(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)

	lowercase  ,					lowercase       =    hf_model.load_state_dict(__SCREAMING_SNAKE_CASE ,							strict=__SCREAMING_SNAKE_CASE							)
	assert len(__SCREAMING_SNAKE_CASE							) == 0
	assert unexpected_keys == ["qformer.embeddings.position_ids"]

	lowercase       =    load_demo_image()
	lowercase       =    vis_processors['eval'](__SCREAMING_SNAKE_CASE							).unsqueeze(0							).to(__SCREAMING_SNAKE_CASE							)
	lowercase       =    tokenizer(['\n'] ,							return_tensors='pt'							).input_ids.to(__SCREAMING_SNAKE_CASE							)

	# create processor
	lowercase       =    BlipImageProcessor(
	    size={'height': image_size, 'width': image_size} ,							image_mean=__SCREAMING_SNAKE_CASE ,							image_std=__SCREAMING_SNAKE_CASE							)
	lowercase       =    BlipaProcessor(image_processor=__SCREAMING_SNAKE_CASE ,							tokenizer=__SCREAMING_SNAKE_CASE							)
	lowercase       =    processor(images=__SCREAMING_SNAKE_CASE ,							return_tensors='pt'							).pixel_values.to(__SCREAMING_SNAKE_CASE							)

	# make sure processor creates exact same pixel values
	assert torch.allclose(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)

	original_model.to(__SCREAMING_SNAKE_CASE							)
	hf_model.to(__SCREAMING_SNAKE_CASE							)
	with torch.no_grad():
		if "opt" in model_name:
			lowercase       =    original_model({'image': original_pixel_values, 'text_input': ['']}							).logits
			lowercase       =    hf_model(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							).logits
		else:
			lowercase       =    original_model(
			    {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}							).logits
			lowercase       =    input_ids.masked_fill(input_ids == tokenizer.pad_token_id ,							-100							)
			lowercase       =    hf_model(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							labels=__SCREAMING_SNAKE_CASE							).logits

	assert original_logits.shape == logits.shape
	print('First values of original logits:' ,							original_logits[0, :3, :3]							)
	print('First values of HF logits:' ,							logits[0, :3, :3]							)

	# assert values
	if model_name == "blip2-flan-t5-xl":
		lowercase       =    torch.tensor(
		    [[-41.58_50, -4.44_40, -8.99_22], [-47.43_22, -5.91_43, -1.73_40]] ,							device=__SCREAMING_SNAKE_CASE							)
		assert torch.allclose(logits[0, :3, :3] ,							__SCREAMING_SNAKE_CASE ,							atol=1e-4							)
	elif model_name == "blip2-flan-t5-xl-coco":
		lowercase       =    torch.tensor(
		    [[-57.01_09, -9.89_67, -12.62_80], [-68.65_78, -12.71_91, -10.50_65]] ,							device=__SCREAMING_SNAKE_CASE							)
	else:
		# cast to same type
		lowercase       =    logits.dtype
		assert torch.allclose(original_logits.to(__SCREAMING_SNAKE_CASE							) ,							__SCREAMING_SNAKE_CASE ,							atol=1e-2							)
	print('Looks ok!'							)

	print('Generating a caption...'							)
	lowercase       =    ''
	lowercase       =    tokenizer(__SCREAMING_SNAKE_CASE ,							return_tensors='pt'							).input_ids.to(__SCREAMING_SNAKE_CASE							)

	lowercase       =    original_model.generate({'image': original_pixel_values}							)
	lowercase       =    hf_model.generate(
	    __SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							do_sample=__SCREAMING_SNAKE_CASE ,							num_beams=5 ,							max_length=30 ,							min_length=1 ,							top_p=0.9 ,							repetition_penalty=1.0 ,							length_penalty=1.0 ,							temperature=1 ,							)
	print('Original generation:' ,							__SCREAMING_SNAKE_CASE							)
	lowercase       =    input_ids.shape[1]
	lowercase       =    processor.batch_decode(outputs[:, prompt_length:] ,							skip_special_tokens=__SCREAMING_SNAKE_CASE							)
	lowercase       =    [text.strip() for text in output_text]
	print('HF generation:' ,							__SCREAMING_SNAKE_CASE							)

	if pytorch_dump_folder_path is not None:
		processor.save_pretrained(__SCREAMING_SNAKE_CASE							)
		hf_model.save_pretrained(__SCREAMING_SNAKE_CASE							)

	if push_to_hub:
		processor.push_to_hub(F'''nielsr/{model_name}'''							)
		hf_model.push_to_hub(F'''nielsr/{model_name}'''							)


if __name__ == "__main__":
	UpperCAmelCase	      = argparse.ArgumentParser()
	UpperCAmelCase	      = [
	    '''blip2-opt-2.7b''',
	    '''blip2-opt-6.7b''',
	    '''blip2-opt-2.7b-coco''',
	    '''blip2-opt-6.7b-coco''',
	    '''blip2-flan-t5-xl''',
	    '''blip2-flan-t5-xl-coco''',
	    '''blip2-flan-t5-xxl''',
	]
	parser.add_argument(
	    '''--model_name''',
	    default='''blip2-opt-2.7b''',
	    choices=choices,
	    type=str,
	    help='''Path to hf config.json of model to convert''',
	)
	parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
	parser.add_argument(
	    '''--push_to_hub''',
	    action='''store_true''',
	    help='''Whether to push the model and processor to the hub after converting''',
	)

	UpperCAmelCase	      = parser.parse_args()

	convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)



import sys
from collections import defaultdict







class        A_   :




	'''simple docstring'''


	def __init__(       self       ):
		lowercase       =    []





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case       ):
		return self.node_position[vertex]





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case       ):
		lowercase       =    pos





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    snake_case       ):
		if start > size // 2 - 1:
			return
		else:
			if 2 * start + 2 >= size:
				lowercase       =    2 * start + 1
			else:
				if heap[2 * start + 1] < heap[2 * start + 2]:
					lowercase       =    2 * start + 1
				else:
					lowercase       =    2 * start + 2
			if heap[smallest_child] < heap[start]:
				lowercase  ,					lowercase       =    heap[smallest_child], positions[smallest_child]
				lowercase  ,					lowercase       =    (
				    heap[start],
				    positions[start],
				)
				lowercase  ,					lowercase       =    temp, tempa

				lowercase       =    self.get_position(positions[smallest_child]       )
				self.set_position(
				    positions[smallest_child]		,    self.get_position(positions[start]       )       )
				self.set_position(positions[start]		,    snake_case       )

				self.top_to_bottom(snake_case		,    snake_case		,    snake_case		,    snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    snake_case       ):
		lowercase       =    position[index]

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

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





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case       ):
		lowercase       =    len(snake_case       ) // 2 - 1
		for i in range(snake_case		,    -1		,    -1       ):
			self.top_to_bottom(snake_case		,    snake_case		,    len(snake_case       )		,    snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case       ):
		lowercase       =    positions[0]
		lowercase       =    sys.maxsize
		self.top_to_bottom(snake_case		,    0		,    len(snake_case       )		,    snake_case       )
		return temp






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	lowercase       =    Heap()

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

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

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

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


if __name__ == "__main__":  # pragma: no cover
	# < --------- Prims Algorithm --------- >
	UpperCAmelCase	      = int(input('''Enter number of edges: ''').strip())
	UpperCAmelCase	      = defaultdict(list)
	for _ in range(edges_number):
		UpperCAmelCase	      = [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))

