




"""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        _UpperCamelCase     (      )      ->      Optional[Any]:


			"""simple docstring"""
			__UpperCAmelCase						:			Union[str, Any]												=       'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
			__UpperCAmelCase						:			Any												=       Image.open(requests.get(UpperCAmelCase_					,							stream=UpperCAmelCase_				).raw				).convert("RGB"				)

			return image

def        _UpperCamelCase     (      UpperCamelCase				)      ->      Tuple:


			"""simple docstring"""
			__UpperCAmelCase						:			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        _UpperCamelCase     (      UpperCamelCase					,							UpperCamelCase					,							UpperCamelCase				)      ->      int:


			"""simple docstring"""
			__UpperCAmelCase						:			List[Any]												=       dct.pop(UpperCAmelCase_				)
			__UpperCAmelCase						:			Optional[int]												=       val

def        _UpperCamelCase     (      UpperCamelCase					,							UpperCamelCase				)      ->      List[Any]:


			"""simple docstring"""
			for i in range(config.vision_config.num_hidden_layers				):
						# read in original q and v biases
						__UpperCAmelCase						:			Tuple												=       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						:			List[Any]												=       torch.cat((q_bias, torch.zeros_like(UpperCAmelCase_					,							requires_grad=UpperCAmelCase_				), v_bias)				)
						__UpperCAmelCase						:			Any												=       qkv_bias

def        _UpperCamelCase     (      UpperCamelCase					,							UpperCamelCase				)      ->      Optional[int]:


			"""simple docstring"""
			__UpperCAmelCase						:			int												=       364 if 'coco' in model_name else 224
			__UpperCAmelCase						:			List[Any]												=       BlipaVisionConfig(image_size=UpperCAmelCase_				).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:
						__UpperCAmelCase						:			Union[str, Any]												=       OPTConfig.from_pretrained("facebook/opt-2.7b"					,							eos_token_id=UpperCAmelCase_				).to_dict()
			elif "opt-6.7b" in model_name:
						__UpperCAmelCase						:			List[str]												=       OPTConfig.from_pretrained("facebook/opt-6.7b"					,							eos_token_id=UpperCAmelCase_				).to_dict()
			elif "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						:			List[str]												=       TaConfig.from_pretrained("google/flan-t5-xxl"					,							dense_act_fn="gelu"					,							bos_token_id=1				).to_dict()

			__UpperCAmelCase						:			Tuple												=       BlipaConfig(vision_config=UpperCAmelCase_					,							text_config=UpperCAmelCase_				)

			return config, image_size

@torch.no_grad()
def        _UpperCamelCase     (      UpperCamelCase					,							UpperCamelCase=None					,							UpperCamelCase=False				)      ->      Dict:


			"""simple docstring"""
			__UpperCAmelCase						:			Tuple												=       (
			    AutoTokenizer.from_pretrained("facebook/opt-2.7b"				)
			    if 'opt' in model_name
			    else AutoTokenizer.from_pretrained("google/flan-t5-xl"				)
			)
			__UpperCAmelCase						:			Union[str, Any]												=       tokenizer("\n"					,							add_special_tokens=UpperCAmelCase_				).input_ids[0]
			__UpperCAmelCase						:			str												=       get_blipa_config(UpperCAmelCase_					,							eos_token_id=UpperCAmelCase_				)

			__UpperCAmelCase						:			Tuple												=       BlipaForConditionalGeneration(UpperCAmelCase_				).eval()

			__UpperCAmelCase						:			int												=       {
			    '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'),
			}

			__UpperCAmelCase						:			int												=       model_name_to_original[model_name]

			# load original model
			print("Loading original model..."				)
			__UpperCAmelCase						:			Any												=       'cuda' if torch.cuda.is_available() else 'cpu'
			__UpperCAmelCase						:			str												=       load_model_and_preprocess(
			    name=UpperCAmelCase_					,							model_type=UpperCAmelCase_					,							is_eval=UpperCAmelCase_					,							device=UpperCAmelCase_				)
			original_model.eval()
			print("Done!"				)

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

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

			# read in qv biases
			read_in_q_v_bias(UpperCAmelCase_					,							UpperCAmelCase_				)

			__UpperCAmelCase						:			List[str]												=       hf_model.load_state_dict(UpperCAmelCase_					,							strict=UpperCAmelCase_				)
			assert len(UpperCAmelCase_				) == 0
			assert unexpected_keys == ["qformer.embeddings.position_ids"]

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

			# create processor
			__UpperCAmelCase						:			Tuple												=       BlipImageProcessor(
			    size={"height": image_size, "width": image_size}					,							image_mean=UpperCAmelCase_					,							image_std=UpperCAmelCase_				)
			__UpperCAmelCase						:			List[str]												=       BlipaProcessor(image_processor=UpperCAmelCase_					,							tokenizer=UpperCAmelCase_				)
			__UpperCAmelCase						:			str												=       processor(images=UpperCAmelCase_					,							return_tensors="pt"				).pixel_values.to(UpperCAmelCase_				)

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

			original_model.to(UpperCAmelCase_				)
			hf_model.to(UpperCAmelCase_				)
			with torch.no_grad():
						if "opt" in model_name:
									__UpperCAmelCase						:			Optional[Any]												=       original_model({"image": original_pixel_values, "text_input": [""]}				).logits
									__UpperCAmelCase						:			List[str]												=       hf_model(UpperCAmelCase_					,							UpperCAmelCase_				).logits
						else:
									__UpperCAmelCase						:			Tuple												=       original_model(
									    {"image": original_pixel_values, "text_input": ["\n"], "text_output": ["\n"]}				).logits
									__UpperCAmelCase						:			List[str]												=       input_ids.masked_fill(input_ids == tokenizer.pad_token_id					,							-100				)
									__UpperCAmelCase						:			Dict												=       hf_model(UpperCAmelCase_					,							UpperCAmelCase_					,							labels=UpperCAmelCase_				).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":
						__UpperCAmelCase						:			List[Any]												=       torch.tensor(
						    [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]					,							device=UpperCAmelCase_				)
						assert torch.allclose(logits[0, :3, :3]					,							UpperCAmelCase_					,							atol=1e-4				)
			elif model_name == "blip2-flan-t5-xl-coco":
						__UpperCAmelCase						:			Tuple												=       torch.tensor(
						    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]					,							device=UpperCAmelCase_				)
			else:
						# cast to same type
						__UpperCAmelCase						:			Optional[int]												=       logits.dtype
						assert torch.allclose(original_logits.to(UpperCAmelCase_				)					,							UpperCAmelCase_					,							atol=1e-2				)
			print("Looks ok!"				)

			print("Generating a caption..."				)
			__UpperCAmelCase						:			str												=       ''
			__UpperCAmelCase						:			str												=       tokenizer(UpperCAmelCase_					,							return_tensors="pt"				).input_ids.to(UpperCAmelCase_				)

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

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

			if push_to_hub:
						processor.push_to_hub(f"nielsr/{model_name}"				)
						hf_model.push_to_hub(f"nielsr/{model_name}"				)


if __name__ == "__main__":
					A				=							argparse.ArgumentParser()
					A				=							[
					    """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""",
					)

					A				=							parser.parse_args()

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






'''simple docstring'''



import json
from typing import List, Optional, Tuple

from tokenizers import normalizers
from tokenizers.pre_tokenizers import BertPreTokenizer, PreTokenizer

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_roformer import RoFormerTokenizer
from .tokenization_utils import JiebaPreTokenizer


A__						:   str				=    logging.get_logger(__name__)

A__						:   str				=    {"""vocab_file""": """vocab.txt""", """tokenizer_file""": """tokenizer.json"""}

A__						:   Tuple				=    {
    """vocab_file""": {
        """junnyu/roformer_chinese_small""": """https://huggingface.co/junnyu/roformer_chinese_small/resolve/main/vocab.txt""",
        """junnyu/roformer_chinese_base""": """https://huggingface.co/junnyu/roformer_chinese_base/resolve/main/vocab.txt""",
        """junnyu/roformer_chinese_char_small""": (
            """https://huggingface.co/junnyu/roformer_chinese_char_small/resolve/main/vocab.txt"""
        ),
        """junnyu/roformer_chinese_char_base""": (
            """https://huggingface.co/junnyu/roformer_chinese_char_base/resolve/main/vocab.txt"""
        ),
        """junnyu/roformer_small_discriminator""": (
            """https://huggingface.co/junnyu/roformer_small_discriminator/resolve/main/vocab.txt"""
        ),
        """junnyu/roformer_small_generator""": (
            """https://huggingface.co/junnyu/roformer_small_generator/resolve/main/vocab.txt"""
        ),
    }
}

