"""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 				_UpperCAmelCase		(							):





		"""simple docstring"""




		lowerCAmelCase__   = """https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
		lowerCAmelCase__   = Image.open(requests.get(lowerCamelCase__      ,	stream=lowerCamelCase__  ).raw  ).convert("""RGB"""  )

		return image







def 				_UpperCAmelCase		(							lowerCamelCase__  ):





		"""simple docstring"""




		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 				_UpperCAmelCase		(							lowerCamelCase__      ,	lowerCamelCase__      ,	lowerCamelCase__  ):





		"""simple docstring"""




		lowerCAmelCase__   = dct.pop(lowerCamelCase__  )
		lowerCAmelCase__   = val







def 				_UpperCAmelCase		(							lowerCamelCase__      ,	lowerCamelCase__  ):





		"""simple docstring"""




		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(lowerCamelCase__      ,	requires_grad=lowerCamelCase__  ), v_bias)  )
				lowerCAmelCase__   = qkv_bias







def 				_UpperCAmelCase		(							lowerCamelCase__      ,	lowerCamelCase__  ):





		"""simple docstring"""




		lowerCAmelCase__   = 364 if """coco""" in model_name else 224
		lowerCAmelCase__   = BlipaVisionConfig(image_size=lowerCamelCase__  ).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=lowerCamelCase__  ).to_dict()
		elif "opt-6.7b" in model_name:
				lowerCAmelCase__   = OPTConfig.from_pretrained("""facebook/opt-6.7b"""      ,	eos_token_id=lowerCamelCase__  ).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=lowerCamelCase__      ,	text_config=lowerCamelCase__  )

		return config, image_size







@torch.no_grad()
def 				_UpperCAmelCase		(							lowerCamelCase__      ,	lowerCamelCase__=None      ,	lowerCamelCase__=False  ):





		"""simple docstring"""




		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=lowerCamelCase__  ).input_ids[0]
		lowerCAmelCase__ ,       lowerCAmelCase__   = get_blipa_config(lowerCamelCase__      ,	eos_token_id=lowerCamelCase__  )

		lowerCAmelCase__   = BlipaForConditionalGeneration(lowerCamelCase__  ).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=lowerCamelCase__      ,	model_type=lowerCamelCase__      ,	is_eval=lowerCamelCase__      ,	device=lowerCamelCase__  )
		original_model.eval()
		print("""Done!"""  )

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

		# some keys can be renamed efficiently
		for key, val in state_dict.copy().items():
				lowerCAmelCase__   = state_dict.pop(lowerCamelCase__  )
				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(lowerCamelCase__      ,	lowerCamelCase__  )

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

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

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

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

		original_model.to(lowerCamelCase__  )
		hf_model.to(lowerCamelCase__  )
		with torch.no_grad():
				if "opt" in model_name:
						lowerCAmelCase__   = original_model({"""image""": original_pixel_values, """text_input""": [""""""]}  ).logits
						lowerCAmelCase__   = hf_model(lowerCamelCase__      ,	lowerCamelCase__  ).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      ,	-100  )
						lowerCAmelCase__   = hf_model(lowerCamelCase__      ,	lowerCamelCase__      ,	labels=lowerCamelCase__  ).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.58_50, -4.44_40, -8.99_22], [-47.43_22, -5.91_43, -1.73_40]]      ,	device=lowerCamelCase__  )
				assert torch.allclose(logits[0, :3, :3]      ,	lowerCamelCase__      ,	atol=1e-4  )
		elif model_name == "blip2-flan-t5-xl-coco":
				lowerCAmelCase__   = torch.tensor(
				    [[-57.01_09, -9.89_67, -12.62_80], [-68.65_78, -12.71_91, -10.50_65]]      ,	device=lowerCamelCase__  )
		else:
				# cast to same type
				lowerCAmelCase__   = logits.dtype
				assert torch.allclose(original_logits.to(lowerCamelCase__  )      ,	lowerCamelCase__      ,	atol=1e-2  )
		print("""Looks ok!"""  )

		print("""Generating a caption..."""  )
		lowerCAmelCase__   = """"""
		lowerCAmelCase__   = tokenizer(lowerCamelCase__      ,	return_tensors="""pt"""  ).input_ids.to(lowerCamelCase__  )

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

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

		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     :				str    				=    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     :				Any    				=    parser.parse_args()

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





