




"""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 					a__				(       )		->      Tuple:
			lowerCamelCase       =      """https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
			lowerCamelCase       =      Image.open(requests.get(snake_case__ ,				stream=snake_case__       ).raw       ).convert("""RGB"""       )

			return image



def 					a__				(       snake_case__       )		->      Optional[Any]:
			lowerCamelCase       =      []
			# 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 					a__				(       snake_case__ ,				snake_case__ ,				snake_case__       )		->      List[str]:
			lowerCamelCase       =      dct.pop(snake_case__       )
			lowerCamelCase       =      val



def 					a__				(       snake_case__ ,				snake_case__       )		->      Tuple:


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

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



def 					a__				(       snake_case__ ,				snake_case__       )		->      int:
			lowerCamelCase       =      3_64 if """coco""" in model_name else 2_24
			lowerCamelCase       =      BlipaVisionConfig(image_size=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:
						lowerCamelCase       =      OPTConfig.from_pretrained("""facebook/opt-2.7b""" ,				eos_token_id=snake_case__       ).to_dict()
			elif "opt-6.7b" in model_name:
						lowerCamelCase       =      OPTConfig.from_pretrained("""facebook/opt-6.7b""" ,				eos_token_id=snake_case__       ).to_dict()
			elif "t5-xl" in model_name:
						lowerCamelCase       =      TaConfig.from_pretrained("""google/flan-t5-xl""" ,				dense_act_fn="""gelu""" ,				bos_token_id=1       ).to_dict()
			elif "t5-xxl" in model_name:
						lowerCamelCase       =      TaConfig.from_pretrained("""google/flan-t5-xxl""" ,				dense_act_fn="""gelu""" ,				bos_token_id=1       ).to_dict()

			lowerCamelCase       =      BlipaConfig(vision_config=snake_case__ ,				text_config=snake_case__       )

			return config, image_size



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

			lowerCamelCase       =      BlipaForConditionalGeneration(snake_case__       ).eval()

			lowerCamelCase       =      {
			    """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"""),
			}

			lowerCamelCase   ,     lowerCamelCase       =      model_name_to_original[model_name]

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

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

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

			# read in qv biases
			read_in_q_v_bias(snake_case__ ,				snake_case__       )

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

			lowerCamelCase       =      load_demo_image()
			lowerCamelCase       =      vis_processors["""eval"""](snake_case__       ).unsqueeze(0       ).to(snake_case__       )
			lowerCamelCase       =      tokenizer(["""\n"""] ,				return_tensors="""pt"""       ).input_ids.to(snake_case__       )

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

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

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

			print("""Generating a caption..."""       )
			lowerCamelCase       =      """"""
			lowerCamelCase       =      tokenizer(snake_case__ ,				return_tensors="""pt"""       ).input_ids.to(snake_case__       )

			lowerCamelCase       =      original_model.generate({"""image""": original_pixel_values}       )
			lowerCamelCase       =      hf_model.generate(
			    snake_case__ ,				snake_case__ ,				do_sample=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:""" ,				snake_case__       )
			lowerCamelCase       =      input_ids.shape[1]
			lowerCamelCase       =      processor.batch_decode(outputs[:, prompt_length:] ,				skip_special_tokens=snake_case__       )
			lowerCamelCase       =      [text.strip() for text in output_text]
			print("""HF generation:""" ,				snake_case__       )

			if pytorch_dump_folder_path is not None:
						processor.save_pretrained(snake_case__       )
						hf_model.save_pretrained(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__":
					lowerCAmelCase							:    Any  							=							argparse.ArgumentParser()
					lowerCAmelCase							:    List[Any]  							=							[
					    """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""",
					)

					lowerCAmelCase							:    Dict  							=							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 					a__				(       snake_case__       )		->      List[str]:
			return getitem, k



def 					a__				(       snake_case__ ,				snake_case__       )		->      Optional[Any]:
			return setitem, k, v



def 					a__				(       snake_case__       )		->      str:
			return delitem, k



def 					a__				(       snake_case__ ,				snake_case__ ,				*snake_case__       )		->      Union[str, Any]:
			try:
						return fun(snake_case__ ,				*snake_case__       ), None
			except Exception as e:
						return None, e


lowerCAmelCase							:    Optional[int]  							=							(
    _set("""key_a""", """val_a"""),
    _set("""key_b""", """val_b"""),
)

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

lowerCAmelCase							:    List[Any]  							=							[
    _set("""key_a""", """val_a"""),
    _set("""key_b""", """val_b"""),
    _del("""key_a"""),
    _del("""key_b"""),
    _set("""key_a""", """val_a"""),
    _del("""key_a"""),
]

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

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

lowerCAmelCase							:    Tuple  							=							[
    *[_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 					a__				(       snake_case__       )		->      Any:
			lowerCamelCase       =      HashMap(initial_block_size=4       )
			lowerCamelCase       =      {}
			for _, (fun, *args) in enumerate(snake_case__       ):
						lowerCamelCase   ,     lowerCamelCase       =      _run_operation(snake_case__ ,				snake_case__ ,				*snake_case__       )
						lowerCamelCase   ,     lowerCamelCase       =      _run_operation(snake_case__ ,				snake_case__ ,				*snake_case__       )
						assert my_res == py_res
						assert str(snake_case__       ) == str(snake_case__       )
						assert set(snake_case__       ) == set(snake_case__       )
						assert len(snake_case__       ) == len(snake_case__       )
						assert set(my.items()       ) == set(py.items()       )



def 					a__				(       )		->      int:
			def is_public(snake_case__       ) -> bool:
						return not name.startswith("""_"""       )

			lowerCamelCase       =      {name for name in dir({}       ) if is_public(snake_case__       )}
			lowerCamelCase       =      {name for name in dir(HashMap()       ) if is_public(snake_case__       )}

			assert dict_public_names > hash_public_names



