
"""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__	:      Union[str, Any]     							=					'https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'
    UpperCAmelCase__	:      Union[str, Any]     							=					Image.open(requests.get(_snake_case    ,stream=_snake_case ).raw ).convert('RGB' )

    return image
def 					lowerCamelCase   (     _snake_case ):
    UpperCAmelCase__	:      Tuple     							=					[]
    # 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   (     _snake_case    ,_snake_case    ,_snake_case ):
    UpperCAmelCase__	:      List[Any]     							=					dct.pop(_snake_case )
    UpperCAmelCase__	:      Union[str, Any]     							=					val
def 					lowerCamelCase   (     _snake_case    ,_snake_case ):
    for i in range(config.vision_config.num_hidden_layers ):
        # read in original q and v biases
        UpperCAmelCase__	:      Optional[Any]     							=					state_dict.pop(F'''visual_encoder.blocks.{i}.attn.q_bias''' )
        UpperCAmelCase__	:      Tuple     							=					state_dict.pop(F'''visual_encoder.blocks.{i}.attn.v_bias''' )

        # next, set bias in the state dict
        UpperCAmelCase__	:      Optional[int]     							=					torch.cat((q_bias, torch.zeros_like(_snake_case    ,requires_grad=_snake_case ), v_bias) )
        UpperCAmelCase__	:      str     							=					qkv_bias
def 					lowerCamelCase   (     _snake_case ):
    UpperCAmelCase__	:      List[str]     							=					364 if 'coco' in model_name else 224
    UpperCAmelCase__	:      int     							=					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__	:      Dict     							=					TaConfig.from_pretrained('google/flan-t5-xl'    ,dense_act_fn='gelu'    ,bos_token_id=1 ).to_dict()
    elif "t5-xxl" in model_name:
        UpperCAmelCase__	:      List[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__	:      Union[str, Any]     							=					LlamaConfig.from_pretrained('decapoda-research/llama-7b-hf'    ,vocab_size=32001 ).to_dict()
    elif "vicuna-13b" in model_name:
        UpperCAmelCase__	:      Optional[Any]     							=					LlamaConfig.from_pretrained('decapoda-research/llama-13b-hf'    ,vocab_size=32001 ).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__	:      List[str]     							=					InstructBlipQFormerConfig(vocab_size=30523 ).to_dict()
    UpperCAmelCase__	:      Optional[Any]     							=					InstructBlipConfig(vision_config=_snake_case    ,text_config=_snake_case    ,qformer_config=_snake_case )

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

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

    UpperCAmelCase__							,							UpperCAmelCase__	:      List[Any]     							=					get_blipa_config(_snake_case )
    UpperCAmelCase__	:      List[Any]     							=					InstructBlipForConditionalGeneration(_snake_case ).eval()

    UpperCAmelCase__	:      Optional[int]     							=					{
        '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__	:      Optional[Any]     							=					model_name_to_original[model_name]

    # load original model
    print('Loading original model...' )
    UpperCAmelCase__	:      List[Any]     							=					'cuda:1' if torch.cuda.is_available() else 'cpu'
    UpperCAmelCase__	:      str     							=					'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__	:      int     							=					original_model.state_dict()
    UpperCAmelCase__	:      List[Any]     							=					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__	:      Dict     							=					state_dict.pop(_snake_case )
        if key.startswith('Qformer.bert' ):
            UpperCAmelCase__	:      str     							=					key.replace('Qformer.bert'    ,'qformer' )
        if "attention.self" in key:
            UpperCAmelCase__	:      int     							=					key.replace('self'    ,'attention' )
        if "llm_proj" in key:
            UpperCAmelCase__	:      str     							=					key.replace('llm_proj'    ,'language_projection' )
        if "t5_proj" in key:
            UpperCAmelCase__	:      List[str]     							=					key.replace('t5_proj'    ,'language_projection' )
        if key.startswith('llm_model' ):
            UpperCAmelCase__	:      List[Any]     							=					key.replace('llm_model'    ,'language_model' )
        if key.startswith('t5' ):
            UpperCAmelCase__	:      Any     							=					key.replace('t5'    ,'language' )
        UpperCAmelCase__	:      Optional[Any]     							=					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__	:      List[str]     							=					load_demo_image()
    UpperCAmelCase__	:      List[Any]     							=					'What is unusual about this image?'

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

    # make sure processor creates exact same pixel values
    UpperCAmelCase__	:      Dict     							=					vis_processors['eval'](_snake_case ).unsqueeze(0 ).to(_snake_case )
    UpperCAmelCase__	:      Optional[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__	:      Tuple     							=					original_model({'image': original_pixel_values, 'text_input': [prompt]} ).logits
            UpperCAmelCase__	:      Union[str, Any]     							=					hf_model(**_snake_case ).logits
        else:
            UpperCAmelCase__	:      Optional[Any]     							=					original_model(
                {'image': original_pixel_values, 'text_input': [prompt], 'text_output': ['\n']} ).logits
            UpperCAmelCase__	:      Optional[Any]     							=					tokenizer('\n'    ,return_tensors='pt' ).input_ids.to(_snake_case )
            UpperCAmelCase__	:      Any     							=					label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id    ,-100 )
            UpperCAmelCase__	:      Dict     							=					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__	:      Union[str, 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__	:      Dict     							=					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__	:      Optional[int]     							=					2
    print('Original generation:'    ,_snake_case )
    UpperCAmelCase__	:      Optional[int]     							=					processor.batch_decode(_snake_case    ,skip_special_tokens=_snake_case )
    UpperCAmelCase__	:      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__":
    UpperCamelCase__          =	argparse.ArgumentParser()
    UpperCamelCase__          =	[
        '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__          =	parser.parse_args()

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



"""simple docstring"""



import unittest

from transformers import CamembertTokenizer, CamembertTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow
from transformers.utils import is_torch_available

from ...test_tokenization_common import TokenizerTesterMixin


UpperCamelCase__          =	get_tests_dir('fixtures/test_sentencepiece.model')
UpperCamelCase__          =	get_tests_dir('fixtures/test_sentencepiece_bpe.model')

UpperCamelCase__          =	'pt' if is_torch_available() else 'tf'





@require_sentencepiece
@require_tokenizers
class 				a			(					lowercase						,		unittest.TestCase      ):
      UpperCamelCase							:   Optional[int]           =     CamembertTokenizer
      UpperCamelCase							:   Optional[int]           =     CamembertTokenizerFast
      UpperCamelCase							:   str           =     True
      UpperCamelCase							:   Tuple           =     True



      def  __snake_case							(   self      ):
          super().setUp()

          # We have a SentencePiece fixture for testing
          UpperCAmelCase__	:      int     							=					CamembertTokenizer(UpperCamelCase_      )
          tokenizer.save_pretrained(self.tmpdirname      )



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      Optional[int]     							=					'<pad>'
          UpperCAmelCase__	:      Tuple     							=					1

          self.assertEqual(self.get_tokenizer()._convert_token_to_id(UpperCamelCase_      )     ,	UpperCamelCase_      )
          self.assertEqual(self.get_tokenizer()._convert_id_to_token(UpperCamelCase_      )     ,	UpperCamelCase_      )



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      Optional[Any]     							=					list(self.get_tokenizer().get_vocab().keys()      )

          self.assertEqual(vocab_keys[0]     ,	'<s>NOTUSED'      )
          self.assertEqual(vocab_keys[1]     ,	'<pad>'      )
          self.assertEqual(vocab_keys[-1]     ,	'<mask>'      )
          self.assertEqual(len(UpperCamelCase_      )     ,	1_004      )



      def  __snake_case							(   self      ):
          self.assertEqual(self.get_tokenizer().vocab_size     ,	1_005      )



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      Dict     							=					CamembertTokenizer(UpperCamelCase_      )
          tokenizer.save_pretrained(self.tmpdirname      )
          UpperCAmelCase__	:      str     							=					CamembertTokenizerFast.from_pretrained(self.tmpdirname      )

          UpperCAmelCase__	:      Union[str, Any]     							=					'I was born in 92000, and this is falsé.'

          UpperCAmelCase__	:      Optional[Any]     							=					tokenizer.encode(UpperCamelCase_      )
          UpperCAmelCase__	:      Dict     							=					rust_tokenizer.encode(UpperCamelCase_      )
          self.assertListEqual(UpperCamelCase_     ,	UpperCamelCase_      )

          UpperCAmelCase__	:      str     							=					tokenizer.encode(UpperCamelCase_     ,	add_special_tokens=UpperCamelCase_      )
          UpperCAmelCase__	:      Union[str, Any]     							=					rust_tokenizer.encode(UpperCamelCase_     ,	add_special_tokens=UpperCamelCase_      )
          self.assertListEqual(UpperCamelCase_     ,	UpperCamelCase_      )

          # <unk> tokens are not the same for `rust` than for `slow`.
          # Because spm gives back raw token instead of `unk` in EncodeAsPieces
          # tokens = tokenizer.tokenize(sequence)
          UpperCAmelCase__	:      List[Any]     							=					tokenizer.convert_ids_to_tokens(UpperCamelCase_      )
          UpperCAmelCase__	:      str     							=					rust_tokenizer.tokenize(UpperCamelCase_      )
          self.assertListEqual(UpperCamelCase_     ,	UpperCamelCase_      )



      def  __snake_case							(   self      ):
          if not self.test_rust_tokenizer:
              return

          UpperCAmelCase__	:      Any     							=					self.get_tokenizer()
          UpperCAmelCase__	:      List[str]     							=					self.get_rust_tokenizer()

          UpperCAmelCase__	:      List[str]     							=					'I was born in 92000, and this is falsé.'

          UpperCAmelCase__	:      List[str]     							=					tokenizer.tokenize(UpperCamelCase_      )
          UpperCAmelCase__	:      str     							=					rust_tokenizer.tokenize(UpperCamelCase_      )
          self.assertListEqual(UpperCamelCase_     ,	UpperCamelCase_      )

          UpperCAmelCase__	:      Optional[Any]     							=					tokenizer.encode(UpperCamelCase_     ,	add_special_tokens=UpperCamelCase_      )
          UpperCAmelCase__	:      Union[str, Any]     							=					rust_tokenizer.encode(UpperCamelCase_     ,	add_special_tokens=UpperCamelCase_      )
          self.assertListEqual(UpperCamelCase_     ,	UpperCamelCase_      )

          UpperCAmelCase__	:      str     							=					self.get_rust_tokenizer()
          UpperCAmelCase__	:      Union[str, Any]     							=					tokenizer.encode(UpperCamelCase_      )
          UpperCAmelCase__	:      List[str]     							=					rust_tokenizer.encode(UpperCamelCase_      )
          self.assertListEqual(UpperCamelCase_     ,	UpperCamelCase_      )



      @slow
      def  __snake_case							(   self      ):
          # fmt: off
          UpperCAmelCase__	:      Any     							=					{'input_ids': [[5, 54, 7_196, 297, 30, 23, 776, 18, 11, 3_215, 3_705, 8_252, 22, 3_164, 1_181, 2_116, 29, 16, 813, 25, 791, 3_314, 20, 3_446, 38, 27_575, 120, 6, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [5, 468, 17, 11, 9_088, 20, 1_517, 8, 22_804, 18_818, 10, 38, 629, 607, 607, 142, 19, 7_196, 867, 56, 10_326, 24, 2_267, 20, 416, 5_072, 15_612, 233, 734, 7, 2_399, 27, 16, 3_015, 1_649, 7, 24, 20, 4_338, 2_399, 27, 13, 3_400, 14, 13, 6_189, 8, 930, 9, 6]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}  # noqa: E501
          # fmt: on

          # camembert is a french model. So we also use french texts.
          UpperCAmelCase__	:      Union[str, Any]     							=					[
              'Le transformeur est un modèle d\'apprentissage profond introduit en 2017, '
              'utilisé principalement dans le domaine du traitement automatique des langues (TAL).',
              'À l\'instar des réseaux de neurones récurrents (RNN), les transformeurs sont conçus '
              'pour gérer des données séquentielles, telles que le langage naturel, pour des tâches '
              'telles que la traduction et la synthèse de texte.',
          ]

          self.tokenizer_integration_test_util(
              expected_encoding=UpperCamelCase_     ,	model_name='camembert-base'     ,	revision='3a0641d9a1aeb7e848a74299e7e4c4bca216b4cf'     ,	sequences=UpperCamelCase_     ,	)

