



"""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 					lowerCamelCase_		():
    _UpperCAmelCase				:    int											=       '''https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'''
    _UpperCAmelCase				:    Tuple											=       Image.open(requests.get(snake_case__	,			stream=snake_case__				).raw				).convert('''RGB'''				)

    return image





def 					lowerCamelCase_		(UpperCamelCase__						:						Optional[int]				):
    _UpperCAmelCase				:    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.embeddings.layernorm.weight''')				)
    rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.bias''', '''qformer.embeddings.layernorm.bias''')				)

    # fmt: on
    return rename_keys





def 					lowerCamelCase_		(UpperCamelCase__						:						Dict	,			UpperCamelCase__						:						Tuple	,			UpperCamelCase__						:						Optional[Any]				):
    _UpperCAmelCase				:    str											=       dct.pop(snake_case__				)
    _UpperCAmelCase				:    Optional[Any]											=       val





def 					lowerCamelCase_		(UpperCamelCase__						:						List[str]	,			UpperCamelCase__						:						Optional[Any]				):
    for i in range(config.vision_config.num_hidden_layers				):
        # read in original q and v biases
        _UpperCAmelCase				:    List[Any]											=       state_dict.pop(F'visual_encoder.blocks.{i}.attn.q_bias'				)
        _UpperCAmelCase				:    int											=       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(snake_case__	,			requires_grad=snake_case__				), v_bias)				)
        _UpperCAmelCase				:    Any											=       qkv_bias





def 					lowerCamelCase_		(UpperCamelCase__						:						Tuple				):
    _UpperCAmelCase				:    Optional[Any]											=       364 if '''coco''' in model_name else 224
    _UpperCAmelCase				:    Any											=       InstructBlipVisionConfig(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 "t5-xl" in model_name:
        _UpperCAmelCase				:    List[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				:    Tuple											=       TaConfig.from_pretrained('''google/flan-t5-xxl'''	,			dense_act_fn='''gelu'''	,			bos_token_id=1				).to_dict()
    elif "vicuna-7b" in model_name:
        _UpperCAmelCase				:    Dict											=       LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''	,			vocab_size=3_2001				).to_dict()
    elif "vicuna-13b" in model_name:
        _UpperCAmelCase				:    int											=       LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''	,			vocab_size=3_2001				).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				:    Union[str, Any]											=       InstructBlipQFormerConfig(vocab_size=3_0523				).to_dict()
    _UpperCAmelCase				:    int											=       InstructBlipConfig(vision_config=snake_case__	,			text_config=snake_case__	,			qformer_config=snake_case__				)

    return config, image_size





@torch.no_grad()
def 					lowerCamelCase_		(UpperCamelCase__						:						Tuple	,			UpperCamelCase__						:						List[Any]=None	,			UpperCamelCase__						:						List[Any]=False				):
    _UpperCAmelCase				:    List[Any]											=       AutoTokenizer.from_pretrained('''bert-base-uncased'''	,			truncation_side='''left'''				)
    qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''}				)

    if "t5" in model_name:
        _UpperCAmelCase				:    Tuple											=       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				:    Dict											=       LlamaTokenizerFast.from_pretrained(
            '''huggyllama/llama-7b'''	,			truncation_side='''left'''	,			bos_token='''</s>'''	,			unk_token='''</s>'''				)
        tokenizer.add_special_tokens({'''pad_token''': '''[PAD]'''}				)

    _UpperCAmelCase				,							_UpperCAmelCase				:    int											=       get_blipa_config(snake_case__				)
    _UpperCAmelCase				:    List[str]											=       InstructBlipForConditionalGeneration(snake_case__				).eval()

    _UpperCAmelCase				:    List[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				,							_UpperCAmelCase				:    str											=       model_name_to_original[model_name]

    # load original model
    print('''Loading original model...'''				)
    _UpperCAmelCase				:    Any											=       '''cuda:1''' if torch.cuda.is_available() else '''cpu'''
    _UpperCAmelCase				:    int											=       '''cuda:2''' if torch.cuda.is_available() else '''cpu'''
    _UpperCAmelCase				,							_UpperCAmelCase				,							_UpperCAmelCase				:    Optional[int]											=       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
    _UpperCAmelCase				:    Optional[int]											=       original_model.state_dict()
    _UpperCAmelCase				:    Optional[int]											=       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():
        _UpperCAmelCase				:    Tuple											=       state_dict.pop(snake_case__				)
        if key.startswith('''Qformer.bert'''				):
            _UpperCAmelCase				:    Any											=       key.replace('''Qformer.bert'''	,			'''qformer'''				)
        if "attention.self" in key:
            _UpperCAmelCase				:    Dict											=       key.replace('''self'''	,			'''attention'''				)
        if "llm_proj" in key:
            _UpperCAmelCase				:    Optional[Any]											=       key.replace('''llm_proj'''	,			'''language_projection'''				)
        if "t5_proj" in key:
            _UpperCAmelCase				:    Optional[int]											=       key.replace('''t5_proj'''	,			'''language_projection'''				)
        if key.startswith('''llm_model'''				):
            _UpperCAmelCase				:    Tuple											=       key.replace('''llm_model'''	,			'''language_model'''				)
        if key.startswith('''t5'''				):
            _UpperCAmelCase				:    int											=       key.replace('''t5'''	,			'''language'''				)
        _UpperCAmelCase				:    Optional[int]											=       val

    # read in qv biases
    read_in_q_v_bias(snake_case__	,			snake_case__				)

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

    _UpperCAmelCase				:    Dict											=       load_demo_image()
    _UpperCAmelCase				:    List[Any]											=       '''What is unusual about this image?'''

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

    # make sure processor creates exact same pixel values
    _UpperCAmelCase				:    Optional[Any]											=       vis_processors['''eval'''](snake_case__				).unsqueeze(0				).to(snake_case__				)
    _UpperCAmelCase				:    int											=       inputs.pixel_values
    assert torch.allclose(original_pixel_values.to(pixel_values.device				)	,			snake_case__				)

    original_model.to(snake_case__				)
    hf_model.to(snake_case__				)
    with torch.no_grad():
        if "vicuna" in model_name:
            _UpperCAmelCase				:    List[str]											=       original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}				).logits
            _UpperCAmelCase				:    Union[str, Any]											=       hf_model(**snake_case__				).logits
        else:
            _UpperCAmelCase				:    int											=       original_model(
                {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}				).logits
            _UpperCAmelCase				:    Optional[int]											=       tokenizer('''\n'''	,			return_tensors='''pt'''				).input_ids.to(snake_case__				)
            _UpperCAmelCase				:    Optional[Any]											=       label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id	,			-100				)
            _UpperCAmelCase				:    Optional[int]											=       hf_model(**snake_case__	,			labels=snake_case__				).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[str]											=       1E-4 if '''vicuna''' in model_name else 1E-5
    assert torch.allclose(original_logits.to(logits.device				)	,			snake_case__	,			atol=snake_case__				)
    print('''Looks ok!'''				)

    print('''Generating with original model...'''				)
    _UpperCAmelCase				:    Optional[int]											=       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				:    Optional[int]											=       hf_model.generate(
        **snake_case__	,			do_sample=snake_case__	,			num_beams=5	,			max_length=256	,			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				:    Dict											=       2
    print('''Original generation:'''	,			snake_case__				)
    _UpperCAmelCase				:    Tuple											=       processor.batch_decode(snake_case__	,			skip_special_tokens=snake_case__				)
    _UpperCAmelCase				:    Optional[int]											=       [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'Salesforce/{model_name}'				)
        hf_model.push_to_hub(F'Salesforce/{model_name}'				)


if __name__ == "__main__":
  _lowerCAmelCase       :List[Any]	=   argparse.ArgumentParser()
  _lowerCAmelCase       :Optional[Any]	=   [
      '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',
  )

  _lowerCAmelCase       :Optional[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.models.biogpt.tokenization_biogpt import VOCAB_FILES_NAMES, BioGptTokenizer
from transformers.testing_utils import slow

from ...test_tokenization_common import TokenizerTesterMixin







class 					lowerCAmelCase_					(   _lowercase		,       unittest.TestCase     ):



						'''simple docstring'''


						_lowerCamelCase:     List[str]   =     BioGptTokenizer
						_lowerCamelCase:     Tuple   =     False


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[str]   )  -> Optional[int]:
							super().setUp()

							# Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
							A     = [
							    'l',
							    'o',
							    'w',
							    'e',
							    'r',
							    's',
							    't',
							    'i',
							    'd',
							    'n',
							    'w</w>',
							    'r</w>',
							    't</w>',
							    'lo',
							    'low',
							    'er</w>',
							    'low</w>',
							    'lowest</w>',
							    'newer</w>',
							    'wider</w>',
							    '<unk>',
							]
							A     = dict(zip(A_	,range(len(A_   )   )   )   )
							A     = ['l o 123', 'lo w 1456', 'e r</w> 1789', '']

							A     = os.path.join(self.tmpdirname	,VOCAB_FILES_NAMES['vocab_file']   )
							A     = os.path.join(self.tmpdirname	,VOCAB_FILES_NAMES['merges_file']   )
							with open(self.vocab_file	,'w'   ) as fp:
								fp.write(json.dumps(A_   )   )
							with open(self.merges_file	,'w'   ) as fp:
								fp.write('\n'.join(A_   )   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   Dict	,A_  :   Tuple   )  -> int:
							A     = 'lower newer'
							A     = 'lower newer'
							return input_text, output_text


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[Any]   )  -> Any:
							A     = BioGptTokenizer(self.vocab_file	,self.merges_file   )

							A     = 'lower'
							A     = ['low', 'er</w>']
							A     = tokenizer.tokenize(A_   )
							self.assertListEqual(A_	,A_   )

							A     = tokens + ['<unk>']
							A     = [14, 15, 20]
							self.assertListEqual(tokenizer.convert_tokens_to_ids(A_   )	,A_   )




						@slow
						def     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]   )  -> Optional[Any]:
							A     = BioGptTokenizer.from_pretrained('microsoft/biogpt'   )

							A     = tokenizer.encode('sequence builders'	,add_special_tokens=A_   )
							A     = tokenizer.encode('multi-sequence build'	,add_special_tokens=A_   )

							A     = tokenizer.build_inputs_with_special_tokens(A_   )
							A     = tokenizer.build_inputs_with_special_tokens(A_	,A_   )

							self.assertTrue(encoded_sentence == [2] + text   )
							self.assertTrue(encoded_pair == [2] + text + [2] + text_a   )