


import collections
from typing import List, Optional, Union

from ...tokenization_utils_base import BatchEncoding
from ...utils import TensorType, add_end_docstrings, add_start_docstrings, logging
from ..bert.tokenization_bert import BertTokenizer


__UpperCamelCase				:  List[Any]      					=							logging.get_logger(__name__)

__UpperCamelCase				:  Any      					=							{'vocab_file': 'vocab.txt', 'tokenizer_file': 'tokenizer.json'}

__UpperCamelCase				:  List[Any]      					=							{
    'vocab_file': {
        'facebook/dpr-ctx_encoder-single-nq-base': (
            'https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/vocab.txt'
        ),
        'facebook/dpr-ctx_encoder-multiset-base': (
            'https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/vocab.txt'
        ),
    },
    'tokenizer_file': {
        'facebook/dpr-ctx_encoder-single-nq-base': (
            'https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/tokenizer.json'
        ),
        'facebook/dpr-ctx_encoder-multiset-base': (
            'https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/tokenizer.json'
        ),
    },
}
__UpperCamelCase				:  List[Any]      					=							{
    'vocab_file': {
        'facebook/dpr-question_encoder-single-nq-base': (
            'https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/vocab.txt'
        ),
        'facebook/dpr-question_encoder-multiset-base': (
            'https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/vocab.txt'
        ),
    },
    'tokenizer_file': {
        'facebook/dpr-question_encoder-single-nq-base': (
            'https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/tokenizer.json'
        ),
        'facebook/dpr-question_encoder-multiset-base': (
            'https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/tokenizer.json'
        ),
    },
}
__UpperCamelCase				:  List[Any]      					=							{
    'vocab_file': {
        'facebook/dpr-reader-single-nq-base': (
            'https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/vocab.txt'
        ),
        'facebook/dpr-reader-multiset-base': (
            'https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/vocab.txt'
        ),
    },
    'tokenizer_file': {
        'facebook/dpr-reader-single-nq-base': (
            'https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/tokenizer.json'
        ),
        'facebook/dpr-reader-multiset-base': (
            'https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/tokenizer.json'
        ),
    },
}

__UpperCamelCase				:  Dict      					=							{
    'facebook/dpr-ctx_encoder-single-nq-base': 512,
    'facebook/dpr-ctx_encoder-multiset-base': 512,
}
__UpperCamelCase				:  List[str]      					=							{
    'facebook/dpr-question_encoder-single-nq-base': 512,
    'facebook/dpr-question_encoder-multiset-base': 512,
}
__UpperCamelCase				:  Union[str, Any]      					=							{
    'facebook/dpr-reader-single-nq-base': 512,
    'facebook/dpr-reader-multiset-base': 512,
}


__UpperCamelCase				:  Optional[int]      					=							{
    'facebook/dpr-ctx_encoder-single-nq-base': {'do_lower_case': True},
    'facebook/dpr-ctx_encoder-multiset-base': {'do_lower_case': True},
}
__UpperCamelCase				:  Dict      					=							{
    'facebook/dpr-question_encoder-single-nq-base': {'do_lower_case': True},
    'facebook/dpr-question_encoder-multiset-base': {'do_lower_case': True},
}
__UpperCamelCase				:  str      					=							{
    'facebook/dpr-reader-single-nq-base': {'do_lower_case': True},
    'facebook/dpr-reader-multiset-base': {'do_lower_case': True},
}





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     Dict			= VOCAB_FILES_NAMES
    A_  :     str			= CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
    A_  :     List[Any]			= CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    A_  :     Any			= CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     Union[str, Any]			= VOCAB_FILES_NAMES
    A_  :     Optional[Any]			= QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
    A_  :     Any			= QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    A_  :     Tuple			= QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION





__UpperCamelCase				:  Any      					=							collections.namedtuple(
    'DPRSpanPrediction', ['span_score', 'relevance_score', 'doc_id', 'start_index', 'end_index', 'text']
)

