






"""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 _a			(		):







      """simple docstring"""






      UpperCAmelCase          =	"""https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"""
      UpperCAmelCase          =	Image.open(requests.get(_snake_case				,	stream=_snake_case							).raw							).convert("""RGB"""							)

      return image

def _a			(		_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	[]
      # 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 _a			(		_snake_case				,	_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	dct.pop(_snake_case							)
      UpperCAmelCase          =	val

def _a			(		_snake_case				,	_snake_case							):







      """simple docstring"""






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

            # next, set bias in the state dict
            UpperCAmelCase          =	torch.cat((q_bias, torch.zeros_like(_snake_case				,	requires_grad=_snake_case							), v_bias)							)
            UpperCAmelCase          =	qkv_bias

def _a			(		_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	364 if """coco""" in model_name else 224
      UpperCAmelCase          =	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          =	TaConfig.from_pretrained("""google/flan-t5-xl"""				,	dense_act_fn="""gelu"""				,	bos_token_id=1							).to_dict()
      elif "t5-xxl" in model_name:
            UpperCAmelCase          =	TaConfig.from_pretrained("""google/flan-t5-xxl"""				,	dense_act_fn="""gelu"""				,	bos_token_id=1							).to_dict()
      elif "vicuna-7b" in model_name:
            UpperCAmelCase          =	LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf"""				,	vocab_size=3_2001							).to_dict()
      elif "vicuna-13b" in model_name:
            UpperCAmelCase          =	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          =	InstructBlipQFormerConfig(vocab_size=3_0523							).to_dict()
      UpperCAmelCase          =	InstructBlipConfig(vision_config=_snake_case				,	text_config=_snake_case				,	qformer_config=_snake_case							)

      return config, image_size

@torch.no_grad()
def _a			(		_snake_case				,	_snake_case=None				,	_snake_case=False							):







      """simple docstring"""






      UpperCAmelCase          =	AutoTokenizer.from_pretrained("""bert-base-uncased"""				,	truncation_side="""left"""							)
      qformer_tokenizer.add_special_tokens({"""bos_token""": """[DEC]"""}							)

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

      UpperCAmelCase      , UpperCAmelCase          =	get_blipa_config(_snake_case							)
      UpperCAmelCase          =	InstructBlipForConditionalGeneration(_snake_case							).eval()

      UpperCAmelCase          =	{
          """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          =	model_name_to_original[model_name]

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

      # create processor
      UpperCAmelCase          =	BlipImageProcessor(
          size={"""height""": image_size, """width""": image_size}				,	image_mean=_snake_case				,	image_std=_snake_case							)
      UpperCAmelCase          =	InstructBlipProcessor(
          image_processor=_snake_case				,	tokenizer=_snake_case				,	qformer_tokenizer=_snake_case				,	)
      UpperCAmelCase          =	processor(images=_snake_case				,	text=_snake_case				,	return_tensors="""pt"""							).to(_snake_case							)

      # make sure processor creates exact same pixel values
      UpperCAmelCase          =	vis_processors["""eval"""](_snake_case							).unsqueeze(0							).to(_snake_case							)
      UpperCAmelCase          =	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          =	original_model({"""image""": original_pixel_values, """text_input""": [prompt]}							).logits
                  UpperCAmelCase          =	hf_model(**_snake_case							).logits
            else:
                  UpperCAmelCase          =	original_model(
                      {"""image""": original_pixel_values, """text_input""": [prompt], """text_output""": ["""\n"""]}							).logits
                  UpperCAmelCase          =	tokenizer("""\n"""				,	return_tensors="""pt"""							).input_ids.to(_snake_case							)
                  UpperCAmelCase          =	label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id				,	-100							)
                  UpperCAmelCase          =	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          =	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          =	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          =	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          =	2
      print("""Original generation:"""				,	_snake_case							)
      UpperCAmelCase          =	processor.batch_decode(_snake_case				,	skip_special_tokens=_snake_case							)
      UpperCAmelCase          =	[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 os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


_UpperCamelCase					      =  logging.get_logger(__name__)

_UpperCamelCase					      =  """▁"""

_UpperCamelCase					      =  {"""vocab_file""": """sentencepiece.bpe.model""", """monolingual_vocab_file""": """dict.txt"""}

_UpperCamelCase					      =  {
    """vocab_file""": {
        """vinai/bartpho-syllable""": """https://huggingface.co/vinai/bartpho-syllable/resolve/main/sentencepiece.bpe.model""",
    },
    """monolingual_vocab_file""": {
        """vinai/bartpho-syllable""": """https://huggingface.co/vinai/bartpho-syllable/resolve/main/dict.txt""",
    },
}

_UpperCamelCase					      =  {"""vinai/bartpho-syllable""": 1024}







