


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 		lowercase_  (     ):


			'''simple docstring'''







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

			return image


def 		lowercase_  (     _UpperCamelCase  ):


			'''simple docstring'''







			__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 		lowercase_  (     _UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase  ):


			'''simple docstring'''







			__lowercase 					=   dct.pop(_UpperCamelCase  )
			__lowercase 					=   val


def 		lowercase_  (     _UpperCamelCase       ,			_UpperCamelCase  ):


			'''simple docstring'''







			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(_UpperCamelCase       ,			requires_grad=_UpperCamelCase  ), v_bias)  )
						__lowercase 					=   qkv_bias


def 		lowercase_  (     _UpperCamelCase       ,			_UpperCamelCase  ):


			'''simple docstring'''







			__lowercase 					=   3_64 if '''coco''' in model_name else 2_24
			__lowercase 					=   BlipaVisionConfig(image_size=_UpperCamelCase  ).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=_UpperCamelCase  ).to_dict()
			elif "opt-6.7b" in model_name:
						__lowercase 					=   OPTConfig.from_pretrained('''facebook/opt-6.7b'''       ,			eos_token_id=_UpperCamelCase  ).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=_UpperCamelCase       ,			text_config=_UpperCamelCase  )

			return config, image_size


@torch.no_grad()
def 		lowercase_  (     _UpperCamelCase       ,			_UpperCamelCase=None       ,			_UpperCamelCase=False  ):


			'''simple docstring'''







			__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=_UpperCamelCase  ).input_ids[0]
			__lowercase					, __lowercase 					=   get_blipa_config(_UpperCamelCase       ,			eos_token_id=_UpperCamelCase  )

			__lowercase 					=   BlipaForConditionalGeneration(_UpperCamelCase  ).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=_UpperCamelCase       ,			model_type=_UpperCamelCase       ,			is_eval=_UpperCamelCase       ,			device=_UpperCamelCase  )
			original_model.eval()
			print('''Done!'''  )

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

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

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

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

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

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

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

			print('''Generating a caption...'''  )
			__lowercase 					=   ''''''
			__lowercase 					=   tokenizer(_UpperCamelCase       ,			return_tensors='''pt'''  ).input_ids.to(_UpperCamelCase  )

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

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

			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 : Dict		=				argparse.ArgumentParser()
	a : Dict		=				[
	    '''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 : Tuple		=				parser.parse_args()

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






import os

from bleurt import score  # From: git+https://github.com/google-research/bleurt.git

import datasets


a : Optional[int]		=				datasets.logging.get_logger(__name__)


a : Tuple		=				'''\
@inproceedings{bleurt,
  title={BLEURT: Learning Robust Metrics for Text Generation},
  author={Thibault Sellam and Dipanjan Das and Ankur P. Parikh},
  booktitle={ACL},
  year={2020},
  url={https://arxiv.org/abs/2004.04696}
}
'''

a : Union[str, Any]		=				'''\
BLEURT a learnt evaluation metric for Natural Language Generation. It is built using multiple phases of transfer learning starting from a pretrained BERT model (Devlin et al. 2018)
and then employing another pre-training phrase using synthetic data. Finally it is trained on WMT human annotations. You may run BLEURT out-of-the-box or fine-tune
it for your specific application (the latter is expected to perform better).

See the project\'s README at https://github.com/google-research/bleurt#readme for more information.
'''

a : Union[str, Any]		=				'''
BLEURT score.

Args:
    `predictions` (list of str): prediction/candidate sentences
    `references` (list of str): reference sentences
    `checkpoint` BLEURT checkpoint. Will default to BLEURT-tiny if None.

Returns:
    \'scores\': List of scores.
Examples:

    >>> predictions = ["hello there", "general kenobi"]
    >>> references = ["hello there", "general kenobi"]
    >>> bleurt = datasets.load_metric("bleurt")
    >>> results = bleurt.compute(predictions=predictions, references=references)
    >>> print([round(v, 2) for v in results["scores"]])
    [1.03, 1.04]
'''

a : Tuple		=				{
    '''bleurt-tiny-128''': '''https://storage.googleapis.com/bleurt-oss/bleurt-tiny-128.zip''',
    '''bleurt-tiny-512''': '''https://storage.googleapis.com/bleurt-oss/bleurt-tiny-512.zip''',
    '''bleurt-base-128''': '''https://storage.googleapis.com/bleurt-oss/bleurt-base-128.zip''',
    '''bleurt-base-512''': '''https://storage.googleapis.com/bleurt-oss/bleurt-base-512.zip''',
    '''bleurt-large-128''': '''https://storage.googleapis.com/bleurt-oss/bleurt-large-128.zip''',
    '''bleurt-large-512''': '''https://storage.googleapis.com/bleurt-oss/bleurt-large-512.zip''',
    '''BLEURT-20-D3''': '''https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D3.zip''',
    '''BLEURT-20-D6''': '''https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D6.zip''',
    '''BLEURT-20-D12''': '''https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D12.zip''',
    '''BLEURT-20''': '''https://storage.googleapis.com/bleurt-oss-21/BLEURT-20.zip''',
}
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION     ,     _KWARGS_DESCRIPTION				)
class 	lowerCamelCase_       (				datasets.Metric				):



					'''simple docstring'''







					def A      (				self						)   ->			Union[str, Any]:







								'''simple docstring'''

								return datasets.MetricInfo(
								    description=_DESCRIPTION ,       citation=_CITATION ,       homepage='''https://github.com/google-research/bleurt''' ,       inputs_description=_KWARGS_DESCRIPTION ,       features=datasets.Features(
								        {
								            '''predictions''': datasets.Value('''string''' ,       id='''sequence'''						),
								            '''references''': datasets.Value('''string''' ,       id='''sequence'''						),
								        }						) ,       codebase_urls=['''https://github.com/google-research/bleurt'''] ,       reference_urls=['''https://github.com/google-research/bleurt''', '''https://arxiv.org/abs/2004.04696'''] ,       )






					def A      (				self ,       snake_case_						)   ->			List[str]:







								'''simple docstring'''

								if self.config_name == "default":
											logger.warning(
											    '''Using default BLEURT-Base checkpoint for sequence maximum length 128. '''
											    '''You can use a bigger model for better results with e.g.: datasets.load_metric(\'bleurt\', \'bleurt-large-512\').'''						)
											__lowercase 					=   '''bleurt-base-128'''

								if self.config_name.lower() in CHECKPOINT_URLS:
											__lowercase 					=   self.config_name.lower()

								elif self.config_name.upper() in CHECKPOINT_URLS:
											__lowercase 					=   self.config_name.upper()

								else:
											raise KeyError(
											    F'{self.config_name} model not found. You should supply the name of a model checkpoint for bleurt in {CHECKPOINT_URLS.keys()}'						)

								# download the model checkpoint specified by self.config_name and set up the scorer
								__lowercase 					=   dl_manager.download_and_extract(CHECKPOINT_URLS[checkpoint_name]						)
								__lowercase 					=   score.BleurtScorer(os.path.join(snake_case_ ,       snake_case_						)						)






					def A      (				self ,       snake_case_ ,       snake_case_						)   ->			Tuple:







								'''simple docstring'''

								__lowercase 					=   self.scorer.score(references=snake_case_ ,       candidates=snake_case_						)
								return {"scores": scores}