__UpperCamelCase				:  Optional[Any]      					=							collections.namedtuple('DPRReaderOutput', ['start_logits', 'end_logits', 'relevance_logits'])


__UpperCamelCase				:  List[str]      					=							R'\n    Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`.\n    It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers),\n    using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)`\n    with the format:\n\n    ```\n    [CLS] <question token ids> [SEP] <titles ids> [SEP] <texts ids>\n    ```\n\n    Args:\n        questions (`str` or `List[str]`):\n            The questions to be encoded. You can specify one question for many passages. In this case, the question\n            will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in\n            `titles` or `texts`.\n        titles (`str` or `List[str]`):\n            The passages titles to be encoded. This can be a string or a list of strings if there are several passages.\n        texts (`str` or `List[str]`):\n            The passages texts to be encoded. This can be a string or a list of strings if there are several passages.\n        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):\n            Activates and controls padding. Accepts the following values:\n\n            - `True` or `\'longest\'`: Pad to the longest sequence in the batch (or no padding if only a single sequence\n              if provided).\n            - `\'max_length\'`: Pad to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided.\n            - `False` or `\'do_not_pad\'` (default): No padding (i.e., can output a batch with sequences of different\n              lengths).\n        truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):\n            Activates and controls truncation. Accepts the following values:\n\n            - `True` or `\'longest_first\'`: Truncate to a maximum length specified with the argument `max_length` or to\n              the maximum acceptable input length for the model if that argument is not provided. This will truncate\n              token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch\n              of pairs) is provided.\n            - `\'only_first\'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided. This will only truncate the first\n              sequence of a pair if a pair of sequences (or a batch of pairs) is provided.\n            - `\'only_second\'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided. This will only truncate the\n              second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.\n            - `False` or `\'do_not_truncate\'` (default): No truncation (i.e., can output batch with sequence lengths\n              greater than the model maximum admissible input size).\n        max_length (`int`, *optional*):\n                Controls the maximum length to use by one of the truncation/padding parameters.\n\n                If left unset or set to `None`, this will use the predefined model maximum length if a maximum length\n                is required by one of the truncation/padding parameters. If the model has no specific maximum input\n                length (like XLNet) truncation/padding to a maximum length will be deactivated.\n        return_tensors (`str` or [`~utils.TensorType`], *optional*):\n                If set, will return tensors instead of list of python integers. Acceptable values are:\n\n                - `\'tf\'`: Return TensorFlow `tf.constant` objects.\n                - `\'pt\'`: Return PyTorch `torch.Tensor` objects.\n                - `\'np\'`: Return Numpy `np.ndarray` objects.\n        return_attention_mask (`bool`, *optional*):\n            Whether or not to return the attention mask. If not set, will return the attention mask according to the\n            specific tokenizer\'s default, defined by the `return_outputs` attribute.\n\n            [What are attention masks?](../glossary#attention-mask)\n\n    Returns:\n        `Dict[str, List[List[int]]]`: A dictionary with the following keys:\n\n        - `input_ids`: List of token ids to be fed to a model.\n        - `attention_mask`: List of indices specifying which tokens should be attended to by the model.\n    '