"""simple docstring"""




import json
import os
import unittest

from transformers import MgpstrTokenizer
from transformers.models.mgp_str.tokenization_mgp_str import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin

@require_tokenizers
class    a_   (  __UpperCamelCase			,		unittest.TestCase      ):
					UpperCamelCase_			:						Optional[int]      =     MgpstrTokenizer
					UpperCamelCase_			:						Optional[int]      =     False
					UpperCamelCase_			:						Tuple      =     {}
					UpperCamelCase_			:						List[str]      =     False




					def     _SCREAMING_SNAKE_CASE    (						self    :  Dict  ):
							super().setUp()

							# fmt: off
							lowerCAmelCase__   = ["""[GO]""", """[s]""", """0""", """1""", """2""", """3""", """4""", """5""", """6""", """7""", """8""", """9""", """a""", """b""", """c""", """d""", """e""", """f""", """g""", """h""", """i""", """j""", """k""", """l""", """m""", """n""", """o""", """p""", """q""", """r""", """s""", """t""", """u""", """v""", """w""", """x""", """y""", """z"""]
							# fmt: on
							lowerCAmelCase__   = dict(zip(snake_case__	,    range(len(snake_case__  )  )  )  )

							lowerCAmelCase__   = os.path.join(self.tmpdirname	,    VOCAB_FILES_NAMES["""vocab_file"""]  )
							with open(self.vocab_file	,    """w"""	,    encoding="""utf-8"""  ) as fp:
									fp.write(json.dumps(snake_case__  ) + """\n"""  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  int	,    **snake_case__    :  Tuple  ):
							return MgpstrTokenizer.from_pretrained(self.tmpdirname	,    **snake_case__  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  List[str]	,    snake_case__    :  Dict  ):
							lowerCAmelCase__   = """tester"""
							lowerCAmelCase__   = """tester"""
							return input_text, output_text




					@unittest.skip("""MGP-STR always lower cases letters."""  )
					def     _SCREAMING_SNAKE_CASE    (						self    :  Optional[int]  ):
							pass




					def     _SCREAMING_SNAKE_CASE    (						self    :  Optional[int]  ):
							lowerCAmelCase__   = self.get_tokenizers(do_lower_case=snake_case__  )
							for tokenizer in tokenizers:
									with self.subTest(F"""{tokenizer.__class__.__name__}"""  ):
											lowerCAmelCase__   = """[SPECIAL_TOKEN]"""

											tokenizer.add_special_tokens({"""cls_token""": special_token}  )
											lowerCAmelCase__   = tokenizer.encode([special_token]	,    add_special_tokens=snake_case__  )
											self.assertEqual(len(snake_case__  )	,    1  )

											lowerCAmelCase__   = tokenizer.decode(snake_case__	,    skip_special_tokens=snake_case__  )
											self.assertTrue(special_token not in decoded  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  List[str]  ):
							lowerCAmelCase__   = self.get_tokenizers()
							for tokenizer in tokenizers:
									with self.subTest(F"""{tokenizer.__class__.__name__}"""  ):
											lowerCAmelCase__ ,       lowerCAmelCase__   = self.get_input_output_texts(snake_case__  )

											lowerCAmelCase__   = tokenizer.tokenize(snake_case__  )
											lowerCAmelCase__   = tokenizer.convert_tokens_to_ids(snake_case__  )
											lowerCAmelCase__   = tokenizer.encode(snake_case__	,    add_special_tokens=snake_case__  )
											self.assertListEqual(snake_case__	,    snake_case__  )

											lowerCAmelCase__   = tokenizer.convert_ids_to_tokens(snake_case__  )
											self.assertNotEqual(len(snake_case__  )	,    0  )
											lowerCAmelCase__   = tokenizer.decode(snake_case__  )
											self.assertIsInstance(snake_case__	,    snake_case__  )

											self.assertEqual(text_a.replace(""" """	,    """"""  )	,    snake_case__  )




					@unittest.skip("""MGP-STR tokenizer only handles one sequence."""  )
					def     _SCREAMING_SNAKE_CASE    (						self    :  Optional[Any]  ):
							pass




					@unittest.skip("""inputs cannot be pretokenized in MgpstrTokenizer"""  )
					def     _SCREAMING_SNAKE_CASE    (						self    :  List[Any]  ):
							pass




