

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

	return image

def     __SCREAMING_SNAKE_CASE    ( A_					):
	lowerCAmelCase__      :		Optional[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     __SCREAMING_SNAKE_CASE    ( A_		,	A_		,	A_					):
	lowerCAmelCase__      :		Dict    							=      dct.pop(A_					)
	lowerCAmelCase__      :		int    							=      val

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	for i in range(config.vision_config.num_hidden_layers					):
		# read in original q and v biases
		lowerCAmelCase__      :		List[str]    							=      state_dict.pop(f'visual_encoder.blocks.{i}.attn.q_bias'					)
		lowerCAmelCase__      :		Tuple    							=      state_dict.pop(f'visual_encoder.blocks.{i}.attn.v_bias'					)

		# next, set bias in the state dict
		lowerCAmelCase__      :		List[str]    							=      torch.cat((q_bias, torch.zeros_like(A_		,	requires_grad=A_					), v_bias)					)
		lowerCAmelCase__      :		str    							=      qkv_bias

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	lowerCAmelCase__      :		Any    							=      3_64 if '''coco''' in model_name else 2_24
	lowerCAmelCase__      :		List[str]    							=      BlipaVisionConfig(image_size=A_					).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__      :		List[str]    							=      OPTConfig.from_pretrained('''facebook/opt-2.7b'''		,	eos_token_id=A_					).to_dict()
	elif "opt-6.7b" in model_name:
		lowerCAmelCase__      :		Dict    							=      OPTConfig.from_pretrained('''facebook/opt-6.7b'''		,	eos_token_id=A_					).to_dict()
	elif "t5-xl" in model_name:
		lowerCAmelCase__      :		int    							=      TaConfig.from_pretrained('''google/flan-t5-xl'''		,	dense_act_fn='''gelu'''		,	bos_token_id=1					).to_dict()
	elif "t5-xxl" in model_name:
		lowerCAmelCase__      :		List[Any]    							=      TaConfig.from_pretrained('''google/flan-t5-xxl'''		,	dense_act_fn='''gelu'''		,	bos_token_id=1					).to_dict()

	lowerCAmelCase__      :		str    							=      BlipaConfig(vision_config=A_		,	text_config=A_					)

	return config, image_size

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

	lowerCAmelCase__      :		List[Any]    							=      BlipaForConditionalGeneration(A_					).eval()

	lowerCAmelCase__      :		str    							=      {
	    '''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__      :		Any    							=      model_name_to_original[model_name]

	# load original model
	print('''Loading original model...'''					)
	lowerCAmelCase__      :		Union[str, Any]    							=      '''cuda''' if torch.cuda.is_available() else '''cpu'''
	lowerCAmelCase__							,lowerCAmelCase__							,lowerCAmelCase__      :		Union[str, Any]    							=      load_model_and_preprocess(
	    name=A_		,	model_type=A_		,	is_eval=A_		,	device=A_					)
	original_model.eval()
	print('''Done!'''					)

	# update state dict keys
	lowerCAmelCase__      :		Optional[Any]    							=      original_model.state_dict()
	lowerCAmelCase__      :		List[Any]    							=      create_rename_keys(A_					)
	for src, dest in rename_keys:
		rename_key(A_		,	A_		,	A_					)

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

	# read in qv biases
	read_in_q_v_bias(A_		,	A_					)

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

	lowerCAmelCase__      :		Dict    							=      load_demo_image()
	lowerCAmelCase__      :		Optional[Any]    							=      vis_processors['''eval'''](A_					).unsqueeze(0					).to(A_					)
	lowerCAmelCase__      :		List[Any]    							=      tokenizer(['''\n''']		,	return_tensors='''pt'''					).input_ids.to(A_					)

	# create processor
	lowerCAmelCase__      :		Optional[Any]    							=      BlipImageProcessor(
	    size={'''height''': image_size, '''width''': image_size}		,	image_mean=A_		,	image_std=A_					)
	lowerCAmelCase__      :		List[str]    							=      BlipaProcessor(image_processor=A_		,	tokenizer=A_					)
	lowerCAmelCase__      :		Any    							=      processor(images=A_		,	return_tensors='''pt'''					).pixel_values.to(A_					)

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

	original_model.to(A_					)
	hf_model.to(A_					)
	with torch.no_grad():
		if "opt" in model_name:
			lowerCAmelCase__      :		Any    							=      original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}					).logits
			lowerCAmelCase__      :		List[Any]    							=      hf_model(A_		,	A_					).logits
		else:
			lowerCAmelCase__      :		Dict    							=      original_model(
			    {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}					).logits
			lowerCAmelCase__      :		List[Any]    							=      input_ids.masked_fill(input_ids == tokenizer.pad_token_id		,	-1_00					)
			lowerCAmelCase__      :		Dict    							=      hf_model(A_		,	A_		,	labels=A_					).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__      :		Optional[Any]    							=      torch.tensor(
		    [[-41.5_850, -4.4_440, -8.9_922], [-47.4_322, -5.9_143, -1.7_340]]		,	device=A_					)
		assert torch.allclose(logits[0, :3, :3]		,	A_		,	atol=1e-4					)
	elif model_name == "blip2-flan-t5-xl-coco":
		lowerCAmelCase__      :		Dict    							=      torch.tensor(
		    [[-57.0_109, -9.8_967, -12.6_280], [-68.6_578, -12.7_191, -10.5_065]]		,	device=A_					)
	else:
		# cast to same type
		lowerCAmelCase__      :		List[str]    							=      logits.dtype
		assert torch.allclose(original_logits.to(A_					)		,	A_		,	atol=1e-2					)
	print('''Looks ok!'''					)

	print('''Generating a caption...'''					)
	lowerCAmelCase__      :		Union[str, Any]    							=      ''''''
	lowerCAmelCase__      :		Union[str, Any]    							=      tokenizer(A_		,	return_tensors='''pt'''					).input_ids.to(A_					)

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

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

	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    : Optional[Any]				 =						argparse.ArgumentParser()
							__UpperCamelCase    : Tuple				 =						[
							    '''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    : Optional[Any]				 =						parser.parse_args()

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







"""simple docstring"""





import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
							from .tokenization_albert import AlbertTokenizer
else:
							__UpperCamelCase    : Any				 =						None

__UpperCamelCase    : Optional[int]				 =						logging.get_logger(__name__)
__UpperCamelCase    : List[Any]				 =						{'''vocab_file''': '''spiece.model''', '''tokenizer_file''': '''tokenizer.json'''}

__UpperCamelCase    : str				 =						{
    '''vocab_file''': {
        '''albert-base-v1''': '''https://huggingface.co/albert-base-v1/resolve/main/spiece.model''',
        '''albert-large-v1''': '''https://huggingface.co/albert-large-v1/resolve/main/spiece.model''',
        '''albert-xlarge-v1''': '''https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model''',
        '''albert-xxlarge-v1''': '''https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model''',
        '''albert-base-v2''': '''https://huggingface.co/albert-base-v2/resolve/main/spiece.model''',
        '''albert-large-v2''': '''https://huggingface.co/albert-large-v2/resolve/main/spiece.model''',
        '''albert-xlarge-v2''': '''https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model''',
        '''albert-xxlarge-v2''': '''https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model''',
    },
    '''tokenizer_file''': {
        '''albert-base-v1''': '''https://huggingface.co/albert-base-v1/resolve/main/tokenizer.json''',
        '''albert-large-v1''': '''https://huggingface.co/albert-large-v1/resolve/main/tokenizer.json''',
        '''albert-xlarge-v1''': '''https://huggingface.co/albert-xlarge-v1/resolve/main/tokenizer.json''',
        '''albert-xxlarge-v1''': '''https://huggingface.co/albert-xxlarge-v1/resolve/main/tokenizer.json''',
        '''albert-base-v2''': '''https://huggingface.co/albert-base-v2/resolve/main/tokenizer.json''',
        '''albert-large-v2''': '''https://huggingface.co/albert-large-v2/resolve/main/tokenizer.json''',
        '''albert-xlarge-v2''': '''https://huggingface.co/albert-xlarge-v2/resolve/main/tokenizer.json''',
        '''albert-xxlarge-v2''': '''https://huggingface.co/albert-xxlarge-v2/resolve/main/tokenizer.json''',
    },
}

__UpperCamelCase    : Any				 =						{
    '''albert-base-v1''': 5_1_2,
    '''albert-large-v1''': 5_1_2,
    '''albert-xlarge-v1''': 5_1_2,
    '''albert-xxlarge-v1''': 5_1_2,
    '''albert-base-v2''': 5_1_2,
    '''albert-large-v2''': 5_1_2,
    '''albert-xlarge-v2''': 5_1_2,
    '''albert-xxlarge-v2''': 5_1_2,
}

__UpperCamelCase    : Union[str, Any]				 =						'''▁'''







class 							SCREAMING_SNAKE_CASE   (						a_			):






					"""simple docstring"""



					lowercase__												=			VOCAB_FILES_NAMES
					lowercase__												=			PRETRAINED_VOCAB_FILES_MAP
					lowercase__												=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					lowercase__												=			AlbertTokenizer







					def __init__(			self     :							Union[str, Any]  ,lowercase_     :							List[str]=None  ,lowercase_     :							Union[str, Any]=None  ,lowercase_     :							str=True  ,lowercase_     :							Optional[Any]=True  ,lowercase_     :							str=False  ,lowercase_     :							Tuple="[CLS]"  ,lowercase_     :							Optional[int]="[SEP]"  ,lowercase_     :							Optional[Any]="<unk>"  ,lowercase_     :							List[Any]="[SEP]"  ,lowercase_     :							Optional[int]="<pad>"  ,lowercase_     :							List[Any]="[CLS]"  ,lowercase_     :							Optional[int]="[MASK]"  ,**lowercase_     :							Optional[int]  ,):
						# Mask token behave like a normal word, i.e. include the space before it and
						# is included in the raw text, there should be a match in a non-normalized sentence.
						lowerCAmelCase__      :		Optional[Any]    							=      (
						    AddedToken(lowercase_  ,lstrip=lowercase_  ,rstrip=lowercase_  ,normalized=lowercase_						)
						    if isinstance(lowercase_  ,lowercase_						)
						    else mask_token
						)

						super().__init__(
						    lowercase_  ,tokenizer_file=lowercase_  ,do_lower_case=lowercase_  ,remove_space=lowercase_  ,keep_accents=lowercase_  ,bos_token=lowercase_  ,eos_token=lowercase_  ,unk_token=lowercase_  ,sep_token=lowercase_  ,pad_token=lowercase_  ,cls_token=lowercase_  ,mask_token=lowercase_  ,**lowercase_  ,)

						lowerCAmelCase__      :		List[Any]    							=      do_lower_case
						lowerCAmelCase__      :		str    							=      remove_space
						lowerCAmelCase__      :		Optional[int]    							=      keep_accents
						lowerCAmelCase__      :		Dict    							=      vocab_file
						lowerCAmelCase__      :		Optional[int]    							=      False if not self.vocab_file else True







					def 			__lowerCAmelCase					(			self     :							Dict  ,lowercase_     :							List[int]  ,lowercase_     :							Optional[List[int]] = None						):
						lowerCAmelCase__      :		List[Any]    							=      [self.sep_token_id]
						lowerCAmelCase__      :		Dict    							=      [self.cls_token_id]
						if token_ids_a is None:
							return cls + token_ids_a + sep
						return cls + token_ids_a + sep + token_ids_a + sep







					def 			__lowerCAmelCase					(			self     :							Union[str, Any]  ,lowercase_     :							List[int]  ,lowercase_     :							Optional[List[int]] = None						):
						lowerCAmelCase__      :		List[str]    							=      [self.sep_token_id]
						lowerCAmelCase__      :		List[Any]    							=      [self.cls_token_id]

						if token_ids_a is None:
							return len(cls + token_ids_a + sep						) * [0]
						return len(cls + token_ids_a + sep						) * [0] + len(token_ids_a + sep						) * [1]







					def 			__lowerCAmelCase					(			self     :							Optional[int]  ,lowercase_     :							str  ,lowercase_     :							Optional[str] = None						):
						if not self.can_save_slow_tokenizer:
							raise ValueError(
							    '''Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '''
							    '''tokenizer.'''						)

						if not os.path.isdir(lowercase_						):
							logger.error(F'Vocabulary path ({save_directory}) should be a directory'						)
							return
						lowerCAmelCase__      :		Union[str, Any]    							=      os.path.join(
						    lowercase_  ,(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']						)

						if os.path.abspath(self.vocab_file						) != os.path.abspath(lowercase_						):
							copyfile(self.vocab_file  ,lowercase_						)

						return (out_vocab_file,)