@add_start_docstrings(snake_case__						)
class  lowerCAmelCase__:






    '''simple docstring'''





    def __call__(	self				:   str       ,   __snake_case				:   Tuple       ,   __snake_case				:   Optional[str] = None       ,   __snake_case				:   Optional[str] = None       ,   __snake_case				:   Union[bool, str] = False       ,   __snake_case				:   Union[bool, str] = False       ,   __snake_case				:   Optional[int] = None       ,   __snake_case				:   Optional[Union[str, TensorType]] = None       ,   __snake_case				:   Optional[bool] = None       ,   **__snake_case				:   Optional[int]       ,   ):





       '''simple docstring'''
       if titles is None and texts is None:
          return super().__call__(
              __snake_case       ,   padding=__snake_case       ,   truncation=__snake_case       ,   max_length=__snake_case       ,   return_tensors=__snake_case       ,   return_attention_mask=__snake_case       ,   **__snake_case       ,   )
       elif titles is None or texts is None:
          UpperCAmelCase_							:      int						=   titles if texts is None else texts
          return super().__call__(
              __snake_case       ,   __snake_case       ,   padding=__snake_case       ,   truncation=__snake_case       ,   max_length=__snake_case       ,   return_tensors=__snake_case       ,   return_attention_mask=__snake_case       ,   **__snake_case       ,   )
       UpperCAmelCase_							:      Dict						=   titles if not isinstance(__snake_case       ,   __snake_case	) else [titles]
       UpperCAmelCase_							:      Optional[Any]						=   texts if not isinstance(__snake_case       ,   __snake_case	) else [texts]
       UpperCAmelCase_							:      Optional[Any]						=   len(__snake_case	)
       UpperCAmelCase_							:      Union[str, Any]						=   questions if not isinstance(__snake_case       ,   __snake_case	) else [questions] * n_passages
       if len(__snake_case	) != len(__snake_case	):
          raise ValueError(
              f'''There should be as many titles than texts but got {len(__snake_case	)} titles and {len(__snake_case	)} texts.'''	)
       UpperCAmelCase_							:      List[str]						=   super().__call__(__snake_case       ,   __snake_case       ,   padding=__snake_case       ,   truncation=__snake_case	)['''input_ids''']
       UpperCAmelCase_							:      str						=   super().__call__(__snake_case       ,   add_special_tokens=__snake_case       ,   padding=__snake_case       ,   truncation=__snake_case	)['''input_ids''']
       UpperCAmelCase_							:      Optional[Any]						=   {
           '''input_ids''': [
               (encoded_question_and_title + encoded_text)[:max_length]
               if max_length is not None and truncation
               else encoded_question_and_title + encoded_text
               for encoded_question_and_title, encoded_text in zip(__snake_case       ,   __snake_case	)
           ]
       }
       if return_attention_mask is not False:
          UpperCAmelCase_							:      Optional[Any]						=   []
          for input_ids in encoded_inputs["input_ids"]:
             attention_mask.append([int(input_id != self.pad_token_id	) for input_id in input_ids]	)
          UpperCAmelCase_							:      Dict						=   attention_mask
       return self.pad(__snake_case       ,   padding=__snake_case       ,   max_length=__snake_case       ,   return_tensors=__snake_case	)





    def 					_lowerCamelCase (	self				:   str       ,   __snake_case				:   BatchEncoding       ,   __snake_case				:   DPRReaderOutput       ,   __snake_case				:   int = 16       ,   __snake_case				:   int = 64       ,   __snake_case				:   int = 4       ,   ):





       '''simple docstring'''
       UpperCAmelCase_							:      List[str]						=   reader_input['''input_ids''']
       UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_							:      Dict						=   reader_output[:3]
       UpperCAmelCase_							:      Optional[int]						=   len(__snake_case	)
       UpperCAmelCase_							:      str						=   sorted(range(__snake_case	)       ,   reverse=__snake_case       ,   key=relevance_logits.__getitem__	)
       UpperCAmelCase_							:      List[DPRReaderOutput]						=   []
       for doc_id in sorted_docs:
          UpperCAmelCase_							:      List[str]						=   list(input_ids[doc_id]	)
          # assuming question & title information is at the beginning of the sequence
          UpperCAmelCase_							:      List[Any]						=   sequence_ids.index(self.sep_token_id       ,   2	) + 1  # second sep id
          if sequence_ids[-1] == self.pad_token_id:
             UpperCAmelCase_							:      List[str]						=   sequence_ids.index(self.pad_token_id	)
          else:
             UpperCAmelCase_							:      Dict						=   len(__snake_case	)

          UpperCAmelCase_							:      str						=   self._get_best_spans(
              start_logits=start_logits[doc_id][passage_offset:sequence_len]       ,   end_logits=end_logits[doc_id][passage_offset:sequence_len]       ,   max_answer_length=__snake_case       ,   top_spans=__snake_case       ,   )
          for start_index, end_index in best_spans:
             start_index += passage_offset
             end_index += passage_offset
             nbest_spans_predictions.append(
                 DPRSpanPrediction(
                     span_score=start_logits[doc_id][start_index] + end_logits[doc_id][end_index]       ,   relevance_score=relevance_logits[doc_id]       ,   doc_id=__snake_case       ,   start_index=__snake_case       ,   end_index=__snake_case       ,   text=self.decode(sequence_ids[start_index : end_index + 1]	)       ,   )	)
          if len(__snake_case	) >= num_spans:
             break
       return nbest_spans_predictions[:num_spans]





    def 					_lowerCamelCase (	self				:   Tuple       ,   __snake_case				:   List[int]       ,   __snake_case				:   List[int]       ,   __snake_case				:   int       ,   __snake_case				:   int       ,   ):





       '''simple docstring'''
       UpperCAmelCase_							:      Tuple						=   []
       for start_index, start_score in enumerate(__snake_case	):
          for answer_length, end_score in enumerate(end_logits[start_index : start_index + max_answer_length]	):
             scores.append(((start_index, start_index + answer_length), start_score + end_score)	)
       UpperCAmelCase_							:      Dict						=   sorted(__snake_case       ,   key=lambda __snake_case	: x[1]       ,   reverse=__snake_case	)
       UpperCAmelCase_							:      List[Any]						=   []
       for (start_index, end_index), score in scores:
          if start_index > end_index:
             raise ValueError(f'''Wrong span indices: [{start_index}:{end_index}]'''	)
          UpperCAmelCase_							:      Dict						=   end_index - start_index + 1
          if length > max_answer_length:
             raise ValueError(f'''Span is too long: {length} > {max_answer_length}'''	)
          if any(
              start_index <= prev_start_index <= prev_end_index <= end_index
              or prev_start_index <= start_index <= end_index <= prev_end_index
              for (prev_start_index, prev_end_index) in chosen_span_intervals	):
             continue
          chosen_span_intervals.append((start_index, end_index)	)

          if len(__snake_case	) == top_spans:
             break
       return chosen_span_intervals