A__						:   str				=    {
    """junnyu/roformer_chinese_small""": 1536,
    """junnyu/roformer_chinese_base""": 1536,
    """junnyu/roformer_chinese_char_small""": 512,
    """junnyu/roformer_chinese_char_base""": 512,
    """junnyu/roformer_small_discriminator""": 128,
    """junnyu/roformer_small_generator""": 128,
}


A__						:   Tuple				=    {
    """junnyu/roformer_chinese_small""": {"""do_lower_case""": True},
    """junnyu/roformer_chinese_base""": {"""do_lower_case""": True},
    """junnyu/roformer_chinese_char_small""": {"""do_lower_case""": True},
    """junnyu/roformer_chinese_char_base""": {"""do_lower_case""": True},
    """junnyu/roformer_small_discriminator""": {"""do_lower_case""": True},
    """junnyu/roformer_small_generator""": {"""do_lower_case""": True},
}





class   UpperCAmelCase_    (_UpperCAmelCase    ):





    """simple docstring"""






    lowerCamelCase    :   Union[str, Any]					=     VOCAB_FILES_NAMES
    lowerCamelCase    :   int					=     PRETRAINED_VOCAB_FILES_MAP
    lowerCamelCase    :   Any					=     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    lowerCamelCase    :   Optional[int]					=     PRETRAINED_INIT_CONFIGURATION
    lowerCamelCase    :   Dict					=     RoFormerTokenizer







    def __init__(				self     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=True     ,	SCREAMING_SNAKE_CASE_="[UNK]"     ,	SCREAMING_SNAKE_CASE_="[SEP]"     ,	SCREAMING_SNAKE_CASE_="[PAD]"     ,	SCREAMING_SNAKE_CASE_="[CLS]"     ,	SCREAMING_SNAKE_CASE_="[MASK]"     ,	SCREAMING_SNAKE_CASE_=True     ,	SCREAMING_SNAKE_CASE_=None     ,	**SCREAMING_SNAKE_CASE_     ,	)							->	Optional[Any]:
         super().__init__(
             SCREAMING_SNAKE_CASE_     ,	tokenizer_file=SCREAMING_SNAKE_CASE_     ,	do_lower_case=SCREAMING_SNAKE_CASE_     ,	unk_token=SCREAMING_SNAKE_CASE_     ,	sep_token=SCREAMING_SNAKE_CASE_     ,	pad_token=SCREAMING_SNAKE_CASE_     ,	cls_token=SCREAMING_SNAKE_CASE_     ,	mask_token=SCREAMING_SNAKE_CASE_     ,	tokenize_chinese_chars=SCREAMING_SNAKE_CASE_     ,	strip_accents=SCREAMING_SNAKE_CASE_     ,	**SCREAMING_SNAKE_CASE_     ,	)

         __lowerCamelCase  :      Optional[int]														=						json.loads(self.backend_tokenizer.normalizer.__getstate__()      )
         if (
             pre_tok_state.get('lowercase'     ,	SCREAMING_SNAKE_CASE_      ) != do_lower_case
             or pre_tok_state.get('strip_accents'     ,	SCREAMING_SNAKE_CASE_      ) != strip_accents
         ):
              __lowerCamelCase  :      Optional[int]														=						getattr(SCREAMING_SNAKE_CASE_     ,	pre_tok_state.pop('type'      )      )
              __lowerCamelCase  :      Union[str, Any]														=						do_lower_case
              __lowerCamelCase  :      str														=						strip_accents
              __lowerCamelCase  :      Optional[Any]														=						pre_tok_class(**SCREAMING_SNAKE_CASE_      )

         __lowerCamelCase  :      Tuple														=						do_lower_case







    def __getstate__(				self      )							->	List[str]:
         __lowerCamelCase  :      Union[str, Any]														=						self.__dict__.copy()
         __lowerCamelCase  :      Dict														=						BertPreTokenizer()
         return state







    def __setstate__(				self     ,	SCREAMING_SNAKE_CASE_      )							->	Optional[Any]:
         __lowerCamelCase  :      Optional[int]														=						d
         __lowerCamelCase  :      List[Any]														=						self.__dict__['_tokenizer'].get_vocab()
         __lowerCamelCase  :      Union[str, Any]														=						PreTokenizer.custom(JiebaPreTokenizer(SCREAMING_SNAKE_CASE_      )      )







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_=None      )							->	str:
         __lowerCamelCase  :      Union[str, Any]														=						[self.cls_token_id] + token_ids_a + [self.sep_token_id]

         if token_ids_a:
              output += token_ids_a + [self.sep_token_id]

         return output







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_ = None      )							->	List[int]:
         __lowerCamelCase  :      List[str]														=						[self.sep_token_id]
         __lowerCamelCase  :      Dict														=						[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  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_ = None      )							->	Tuple[str]:
         __lowerCamelCase  :      Optional[Any]														=						self._tokenizer.model.save(SCREAMING_SNAKE_CASE_     ,	name=SCREAMING_SNAKE_CASE_      )
         return tuple(SCREAMING_SNAKE_CASE_      )







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=False     ,	**SCREAMING_SNAKE_CASE_     ,	)							->	Any:
         __lowerCamelCase  :      Tuple														=						BertPreTokenizer()
         return super().save_pretrained(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	**SCREAMING_SNAKE_CASE_      )






