






'''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_float32 (there's also the fix_lavis branch)
# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
# same for Vicuna-13b
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipImageProcessor,
    InstructBlipConfig,
    InstructBlipForConditionalGeneration,
    InstructBlipProcessor,
    InstructBlipQFormerConfig,
    InstructBlipVisionConfig,
    LlamaConfig,
    LlamaTokenizerFast,
    TaConfig,
    TaTokenizerFast,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
def __UpperCAmelCase   (					)    ->   Tuple:
			UpperCAmelCase_      :       Dict          =       'https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'
			UpperCAmelCase_      :       Optional[Any]          =       Image.open(requests.get(lowerCamelCase_			,       stream=lowerCamelCase_		).raw		).convert('''RGB'''		)

			return image
def __UpperCAmelCase   (					A					:			int		)    ->   Any:
			UpperCAmelCase_      :       int          =       []
			# 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.embeddings.layernorm.weight''')		)
			rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.bias''', '''qformer.embeddings.layernorm.bias''')		)

			# fmt: on
			return rename_keys
def __UpperCAmelCase   (					A					:			Tuple			,       A					:			Tuple			,       A					:			Dict		)    ->   Any:
			UpperCAmelCase_      :       List[Any]          =       dct.pop(lowerCamelCase_		)
			UpperCAmelCase_      :       List[Any]          =       val
def __UpperCAmelCase   (					A					:			int			,       A					:			Optional[Any]		)    ->   Optional[Any]:





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

						# next, set bias in the state dict
						UpperCAmelCase_      :       Optional[Any]          =       torch.cat((q_bias, torch.zeros_like(lowerCamelCase_			,       requires_grad=lowerCamelCase_		), v_bias)		)
						UpperCAmelCase_      :       List[str]          =       qkv_bias
def __UpperCAmelCase   (					A					:			Dict		)    ->   Any:
			UpperCAmelCase_      :       Optional[Any]          =       3_6_4 if 'coco' in model_name else 2_2_4
			UpperCAmelCase_      :       Optional[int]          =       InstructBlipVisionConfig(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 "t5-xl" in model_name:
						UpperCAmelCase_      :       Union[str, 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_      :       Optional[Any]          =       TaConfig.from_pretrained('''google/flan-t5-xxl'''			,       dense_act_fn='''gelu'''			,       bos_token_id=1		).to_dict()
			elif "vicuna-7b" in model_name:
						UpperCAmelCase_      :       List[Any]          =       LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''			,       vocab_size=3_2_0_0_1		).to_dict()
			elif "vicuna-13b" in model_name:
						UpperCAmelCase_      :       Dict          =       LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''			,       vocab_size=3_2_0_0_1		).to_dict()
			else:
						raise ValueError('''Model name not supported'''		)

			# the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
			UpperCAmelCase_      :       Optional[int]          =       InstructBlipQFormerConfig(vocab_size=3_0_5_2_3		).to_dict()
			UpperCAmelCase_      :       List[Any]          =       InstructBlipConfig(vision_config=lowerCamelCase_			,       text_config=lowerCamelCase_			,       qformer_config=lowerCamelCase_		)

			return config, image_size
@torch.no_grad()
def __UpperCAmelCase   (					A					:			Dict			,       A					:			Any=None			,       A					:			Any=False		)    ->   Union[str, Any]:
			UpperCAmelCase_      :       Any          =       AutoTokenizer.from_pretrained('''bert-base-uncased'''			,       truncation_side='''left'''		)
			qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''}		)

			if "t5" in model_name:
						UpperCAmelCase_      :       Optional[int]          =       TaTokenizerFast.from_pretrained('''google/flan-t5-xl'''			,       truncation_side='''left'''		)
			elif "vicuna" in model_name:
						# the following was used in the original implementation:
						# tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
						# tokenizer.add_special_tokens({"pad_token": "[PAD]"})
						# tokenizer.add_special_tokens({"bos_token": "</s>"})
						# tokenizer.add_special_tokens({"eos_token": "</s>"})
						# tokenizer.add_special_tokens({"unk_token": "</s>"})
						UpperCAmelCase_      :       List[Any]          =       LlamaTokenizerFast.from_pretrained(
						    '''huggyllama/llama-7b'''			,       truncation_side='''left'''			,       bos_token='''</s>'''			,       unk_token='''</s>'''		)
						tokenizer.add_special_tokens({'''pad_token''': '''[PAD]'''}		)

			UpperCAmelCase_      :       Dict          =       get_blipa_config(lowerCamelCase_		)
			UpperCAmelCase_      :       int          =       InstructBlipForConditionalGeneration(lowerCamelCase_		).eval()

			UpperCAmelCase_      :       str          =       {
			    'instructblip-vicuna-7b': ('blip2_vicuna_instruct', 'vicuna7b'),
			    'instructblip-vicuna-13b': ('blip2_vicuna_instruct', 'vicuna13b'),
			    'instructblip-flan-t5-xl': ('blip2_t5_instruct', 'flant5xl'),
			    'instructblip-flan-t5-xxl': ('blip2_t5_instruct', 'flant5xxl'),
			}

			UpperCAmelCase_      :       Optional[int]          =       model_name_to_original[model_name]

			# load original model
			print('''Loading original model...'''		)
			UpperCAmelCase_      :       Optional[int]          =       'cuda:1' if torch.cuda.is_available() else 'cpu'
			UpperCAmelCase_      :       Any          =       'cuda:2' if torch.cuda.is_available() else 'cpu'
			UpperCAmelCase_      :       Optional[Any]          =       load_model_and_preprocess(
			    name=lowerCamelCase_			,       model_type=lowerCamelCase_			,       is_eval=lowerCamelCase_			,       device=lowerCamelCase_		)
			original_model.eval()
			print('''Done!'''		)

			# update state dict keys
			UpperCAmelCase_      :       Optional[int]          =       original_model.state_dict()
			UpperCAmelCase_      :       Dict          =       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():
						UpperCAmelCase_      :       Optional[int]          =       state_dict.pop(lowerCamelCase_		)
						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 "llm_proj" in key:
									UpperCAmelCase_      :       Any          =       key.replace('''llm_proj'''			,       '''language_projection'''		)
						if "t5_proj" in key:
									UpperCAmelCase_      :       str          =       key.replace('''t5_proj'''			,       '''language_projection'''		)
						if key.startswith('''llm_model'''		):
									UpperCAmelCase_      :       Any          =       key.replace('''llm_model'''			,       '''language_model'''		)
						if key.startswith('''t5'''		):
									UpperCAmelCase_      :       Union[str, Any]          =       key.replace('''t5'''			,       '''language'''		)
						UpperCAmelCase_      :       Optional[int]          =       val

			# read in qv biases
			read_in_q_v_bias(lowerCamelCase_			,       lowerCamelCase_		)

			# note: weights get loaded in torch.float32 by default
			hf_model.load_state_dict(lowerCamelCase_			,       strict=lowerCamelCase_		)

			UpperCAmelCase_      :       str          =       load_demo_image()
			UpperCAmelCase_      :       Any          =       'What is unusual about this image?'

			# create processor
			UpperCAmelCase_      :       Dict          =       BlipImageProcessor(
			    size={'''height''': image_size, '''width''': image_size}			,       image_mean=lowerCamelCase_			,       image_std=lowerCamelCase_		)
			UpperCAmelCase_      :       Dict          =       InstructBlipProcessor(
			    image_processor=lowerCamelCase_			,       tokenizer=lowerCamelCase_			,       qformer_tokenizer=lowerCamelCase_			,       )
			UpperCAmelCase_      :       List[Any]          =       processor(images=lowerCamelCase_			,       text=lowerCamelCase_			,       return_tensors='''pt'''		).to(lowerCamelCase_		)

			# make sure processor creates exact same pixel values
			UpperCAmelCase_      :       str          =       vis_processors['eval'](lowerCamelCase_		).unsqueeze(0		).to(lowerCamelCase_		)
			UpperCAmelCase_      :       List[str]          =       inputs.pixel_values
			assert torch.allclose(original_pixel_values.to(pixel_values.device		)			,       lowerCamelCase_		)

			original_model.to(lowerCamelCase_		)
			hf_model.to(lowerCamelCase_		)
			with torch.no_grad():
						if "vicuna" in model_name:
									UpperCAmelCase_      :       str          =       original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}		).logits
									UpperCAmelCase_      :       Tuple          =       hf_model(**lowerCamelCase_		).logits
						else:
									UpperCAmelCase_      :       Any          =       original_model(
									    {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}		).logits
									UpperCAmelCase_      :       str          =       tokenizer('''\n'''			,       return_tensors='''pt'''		).input_ids.to(lowerCamelCase_		)
									UpperCAmelCase_      :       str          =       label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id			,       -1_0_0		)
									UpperCAmelCase_      :       Dict          =       hf_model(**lowerCamelCase_			,       labels=lowerCamelCase_		).logits

			print('''First values of original logits:'''			,       original_logits[0, :3, :3]		)
			print('''First values of HF logits:'''			,       logits[0, :3, :3]		)

			# assert values
			assert original_logits.shape == logits.shape
			UpperCAmelCase_      :       List[Any]          =       1e-4 if 'vicuna' in model_name else 1e-5
			assert torch.allclose(original_logits.to(logits.device		)			,       lowerCamelCase_			,       atol=lowerCamelCase_		)
			print('''Looks ok!'''		)

			print('''Generating with original model...'''		)
			UpperCAmelCase_      :       List[Any]          =       original_model.generate({'''image''': original_pixel_values, '''prompt''': prompt}			,       num_beams=5		)

			# important: we need to cast the weights of the HF model to the appropriate type
			print('''Generating with HF model...'''		)
			UpperCAmelCase_      :       Union[str, Any]          =       hf_model.generate(
			    **lowerCamelCase_			,       do_sample=lowerCamelCase_			,       num_beams=5			,       max_length=2_5_6			,       min_length=1			,       top_p=0.9			,       repetition_penalty=1.5			,       length_penalty=1.0			,       temperature=1			,       )
			if "vicuna" in model_name:
						# convert output id 0 to 2 (eos_token_id)
						# TODO add this in the generate method?
						UpperCAmelCase_      :       Optional[int]          =       2
			print('''Original generation:'''			,       lowerCamelCase_		)
			UpperCAmelCase_      :       List[Any]          =       processor.batch_decode(lowerCamelCase_			,       skip_special_tokens=lowerCamelCase_		)
			UpperCAmelCase_      :       Optional[Any]          =       [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"Salesforce/{model_name}"		)
						hf_model.push_to_hub(F"Salesforce/{model_name}"		)


if __name__ == "__main__":
						_UpperCamelCase						:						int  =					argparse.ArgumentParser()
						_UpperCamelCase						:						int  =					[
						    '''instructblip-vicuna-7b''',
						    '''instructblip-vicuna-13b''',
						    '''instructblip-flan-t5-xl''',
						    '''instructblip-flan-t5-xxl''',
						]
						parser.add_argument(
						    '--model_name',
						    default='instructblip-flan-t5-xl',
						    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)








import os

import pytest

from datasets import (
    get_dataset_config_info,
    get_dataset_config_names,
    get_dataset_infos,
    get_dataset_split_names,
    inspect_dataset,
    inspect_metric,
)


UpperCamelCase__ :   List[Any]						=							pytest.mark.integration





@pytest.mark.parametrize('path'  ,							['paws', 'csv']  )
def        __UpperCAmelCase						(							lowerCamelCase_ :     Any  ,							lowerCamelCase_ :     int  )       ->  Any:







			"""simple docstring"""





			inspect_dataset(lowerCamelCase_  ,							lowerCamelCase_  )
			SCREAMING_SNAKE_CASE_   : List[str]					  =      path + '.py'
			assert script_name in os.listdir(lowerCamelCase_  )
			assert "__pycache__" not in os.listdir(lowerCamelCase_  )





@pytest.mark.filterwarnings('ignore:inspect_metric is deprecated:FutureWarning'  )
@pytest.mark.filterwarnings('ignore:metric_module_factory is deprecated:FutureWarning'  )
@pytest.mark.parametrize('path'  ,							['accuracy']  )
def        __UpperCAmelCase						(							lowerCamelCase_ :     Optional[int]  ,							lowerCamelCase_ :     Dict  )       ->  Optional[Any]:







			"""simple docstring"""





			inspect_metric(lowerCamelCase_  ,							lowerCamelCase_  )
			SCREAMING_SNAKE_CASE_   : Tuple					  =      path + '.py'
			assert script_name in os.listdir(lowerCamelCase_  )
			assert "__pycache__" not in os.listdir(lowerCamelCase_  )





@pytest.mark.parametrize(
    'path, config_name, expected_splits'  ,							[
        ('squad', 'plain_text', ['train', 'validation']),
        ('dalle-mini/wit', 'dalle-mini--wit', ['train']),
        ('paws', 'labeled_final', ['train', 'test', 'validation']),
    ]  ,							)
def        __UpperCAmelCase						(							lowerCamelCase_ :     List[str]  ,							lowerCamelCase_ :     Dict  ,							lowerCamelCase_ :     Optional[Any]  )       ->  Optional[int]:







			"""simple docstring"""





			SCREAMING_SNAKE_CASE_   : int					  =      get_dataset_config_info(lowerCamelCase_  ,							config_name=lowerCamelCase_  )
			assert info.config_name == config_name
			assert list(info.splits.keys()  ) == expected_splits





@pytest.mark.parametrize(
    'path, config_name, expected_exception'  ,							[
        ('paws', None, ValueError),
    ]  ,							)
def        __UpperCAmelCase						(							lowerCamelCase_ :     Union[str, Any]  ,							lowerCamelCase_ :     Optional[int]  ,							lowerCamelCase_ :     int  )       ->  List[Any]:







			"""simple docstring"""





			with pytest.raises(lowerCamelCase_  ):
						get_dataset_config_info(lowerCamelCase_  ,							config_name=lowerCamelCase_  )





@pytest.mark.parametrize(
    'path, expected'  ,							[
        ('squad', 'plain_text'),
        ('acronym_identification', 'default'),
        ('lhoestq/squad', 'plain_text'),
        ('lhoestq/test', 'default'),
        ('lhoestq/demo1', 'lhoestq--demo1'),
        ('dalle-mini/wit', 'dalle-mini--wit'),
    ]  ,							)
def        __UpperCAmelCase						(							lowerCamelCase_ :     Tuple  ,							lowerCamelCase_ :     List[Any]  )       ->  Dict:







			"""simple docstring"""





			SCREAMING_SNAKE_CASE_   : str					  =      get_dataset_config_names(lowerCamelCase_  )
			assert expected in config_names





@pytest.mark.parametrize(
    'path, expected_configs, expected_splits_in_first_config'  ,							[
        ('squad', ['plain_text'], ['train', 'validation']),
        ('dalle-mini/wit', ['dalle-mini--wit'], ['train']),
        ('paws', ['labeled_final', 'labeled_swap', 'unlabeled_final'], ['train', 'test', 'validation']),
    ]  ,							)
def        __UpperCAmelCase						(							lowerCamelCase_ :     List[Any]  ,							lowerCamelCase_ :     int  ,							lowerCamelCase_ :     Union[str, Any]  )       ->  Optional[int]:







			"""simple docstring"""





			SCREAMING_SNAKE_CASE_   : str					  =      get_dataset_infos(lowerCamelCase_  )
			assert list(infos.keys()  ) == expected_configs
			SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      expected_configs[0]
			assert expected_config in infos
			SCREAMING_SNAKE_CASE_   : Any					  =      infos[expected_config]
			assert info.config_name == expected_config
			assert list(info.splits.keys()  ) == expected_splits_in_first_config





@pytest.mark.parametrize(
    'path, expected_config, expected_splits'  ,							[
        ('squad', 'plain_text', ['train', 'validation']),
        ('dalle-mini/wit', 'dalle-mini--wit', ['train']),
        ('paws', 'labeled_final', ['train', 'test', 'validation']),
    ]  ,							)
def        __UpperCAmelCase						(							lowerCamelCase_ :     Union[str, Any]  ,							lowerCamelCase_ :     List[str]  ,							lowerCamelCase_ :     str  )       ->  str:







			"""simple docstring"""





			SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      get_dataset_infos(lowerCamelCase_  )
			assert expected_config in infos
			SCREAMING_SNAKE_CASE_   : Tuple					  =      infos[expected_config]
			assert info.config_name == expected_config
			assert list(info.splits.keys()  ) == expected_splits





@pytest.mark.parametrize(
    'path, config_name, expected_exception'  ,							[
        ('paws', None, ValueError),
    ]  ,							)
def        __UpperCAmelCase						(							lowerCamelCase_ :     int  ,							lowerCamelCase_ :     int  ,							lowerCamelCase_ :     int  )       ->  int:







			"""simple docstring"""





			with pytest.raises(lowerCamelCase_  ):
						get_dataset_split_names(lowerCamelCase_  ,							config_name=lowerCamelCase_  )



