



'''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       lowerCamelCase__       (						):
  a						:					str        =     'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
  a						:					Optional[Any]        =     Image.open(requests.get(__lowercase ,			stream=__lowercase				).raw				).convert('RGB'				)

  return image
def       lowerCamelCase__       (						_A				):
  a						:					Union[str, 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       lowerCamelCase__       (						_A ,			_A ,			_A				):
  a						:					Union[str, Any]        =     dct.pop(__lowercase				)
  a						:					str        =     val
def       lowerCamelCase__       (						_A ,			_A				):






  for i in range(config.vision_config.num_hidden_layers				):
    # read in original q and v biases
    a						:					Any        =     state_dict.pop(f"""visual_encoder.blocks.{i}.attn.q_bias"""				)
    a						:					int        =     state_dict.pop(f"""visual_encoder.blocks.{i}.attn.v_bias"""				)

    # next, set bias in the state dict
    a						:					Tuple        =     torch.cat((q_bias, torch.zeros_like(__lowercase ,			requires_grad=__lowercase				), v_bias)				)
    a						:					int        =     qkv_bias
def       lowerCamelCase__       (						_A ,			_A				):
  a						:					List[str]        =     364 if 'coco' in model_name else 224
  a						:					Dict        =     BlipaVisionConfig(image_size=__lowercase				).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:
    a						:					Tuple        =     OPTConfig.from_pretrained('facebook/opt-2.7b' ,			eos_token_id=__lowercase				).to_dict()
  elif "opt-6.7b" in model_name:
    a						:					Dict        =     OPTConfig.from_pretrained('facebook/opt-6.7b' ,			eos_token_id=__lowercase				).to_dict()
  elif "t5-xl" in model_name:
    a						:					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:
    a						:					Optional[Any]        =     TaConfig.from_pretrained('google/flan-t5-xxl' ,			dense_act_fn='gelu' ,			bos_token_id=1				).to_dict()

  a						:					List[str]        =     BlipaConfig(vision_config=__lowercase ,			text_config=__lowercase				)

  return config, image_size
@torch.no_grad()
def       lowerCamelCase__       (						_A ,			_A=None ,			_A=False				):
  a						:					List[Any]        =     (
      AutoTokenizer.from_pretrained('facebook/opt-2.7b'				)
      if 'opt' in model_name
      else AutoTokenizer.from_pretrained('google/flan-t5-xl'				)
  )
  a						:					Tuple        =     tokenizer('\n' ,			add_special_tokens=__lowercase				).input_ids[0]
  a   ,	a						:					Optional[int]        =     get_blipa_config(__lowercase ,			eos_token_id=__lowercase				)

  a						:					Dict        =     BlipaForConditionalGeneration(__lowercase				).eval()

  a						:					Union[str, Any]        =     {
      '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'),
  }

  a   ,	a						:					Optional[int]        =     model_name_to_original[model_name]

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

  # update state dict keys
  a						:					Optional[Any]        =     original_model.state_dict()
  a						:					int        =     create_rename_keys(__lowercase				)
  for src, dest in rename_keys:
    rename_key(__lowercase ,			__lowercase ,			__lowercase				)

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

  # read in qv biases
  read_in_q_v_bias(__lowercase ,			__lowercase				)

  a   ,	a						:					Union[str, Any]        =     hf_model.load_state_dict(__lowercase ,			strict=__lowercase				)
  assert len(__lowercase				) == 0
  assert unexpected_keys == ["qformer.embeddings.position_ids"]

  a						:					Union[str, Any]        =     load_demo_image()
  a						:					Optional[int]        =     vis_processors['eval'](__lowercase				).unsqueeze(0				).to(__lowercase				)
  a						:					Tuple        =     tokenizer(['\n'] ,			return_tensors='pt'				).input_ids.to(__lowercase				)

  # create processor
  a						:					Optional[int]        =     BlipImageProcessor(
      size={'height': image_size, 'width': image_size} ,			image_mean=__lowercase ,			image_std=__lowercase				)
  a						:					Optional[int]        =     BlipaProcessor(image_processor=__lowercase ,			tokenizer=__lowercase				)
  a						:					str        =     processor(images=__lowercase ,			return_tensors='pt'				).pixel_values.to(__lowercase				)

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

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

  print('Generating a caption...'				)
  a						:					Optional[int]        =     ''
  a						:					int        =     tokenizer(__lowercase ,			return_tensors='pt'				).input_ids.to(__lowercase				)

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

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

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


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

     lowerCAmelCase:					Tuple						       =						parser.parse_args()

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






'''simple docstring'''
# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..models.auto import AutoProcessor
from ..models.vision_encoder_decoder import VisionEncoderDecoderModel
from ..utils import is_vision_available
from .base import PipelineTool


if is_vision_available():
    from PIL import Image

class      snake_case     (    __lowerCamelCase       ):




 """simple docstring"""



 SCREAMING_SNAKE_CASE_	:					Tuple					="naver-clova-ix/donut-base-finetuned-docvqa"
 SCREAMING_SNAKE_CASE_	:					Dict					=(
     "This is a tool that answers a question about an document (pdf). It takes an input named `document` which "
     "should be the document containing the information, as well as a `question` that is the question about the "
     "document. It returns a text that contains the answer to the question."
 )
 SCREAMING_SNAKE_CASE_	:					List[str]					="document_qa"
 SCREAMING_SNAKE_CASE_	:					Union[str, Any]					=AutoProcessor
 SCREAMING_SNAKE_CASE_	:					Union[str, Any]					=VisionEncoderDecoderModel

 SCREAMING_SNAKE_CASE_	:					List[Any]					=["image", "text"]
 SCREAMING_SNAKE_CASE_	:					Any					=["text"]






 def __init__(       self       :							Optional[int]							,	*__A       :							List[str]							,	**__A       :							List[Any]	):
        if not is_vision_available():
               raise ValueError('Pillow must be installed to use the DocumentQuestionAnsweringTool.'	)

        super().__init__(*__A							,	**__A	)






 def        _lowerCamelCase					(       self       :							Any							,	__A       :							"Image"							,	__A       :							str	):
        __UpperCamelCase       =  '<s_docvqa><s_question>{user_input}</s_question><s_answer>'
        __UpperCamelCase       =  task_prompt.replace('{user_input}'							,	__A	)
        __UpperCamelCase       =  self.pre_processor.tokenizer(
            __A							,	add_special_tokens=__A							,	return_tensors='pt'	).input_ids
        __UpperCamelCase       =  self.pre_processor(__A							,	return_tensors='pt'	).pixel_values

        return {"decoder_input_ids": decoder_input_ids, "pixel_values": pixel_values}






 def        _lowerCamelCase					(       self       :							Union[str, Any]							,	__A       :							Optional[Any]	):
        return self.model.generate(
            inputs['pixel_values'].to(self.device	)							,	decoder_input_ids=inputs['decoder_input_ids'].to(self.device	)							,	max_length=self.model.decoder.config.max_position_embeddings							,	early_stopping=__A							,	pad_token_id=self.pre_processor.tokenizer.pad_token_id							,	eos_token_id=self.pre_processor.tokenizer.eos_token_id							,	use_cache=__A							,	num_beams=1							,	bad_words_ids=[[self.pre_processor.tokenizer.unk_token_id]]							,	return_dict_in_generate=__A							,	).sequences






 def        _lowerCamelCase					(       self       :							Tuple							,	__A       :							List[Any]	):
        __UpperCamelCase       =  self.pre_processor.batch_decode(__A	)[0]
        __UpperCamelCase       =  sequence.replace(self.pre_processor.tokenizer.eos_token							,	''	)
        __UpperCamelCase       =  sequence.replace(self.pre_processor.tokenizer.pad_token							,	''	)
        __UpperCamelCase       =  re.sub(R'<.*?>'							,	''							,	__A							,	count=1	).strip()  # remove first task start token
        __UpperCamelCase       =  self.pre_processor.tokenajson(__A	)

        return sequence["answer"]