@add_end_docstrings(snake_case__						)
class  lowerCAmelCase__(    snake_case__    ,    snake_case__						):






    '''simple docstring'''





    A_  :     str			= VOCAB_FILES_NAMES
    A_  :     Optional[Any]			= READER_PRETRAINED_VOCAB_FILES_MAP
    A_  :     Union[str, Any]			= READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    A_  :     int			= READER_PRETRAINED_INIT_CONFIGURATION
    A_  :     Union[str, Any]			= ['input_ids', 'attention_mask']



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      snake_case_						( ):
   UpperCAmelCase_							:      int						=   '''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
   UpperCAmelCase_							:      Dict						=   Image.open(requests.get(__lowercase			,       stream=__lowercase       ).raw       ).convert('''RGB'''       )

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

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

   UpperCAmelCase_							:      List[Any]						=   BlipaConfig(vision_config=__lowercase			,       text_config=__lowercase       )

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

   UpperCAmelCase_							:      List[Any]						=   BlipaForConditionalGeneration(__lowercase       ).eval()

   UpperCAmelCase_							:      Tuple						=   {
       '''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_		,      UpperCAmelCase_							:      Optional[Any]						=   model_name_to_original[model_name]

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

   # update state dict keys
   UpperCAmelCase_							:      Optional[Any]						=   original_model.state_dict()
   UpperCAmelCase_							:      List[Any]						=   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():
      UpperCAmelCase_							:      Union[str, Any]						=   state_dict.pop(__lowercase       )
      if key.startswith('''Qformer.bert'''       ):
         UpperCAmelCase_							:      Tuple						=   key.replace('''Qformer.bert'''			,       '''qformer'''       )
      if "attention.self" in key:
         UpperCAmelCase_							:      Optional[Any]						=   key.replace('''self'''			,       '''attention'''       )
      if "opt_proj" in key:
         UpperCAmelCase_							:      Any						=   key.replace('''opt_proj'''			,       '''language_projection'''       )
      if "t5_proj" in key:
         UpperCAmelCase_							:      Tuple						=   key.replace('''t5_proj'''			,       '''language_projection'''       )
      if key.startswith('''opt'''       ):
         UpperCAmelCase_							:      Any						=   key.replace('''opt'''			,       '''language'''       )
      if key.startswith('''t5'''       ):
         UpperCAmelCase_							:      Optional[Any]						=   key.replace('''t5'''			,       '''language'''       )
      UpperCAmelCase_							:      List[str]						=   val

   # read in qv biases
   read_in_q_v_bias(__lowercase			,       __lowercase       )

   UpperCAmelCase_		,      UpperCAmelCase_							:      Optional[int]						=   hf_model.load_state_dict(__lowercase			,       strict=__lowercase       )
   assert len(__lowercase       ) == 0
   assert unexpected_keys == ["qformer.embeddings.position_ids"]

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

   # create processor
   UpperCAmelCase_							:      Optional[int]						=   BlipImageProcessor(
       size={'''height''': image_size, '''width''': image_size}			,       image_mean=__lowercase			,       image_std=__lowercase       )
   UpperCAmelCase_							:      Tuple						=   BlipaProcessor(image_processor=__lowercase			,       tokenizer=__lowercase       )
   UpperCAmelCase_							:      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:
         UpperCAmelCase_							:      Tuple						=   original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}       ).logits
         UpperCAmelCase_							:      Optional[int]						=   hf_model(__lowercase			,       __lowercase       ).logits
      else:
         UpperCAmelCase_							:      int						=   original_model(
             {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}       ).logits
         UpperCAmelCase_							:      Optional[int]						=   input_ids.masked_fill(input_ids == tokenizer.pad_token_id			,       -1_0_0       )
         UpperCAmelCase_							:      int						=   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":
      UpperCAmelCase_							:      Tuple						=   torch.tensor(
          [[-4_1.5_8_5_0, -4.4_4_4_0, -8.9_9_2_2], [-4_7.4_3_2_2, -5.9_1_4_3, -1.7_3_4_0]]			,       device=__lowercase       )
      assert torch.allclose(logits[0, :3, :3]			,       __lowercase			,       atol=1e-4       )
   elif model_name == "blip2-flan-t5-xl-coco":
      UpperCAmelCase_							:      Tuple						=   torch.tensor(
          [[-5_7.0_1_0_9, -9.8_9_6_7, -1_2.6_2_8_0], [-6_8.6_5_7_8, -1_2.7_1_9_1, -1_0.5_0_6_5]]			,       device=__lowercase       )
   else:
      # cast to same type
      UpperCAmelCase_							:      Optional[int]						=   logits.dtype
      assert torch.allclose(original_logits.to(__lowercase       )			,       __lowercase			,       atol=1e-2       )
   print('''Looks ok!'''       )

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

   UpperCAmelCase_							:      int						=   original_model.generate({'''image''': original_pixel_values}       )
   UpperCAmelCase_							:      Optional[int]						=   hf_model.generate(
       __lowercase			,       __lowercase			,       do_sample=__lowercase			,       num_beams=5			,       max_length=3_0			,       min_length=1			,       top_p=0.9			,       repetition_penalty=1.0			,       length_penalty=1.0			,       temperature=1			,       )
   print('''Original generation:'''			,       __lowercase       )
   UpperCAmelCase_							:      Tuple						=   input_ids.shape[1]
   UpperCAmelCase_							:      str						=   processor.batch_decode(outputs[:, prompt_length:]			,       skip_special_tokens=__lowercase       )
   UpperCAmelCase_							:      Optional[int]						=   [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__":
       __UpperCamelCase				:  List[Any]      					=							argparse.ArgumentParser()
       __UpperCamelCase				:  Optional[Any]      					=							[
           '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',
       )

       __UpperCamelCase				:  int      					=							parser.parse_args()

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