class  lowerCamelCase__ (	snake_case			):
    SCREAMING_SNAKE_CASE		=      VOCAB_FILES_NAMES
    SCREAMING_SNAKE_CASE		=      PRETRAINED_VOCAB_FILES_MAP
    SCREAMING_SNAKE_CASE		=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    SCREAMING_SNAKE_CASE		=      ['''input_ids''', '''attention_mask''']






    def __init__(			self     ,A     ,A     ,A="<s>"     ,A="</s>"     ,A="</s>"     ,A="<s>"     ,A="<unk>"     ,A="<pad>"     ,A="<mask>"     ,A = None     ,**A     ,):
          # Mask token behave like a normal word, i.e. include the space before it
          UpperCAmelCase          =	AddedToken(A     ,lstrip=A     ,rstrip=A ) if isinstance(A     ,A ) else mask_token

          UpperCAmelCase          =	{} if sp_model_kwargs is None else sp_model_kwargs

          super().__init__(
              bos_token=A     ,eos_token=A     ,unk_token=A     ,sep_token=A     ,cls_token=A     ,pad_token=A     ,mask_token=A     ,sp_model_kwargs=self.sp_model_kwargs     ,**A     ,)

          UpperCAmelCase          =	vocab_file
          UpperCAmelCase          =	monolingual_vocab_file
          UpperCAmelCase          =	spm.SentencePieceProcessor(**self.sp_model_kwargs )
          self.sp_model.Load(str(A ) )

          # Load the reduced vocab

          # Keep order of special tokens for backward compatibility
          UpperCAmelCase          =	{}
          UpperCAmelCase          =	0
          for token in [bos_token, pad_token, eos_token, unk_token, sep_token, cls_token]:
                if str(A ) not in self.fairseq_tokens_to_ids:
                      UpperCAmelCase          =	cnt
                      cnt += 1
          with open(A     ,"""r"""     ,encoding="""utf-8""" ) as f:
                for line in f.readlines():
                      UpperCAmelCase          =	line.strip().split()[0]
                      UpperCAmelCase          =	len(self.fairseq_tokens_to_ids )
          if str(A ) not in self.fairseq_tokens_to_ids:
                UpperCAmelCase          =	len(self.fairseq_tokens_to_ids )

          UpperCAmelCase          =	{v: k for k, v in self.fairseq_tokens_to_ids.items()}






    def __getstate__(			self ):
          UpperCAmelCase          =	self.__dict__.copy()
          UpperCAmelCase          =	None
          UpperCAmelCase          =	self.sp_model.serialized_model_proto()
          return state






    def __setstate__(			self     ,A ):
          UpperCAmelCase          =	d

          # for backward compatibility
          if not hasattr(self     ,"""sp_model_kwargs""" ):
                UpperCAmelCase          =	{}

          UpperCAmelCase          =	spm.SentencePieceProcessor(**self.sp_model_kwargs )
          self.sp_model.LoadFromSerializedProto(self.sp_model_proto )






    def     _UpperCamelCase							(			self     ,A     ,A = None ):

          if token_ids_a is None:
                return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
          UpperCAmelCase          =	[self.cls_token_id]
          UpperCAmelCase          =	[self.sep_token_id]
          return cls + token_ids_a + sep + sep + token_ids_a + sep






    def     _UpperCamelCase							(			self     ,A     ,A = None     ,A = False ):

          if already_has_special_tokens:
                return super().get_special_tokens_mask(
                    token_ids_a=A     ,token_ids_a=A     ,already_has_special_tokens=A )

          if token_ids_a is None:
                return [1] + ([0] * len(A )) + [1]
          return [1] + ([0] * len(A )) + [1, 1] + ([0] * len(A )) + [1]






    def     _UpperCamelCase							(			self     ,A     ,A = None ):
          UpperCAmelCase          =	[self.sep_token_id]
          UpperCAmelCase          =	[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 + sep + token_ids_a + sep ) * [0]






    @property
    def     _UpperCamelCase							(			self ):
          return len(self.fairseq_ids_to_tokens )






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	{self.convert_ids_to_tokens(A ): i for i in range(self.vocab_size )}
          vocab.update(self.added_tokens_encoder )
          return vocab






    def     _UpperCamelCase							(			self     ,A ):
          return self.sp_model.encode(A     ,out_type=A )






    def     _UpperCamelCase							(			self     ,A ):
          if token in self.fairseq_tokens_to_ids:
                return self.fairseq_tokens_to_ids[token]
          else:
                return self.unk_token_id






    def     _UpperCamelCase							(			self     ,A ):
          return self.fairseq_ids_to_tokens[index]






    def     _UpperCamelCase							(			self     ,A ):
          UpperCAmelCase          =	"""""".join(A ).replace(A     ,""" """ ).strip()
          return out_string






    def     _UpperCamelCase							(			self     ,A     ,A = None ):
          if not os.path.isdir(A ):
                logger.error(F'''Vocabulary path ({save_directory}) should be a directory''' )
                return
          UpperCAmelCase          =	os.path.join(
              A     ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""] )
          UpperCAmelCase          =	os.path.join(
              A     ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""monolingual_vocab_file"""]     ,)

          if os.path.abspath(self.vocab_file ) != os.path.abspath(A ) and os.path.isfile(self.vocab_file ):
                copyfile(self.vocab_file     ,A )
          elif not os.path.isfile(self.vocab_file ):
                with open(A     ,"""wb""" ) as fi:
                      UpperCAmelCase          =	self.sp_model.serialized_model_proto()
                      fi.write(A )

          if os.path.abspath(self.monolingual_vocab_file ) != os.path.abspath(
              A ) and os.path.isfile(self.monolingual_vocab_file ):
                copyfile(self.monolingual_vocab_file     ,A )
          elif not os.path.isfile(self.monolingual_vocab_file ):
                with open(A     ,"""w"""     ,encoding="""utf-8""" ) as fp:
                      for token in self.fairseq_tokens_to_ids:
                            if token not in self.all_special_tokens:
                                  fp.write(F'''{str(A )} \n''' )

          return out_vocab_file, out_monolingual_vocab_file



