



'''simple docstring'''




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       __lowerCamelCase  (					)      ->			Dict:
							UpperCAmelCase     :		Union[str, Any]							     =		"""https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
							UpperCAmelCase     :		Optional[int]							     =		Image.open(requests.get(_lowercase	,				stream=_lowercase		).raw		).convert("""RGB"""		)

							return image



def       __lowerCamelCase  (					_lowercase		)      ->			int:
							UpperCAmelCase     :		List[Any]							     =		[]
							# 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       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase		)      ->			Dict:
							UpperCAmelCase     :		List[Any]							     =		dct.pop(_lowercase		)
							UpperCAmelCase     :		Union[str, Any]							     =		val



def       __lowerCamelCase  (					_lowercase	,				_lowercase		)      ->			List[str]:
							for i in range(config.vision_config.num_hidden_layers		):
														# read in original q and v biases
														UpperCAmelCase     :		Any							     =		state_dict.pop(F'''visual_encoder.blocks.{i}.attn.q_bias'''		)
														UpperCAmelCase     :		Tuple							     =		state_dict.pop(F'''visual_encoder.blocks.{i}.attn.v_bias'''		)

														# next, set bias in the state dict
														UpperCAmelCase     :		Union[str, Any]							     =		torch.cat((q_bias, torch.zeros_like(_lowercase	,				requires_grad=_lowercase		), v_bias)		)
														UpperCAmelCase     :		Any							     =		qkv_bias



def       __lowerCamelCase  (					_lowercase	,				_lowercase		)      ->			str:
							UpperCAmelCase     :		Tuple							     =		3_6_4 if """coco""" in model_name else 2_2_4
							UpperCAmelCase     :		Any							     =		BlipaVisionConfig(image_size=_lowercase		).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:
														UpperCAmelCase     :		Optional[int]							     =		OPTConfig.from_pretrained("""facebook/opt-2.7b"""	,				eos_token_id=_lowercase		).to_dict()
							elif "opt-6.7b" in model_name:
														UpperCAmelCase     :		str							     =		OPTConfig.from_pretrained("""facebook/opt-6.7b"""	,				eos_token_id=_lowercase		).to_dict()
							elif "t5-xl" in model_name:
														UpperCAmelCase     :		Optional[Any]							     =		TaConfig.from_pretrained("""google/flan-t5-xl"""	,				dense_act_fn="""gelu"""	,				bos_token_id=1		).to_dict()
							elif "t5-xxl" in model_name:
														UpperCAmelCase     :		Tuple							     =		TaConfig.from_pretrained("""google/flan-t5-xxl"""	,				dense_act_fn="""gelu"""	,				bos_token_id=1		).to_dict()

							UpperCAmelCase     :		Dict							     =		BlipaConfig(vision_config=_lowercase	,				text_config=_lowercase		)

							return config, image_size



@torch.no_grad()
def       __lowerCamelCase  (					_lowercase	,				_lowercase=None	,				_lowercase=False		)      ->			Dict:
							UpperCAmelCase     :		Dict							     =		(
							    AutoTokenizer.from_pretrained("""facebook/opt-2.7b"""		)
							    if """opt""" in model_name
							    else AutoTokenizer.from_pretrained("""google/flan-t5-xl"""		)
							)
							UpperCAmelCase     :		int							     =		tokenizer("""\n"""	,				add_special_tokens=_lowercase		).input_ids[0]
							UpperCAmelCase     ,     UpperCAmelCase     :		Any							     =		get_blipa_config(_lowercase	,				eos_token_id=_lowercase		)

							UpperCAmelCase     :		int							     =		BlipaForConditionalGeneration(_lowercase		).eval()

							UpperCAmelCase     :		int							     =		{
							    """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"""),
							}

							UpperCAmelCase     ,     UpperCAmelCase     :		Optional[Any]							     =		model_name_to_original[model_name]

							# load original model
							print("""Loading original model..."""		)
							UpperCAmelCase     :		Any							     =		"""cuda""" if torch.cuda.is_available() else """cpu"""
							UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		str							     =		load_model_and_preprocess(
							    name=_lowercase	,				model_type=_lowercase	,				is_eval=_lowercase	,				device=_lowercase		)
							original_model.eval()
							print("""Done!"""		)

							# update state dict keys
							UpperCAmelCase     :		Optional[int]							     =		original_model.state_dict()
							UpperCAmelCase     :		str							     =		create_rename_keys(_lowercase		)
							for src, dest in rename_keys:
														rename_key(_lowercase	,				_lowercase	,				_lowercase		)

							# some keys can be renamed efficiently
							for key, val in state_dict.copy().items():
														UpperCAmelCase     :		int							     =		state_dict.pop(_lowercase		)
														if key.startswith("""Qformer.bert"""		):
																					UpperCAmelCase     :		Tuple							     =		key.replace("""Qformer.bert"""	,				"""qformer"""		)
														if "attention.self" in key:
																					UpperCAmelCase     :		List[Any]							     =		key.replace("""self"""	,				"""attention"""		)
														if "opt_proj" in key:
																					UpperCAmelCase     :		Dict							     =		key.replace("""opt_proj"""	,				"""language_projection"""		)
														if "t5_proj" in key:
																					UpperCAmelCase     :		Optional[Any]							     =		key.replace("""t5_proj"""	,				"""language_projection"""		)
														if key.startswith("""opt"""		):
																					UpperCAmelCase     :		Dict							     =		key.replace("""opt"""	,				"""language"""		)
														if key.startswith("""t5"""		):
																					UpperCAmelCase     :		Any							     =		key.replace("""t5"""	,				"""language"""		)
														UpperCAmelCase     :		Tuple							     =		val

							# read in qv biases
							read_in_q_v_bias(_lowercase	,				_lowercase		)

							UpperCAmelCase     ,     UpperCAmelCase     :		List[str]							     =		hf_model.load_state_dict(_lowercase	,				strict=_lowercase		)
							assert len(_lowercase		) == 0
							assert unexpected_keys == ["qformer.embeddings.position_ids"]

							UpperCAmelCase     :		List[str]							     =		load_demo_image()
							UpperCAmelCase     :		Dict							     =		vis_processors["""eval"""](_lowercase		).unsqueeze(0		).to(_lowercase		)
							UpperCAmelCase     :		Optional[int]							     =		tokenizer(["""\n"""]	,				return_tensors="""pt"""		).input_ids.to(_lowercase		)

							# create processor
							UpperCAmelCase     :		Union[str, Any]							     =		BlipImageProcessor(
							    size={"""height""": image_size, """width""": image_size}	,				image_mean=_lowercase	,				image_std=_lowercase		)
							UpperCAmelCase     :		int							     =		BlipaProcessor(image_processor=_lowercase	,				tokenizer=_lowercase		)
							UpperCAmelCase     :		Union[str, Any]							     =		processor(images=_lowercase	,				return_tensors="""pt"""		).pixel_values.to(_lowercase		)

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

							original_model.to(_lowercase		)
							hf_model.to(_lowercase		)
							with torch.no_grad():
														if "opt" in model_name:
																					UpperCAmelCase     :		Optional[Any]							     =		original_model({"""image""": original_pixel_values, """text_input""": [""""""]}		).logits
																					UpperCAmelCase     :		str							     =		hf_model(_lowercase	,				_lowercase		).logits
														else:
																					UpperCAmelCase     :		Tuple							     =		original_model(
																					    {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]}		).logits
																					UpperCAmelCase     :		List[str]							     =		input_ids.masked_fill(input_ids == tokenizer.pad_token_id	,				-1_0_0		)
																					UpperCAmelCase     :		List[str]							     =		hf_model(_lowercase	,				_lowercase	,				labels=_lowercase		).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":
														UpperCAmelCase     :		str							     =		torch.tensor(
														    [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]	,				device=_lowercase		)
														assert torch.allclose(logits[0, :3, :3]	,				_lowercase	,				atol=1e-4		)
							elif model_name == "blip2-flan-t5-xl-coco":
														UpperCAmelCase     :		Tuple							     =		torch.tensor(
														    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]	,				device=_lowercase		)
							else:
														# cast to same type
														UpperCAmelCase     :		Optional[Any]							     =		logits.dtype
														assert torch.allclose(original_logits.to(_lowercase		)	,				_lowercase	,				atol=1e-2		)
							print("""Looks ok!"""		)

							print("""Generating a caption..."""		)
							UpperCAmelCase     :		Tuple							     =		""""""
							UpperCAmelCase     :		Tuple							     =		tokenizer(_lowercase	,				return_tensors="""pt"""		).input_ids.to(_lowercase		)

							UpperCAmelCase     :		str							     =		original_model.generate({"""image""": original_pixel_values}		)
							UpperCAmelCase     :		Optional[int]							     =		hf_model.generate(
							    _lowercase	,				_lowercase	,				do_sample=_lowercase	,				num_beams=5	,				max_length=3_0	,				min_length=1	,				top_p=0.9	,				repetition_penalty=1.0	,				length_penalty=1.0	,				temperature=1	,				)
							print("""Original generation:"""	,				_lowercase		)
							UpperCAmelCase     :		Optional[int]							     =		input_ids.shape[1]
							UpperCAmelCase     :		Optional[Any]							     =		processor.batch_decode(outputs[:, prompt_length:]	,				skip_special_tokens=_lowercase		)
							UpperCAmelCase     :		str							     =		[text.strip() for text in output_text]
							print("""HF generation:"""	,				_lowercase		)

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

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


if __name__ == "__main__":
							a   :			Optional[Any]					  =    argparse.ArgumentParser()
							a   :			List[str]					  =    [
							    """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""",
							)

							a   :			List[str]					  =    parser.parse_args()

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






'''simple docstring'''




from operator import delitem, getitem, setitem

import pytest

from data_structures.hashing.hash_map import HashMap



def       __lowerCamelCase  (					_lowercase		)      ->			Optional[Any]:
							return getitem, k



def       __lowerCamelCase  (					_lowercase	,				_lowercase		)      ->			List[str]:
							return setitem, k, v



def       __lowerCamelCase  (					_lowercase		)      ->			int:
							return delitem, k



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				*_lowercase		)      ->			Optional[Any]:
							try:
														return fun(_lowercase	,				*_lowercase		), None
							except Exception as e:
														return None, e


a   :			List[str]					  =    (
    _set("""key_a""", """val_a"""),
    _set("""key_b""", """val_b"""),
)

a   :			List[Any]					  =    [
    _set("""key_a""", """val_a"""),
    _set("""key_a""", """val_b"""),
]

a   :			int					  =    [
    _set("""key_a""", """val_a"""),
    _set("""key_b""", """val_b"""),
    _del("""key_a"""),
    _del("""key_b"""),
    _set("""key_a""", """val_a"""),
    _del("""key_a"""),
]

a   :			List[Any]					  =    [
    _get("""key_a"""),
    _del("""key_a"""),
    _set("""key_a""", """val_a"""),
    _del("""key_a"""),
    _del("""key_a"""),
    _get("""key_a"""),
]

a   :			Tuple					  =    [
    *[_set(x, x) for x in range(5)],  # guaranteed upsize
]

a   :			Optional[Any]					  =    [
    *[_set(x, x) for x in range(5)],  # guaranteed upsize
    *[_del(x) for x in range(5)],
    _set("""key_a""", """val_b"""),
]



@pytest.mark.parametrize(
    """operations"""	,				(
        pytest.param(_add_items	,				id="""add items"""		),
        pytest.param(_overwrite_items	,				id="""overwrite items"""		),
        pytest.param(_delete_items	,				id="""delete items"""		),
        pytest.param(_access_absent_items	,				id="""access absent items"""		),
        pytest.param(_add_with_resize_up	,				id="""add with resize up"""		),
        pytest.param(_add_with_resize_down	,				id="""add with resize down"""		),
    )	,				)
def       __lowerCamelCase  (					_lowercase		)      ->			Optional[int]:
							UpperCAmelCase     :		List[str]							     =		HashMap(initial_block_size=4		)
							UpperCAmelCase     :		Dict							     =		{}
							for _, (fun, *args) in enumerate(_lowercase		):
														UpperCAmelCase     ,     UpperCAmelCase     :		Union[str, Any]							     =		_run_operation(_lowercase	,				_lowercase	,				*_lowercase		)
														UpperCAmelCase     ,     UpperCAmelCase     :		Any							     =		_run_operation(_lowercase	,				_lowercase	,				*_lowercase		)
														assert my_res == py_res
														assert str(_lowercase		) == str(_lowercase		)
														assert set(_lowercase		) == set(_lowercase		)
														assert len(_lowercase		) == len(_lowercase		)
														assert set(my.items()		) == set(py.items()		)



def       __lowerCamelCase  (					)      ->			List[Any]:
							def is_public(_lowercase		) -> bool:
														return not name.startswith("""_"""		)

							UpperCAmelCase     :		int							     =		{name for name in dir({}		) if is_public(_lowercase		)}
							UpperCAmelCase     :		Any							     =		{name for name in dir(HashMap()		) if is_public(_lowercase		)}

							assert dict_public_names > hash_public_names

