


import logging
import os
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

from filelock import FileLock

from transformers import PreTrainedTokenizer, is_tf_available, is_torch_available


UpperCAmelCase      =   logging.getLogger(__name__)







@dataclass
class lowerCAmelCase_						:

 '''simple docstring'''

 __snake_case												=  42
 __snake_case												=  42
 __snake_case												=  42







@dataclass
class lowerCAmelCase_						:

 '''simple docstring'''

 __snake_case												=  42
 __snake_case												=  42
 __snake_case												=  None
 __snake_case												=  None







class lowerCAmelCase_						(			lowerCamelCase__	):

 '''simple docstring'''

 __snake_case												=  "train"
 __snake_case												=  "dev"
 __snake_case												=  "test"







class lowerCAmelCase_						:

 '''simple docstring'''

 @staticmethod
 def 						UpperCamelCase__					(     _UpperCAmelCase		,				_UpperCAmelCase ):
       raise NotImplementedError



 @staticmethod
 def 						UpperCamelCase__					(     _UpperCAmelCase ):
       raise NotImplementedError





 @staticmethod
 def 						UpperCamelCase__					(     _UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase=False		,				_UpperCAmelCase="[CLS]"		,				_UpperCAmelCase=1		,				_UpperCAmelCase="[SEP]"		,				_UpperCAmelCase=False		,				_UpperCAmelCase=False		,				_UpperCAmelCase=0		,				_UpperCAmelCase=0		,				_UpperCAmelCase=-1_00		,				_UpperCAmelCase=0		,				_UpperCAmelCase=True		,				):
       snake_case_	=				{label: i for i, label in enumerate(_UpperCAmelCase )}

       snake_case_	=				[]
       for ex_index, example in enumerate(_UpperCAmelCase ):
             if ex_index % 1_00_00 == 0:
                   logger.info('''Writing example %d of %d'''		,				_UpperCAmelCase		,				len(_UpperCAmelCase ) )

             snake_case_	=				[]
             snake_case_	=				[]
             for word, label in zip(example.words		,				example.labels ):
                   snake_case_	=				tokenizer.tokenize(_UpperCAmelCase )

                   # bert-base-multilingual-cased sometimes output "nothing ([]) when calling tokenize with just a space.
                   if len(_UpperCAmelCase ) > 0:
                         tokens.extend(_UpperCAmelCase )
                         # Use the real label id for the first token of the word, and padding ids for the remaining tokens
                         label_ids.extend([label_map[label]] + [pad_token_label_id] * (len(_UpperCAmelCase ) - 1) )

            # Account for [CLS] and [SEP] with "- 2" and with "- 3" for RoBERTa.
             snake_case_	=				tokenizer.num_special_tokens_to_add()
             if len(_UpperCAmelCase ) > max_seq_length - special_tokens_count:
                   snake_case_	=				tokens[: (max_seq_length - special_tokens_count)]
                   snake_case_	=				label_ids[: (max_seq_length - special_tokens_count)]

             # The convention in BERT is:
             # (a) For sequence pairs:
             #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
             #  type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1
             # (b) For single sequences:
             #  tokens:   [CLS] the dog is hairy . [SEP]
             #  type_ids:   0   0   0   0  0     0   0
             #
             # Where "type_ids" are used to indicate whether this is the first
             # sequence or the second sequence. The embedding vectors for `type=0` and
             # `type=1` were learned during pre-training and are added to the wordpiece
             # embedding vector (and position vector). This is not *strictly* necessary
             # since the [SEP] token unambiguously separates the sequences, but it makes
             # it easier for the model to learn the concept of sequences.
             #
             # For classification tasks, the first vector (corresponding to [CLS]) is
             # used as the "sentence vector". Note that this only makes sense because
             # the entire model is fine-tuned.
             tokens += [sep_token]
             label_ids += [pad_token_label_id]
             if sep_token_extra:
                   # roberta uses an extra separator b/w pairs of sentences
                   tokens += [sep_token]
                   label_ids += [pad_token_label_id]
             snake_case_	=				[sequence_a_segment_id] * len(_UpperCAmelCase )

             if cls_token_at_end:
                   tokens += [cls_token]
                   label_ids += [pad_token_label_id]
                   segment_ids += [cls_token_segment_id]
             else:
                   snake_case_	=				[cls_token] + tokens
                   snake_case_	=				[pad_token_label_id] + label_ids
                   snake_case_	=				[cls_token_segment_id] + segment_ids

             snake_case_	=				tokenizer.convert_tokens_to_ids(_UpperCAmelCase )

             # The mask has 1 for real tokens and 0 for padding tokens. Only real
             # tokens are attended to.
             snake_case_	=				[1 if mask_padding_with_zero else 0] * len(_UpperCAmelCase )

             # Zero-pad up to the sequence length.
             snake_case_	=				max_seq_length - len(_UpperCAmelCase )
             if pad_on_left:
                   snake_case_	=				([pad_token] * padding_length) + input_ids
                   snake_case_	=				([0 if mask_padding_with_zero else 1] * padding_length) + input_mask
                   snake_case_	=				([pad_token_segment_id] * padding_length) + segment_ids
                   snake_case_	=				([pad_token_label_id] * padding_length) + label_ids
             else:
                   input_ids += [pad_token] * padding_length
                   input_mask += [0 if mask_padding_with_zero else 1] * padding_length
                   segment_ids += [pad_token_segment_id] * padding_length
                   label_ids += [pad_token_label_id] * padding_length

             assert len(_UpperCAmelCase ) == max_seq_length
             assert len(_UpperCAmelCase ) == max_seq_length
             assert len(_UpperCAmelCase ) == max_seq_length
             assert len(_UpperCAmelCase ) == max_seq_length

             if ex_index < 5:
                   logger.info('''*** Example ***''' )
                   logger.info('''guid: %s'''		,				example.guid )
                   logger.info('''tokens: %s'''		,				''' '''.join([str(_UpperCAmelCase ) for x in tokens] ) )
                   logger.info('''input_ids: %s'''		,				''' '''.join([str(_UpperCAmelCase ) for x in input_ids] ) )
                   logger.info('''input_mask: %s'''		,				''' '''.join([str(_UpperCAmelCase ) for x in input_mask] ) )
                   logger.info('''segment_ids: %s'''		,				''' '''.join([str(_UpperCAmelCase ) for x in segment_ids] ) )
                   logger.info('''label_ids: %s'''		,				''' '''.join([str(_UpperCAmelCase ) for x in label_ids] ) )

             if "token_type_ids" not in tokenizer.model_input_names:
                   snake_case_	=				None

             features.append(
                 InputFeatures(
                     input_ids=_UpperCAmelCase		,				attention_mask=_UpperCAmelCase		,				token_type_ids=_UpperCAmelCase		,				label_ids=_UpperCAmelCase ) )
       return features


if is_torch_available():
   import torch
   from torch import nn
   from torch.utils.data import Dataset







   class lowerCAmelCase_						(			lowerCamelCase__	):

    '''simple docstring'''

    __snake_case												=  42
    __snake_case												=  nn.CrossEntropyLoss().ignore_index



    def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase = None		,				_UpperCAmelCase=False		,				_UpperCAmelCase = Split.train		,				):
          # Load data features from cache or dataset file
          snake_case_	=				os.path.join(
              _UpperCAmelCase		,				'''cached_{}_{}_{}'''.format(mode.value		,				tokenizer.__class__.__name__		,				str(_UpperCAmelCase ) )		,				)

          # Make sure only the first process in distributed training processes the dataset,
          # and the others will use the cache.
          snake_case_	=				cached_features_file + '''.lock'''
          with FileLock(_UpperCAmelCase ):
                if os.path.exists(_UpperCAmelCase ) and not overwrite_cache:
                      logger.info(F'''Loading features from cached file {cached_features_file}''' )
                      snake_case_	=				torch.load(_UpperCAmelCase )
                else:
                      logger.info(F'''Creating features from dataset file at {data_dir}''' )
                      snake_case_	=				token_classification_task.read_examples_from_file(_UpperCAmelCase		,				_UpperCAmelCase )
                      # TODO clean up all this to leverage built-in features of tokenizers
                      snake_case_	=				token_classification_task.convert_examples_to_features(
                          _UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				cls_token_at_end=bool(model_type in ['''xlnet'''] )		,				cls_token=tokenizer.cls_token		,				cls_token_segment_id=2 if model_type in ['''xlnet'''] else 0		,				sep_token=tokenizer.sep_token		,				sep_token_extra=_UpperCAmelCase		,				pad_on_left=bool(tokenizer.padding_side == '''left''' )		,				pad_token=tokenizer.pad_token_id		,				pad_token_segment_id=tokenizer.pad_token_type_id		,				pad_token_label_id=self.pad_token_label_id		,				)
                      logger.info(F'''Saving features into cached file {cached_features_file}''' )
                      torch.save(self.features		,				_UpperCAmelCase )



    def __len__(     self ):
          return len(self.features )





    def __getitem__(     self		,				_UpperCAmelCase ):
          return self.features[i]


if is_tf_available():
   import tensorflow as tf







   class lowerCAmelCase_						:

    '''simple docstring'''

    __snake_case												=  42
    __snake_case												=  -1_0_0



    def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase = None		,				_UpperCAmelCase=False		,				_UpperCAmelCase = Split.train		,				):
          snake_case_	=				token_classification_task.read_examples_from_file(_UpperCAmelCase		,				_UpperCAmelCase )
          # TODO clean up all this to leverage built-in features of tokenizers
          snake_case_	=				token_classification_task.convert_examples_to_features(
              _UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				cls_token_at_end=bool(model_type in ['''xlnet'''] )		,				cls_token=tokenizer.cls_token		,				cls_token_segment_id=2 if model_type in ['''xlnet'''] else 0		,				sep_token=tokenizer.sep_token		,				sep_token_extra=_UpperCAmelCase		,				pad_on_left=bool(tokenizer.padding_side == '''left''' )		,				pad_token=tokenizer.pad_token_id		,				pad_token_segment_id=tokenizer.pad_token_type_id		,				pad_token_label_id=self.pad_token_label_id		,				)

          def gen():
                for ex in self.features:
                      if ex.token_type_ids is None:
                            yield (
                                {"input_ids": ex.input_ids, "attention_mask": ex.attention_mask},
                                ex.label_ids,
                            )
                      else:
                            yield (
                                {
                                    "input_ids": ex.input_ids,
                                    "attention_mask": ex.attention_mask,
                                    "token_type_ids": ex.token_type_ids,
                                },
                                ex.label_ids,
                            )

          if "token_type_ids" not in tokenizer.model_input_names:
                snake_case_	=				tf.data.Dataset.from_generator(
                    _UpperCAmelCase		,				({'''input_ids''': tf.intaa, '''attention_mask''': tf.intaa}, tf.intaa)		,				(
                        {'''input_ids''': tf.TensorShape([None] ), '''attention_mask''': tf.TensorShape([None] )},
                        tf.TensorShape([None] ),
                    )		,				)
          else:
                snake_case_	=				tf.data.Dataset.from_generator(
                    _UpperCAmelCase		,				({'''input_ids''': tf.intaa, '''attention_mask''': tf.intaa, '''token_type_ids''': tf.intaa}, tf.intaa)		,				(
                        {
                            '''input_ids''': tf.TensorShape([None] ),
                            '''attention_mask''': tf.TensorShape([None] ),
                            '''token_type_ids''': tf.TensorShape([None] ),
                        },
                        tf.TensorShape([None] ),
                    )		,				)



    def 						UpperCamelCase__					(     self ):
          snake_case_	=				self.dataset.apply(tf.data.experimental.assert_cardinality(len(self.features ) ) )

          return self.dataset



    def __len__(     self ):
          return len(self.features )





    def __getitem__(     self		,				_UpperCAmelCase ):
          return self.features[i]



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




      """simple docstring"""

      snake_case_	=				'''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
      snake_case_	=				Image.open(requests.get(SCREAMING_SNAKE_CASE , stream=SCREAMING_SNAKE_CASE ).raw ).convert('''RGB''' )

      return image






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE )->  Any:




      """simple docstring"""

      snake_case_	=				[]
      # 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   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  Tuple:




      """simple docstring"""

      snake_case_	=				dct.pop(SCREAMING_SNAKE_CASE )
      snake_case_	=				val






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  Dict:




      """simple docstring"""

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

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






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  List[str]:




      """simple docstring"""

      snake_case_	=				364 if '''coco''' in model_name else 224
      snake_case_	=				BlipaVisionConfig(image_size=SCREAMING_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 "opt-2.7b" in model_name:
            snake_case_	=				OPTConfig.from_pretrained('''facebook/opt-2.7b''' , eos_token_id=SCREAMING_SNAKE_CASE ).to_dict()
      elif "opt-6.7b" in model_name:
            snake_case_	=				OPTConfig.from_pretrained('''facebook/opt-6.7b''' , eos_token_id=SCREAMING_SNAKE_CASE ).to_dict()
      elif "t5-xl" in model_name:
            snake_case_	=				TaConfig.from_pretrained('''google/flan-t5-xl''' , dense_act_fn='''gelu''' , bos_token_id=1 ).to_dict()
      elif "t5-xxl" in model_name:
            snake_case_	=				TaConfig.from_pretrained('''google/flan-t5-xxl''' , dense_act_fn='''gelu''' , bos_token_id=1 ).to_dict()

      snake_case_	=				BlipaConfig(vision_config=SCREAMING_SNAKE_CASE , text_config=SCREAMING_SNAKE_CASE )

      return config, image_size






@torch.no_grad()
def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE=None , SCREAMING_SNAKE_CASE=False )->  List[Any]:




      """simple docstring"""

      snake_case_	=				(
          AutoTokenizer.from_pretrained('''facebook/opt-2.7b''' )
          if '''opt''' in model_name
          else AutoTokenizer.from_pretrained('''google/flan-t5-xl''' )
      )
      snake_case_	=				tokenizer('''\n''' , add_special_tokens=SCREAMING_SNAKE_CASE ).input_ids[0]
      snake_case_	,   snake_case_	=				get_blipa_config(SCREAMING_SNAKE_CASE , eos_token_id=SCREAMING_SNAKE_CASE )

      snake_case_	=				BlipaForConditionalGeneration(SCREAMING_SNAKE_CASE ).eval()

      snake_case_	=				{
          '''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'''),
      }

      snake_case_	,   snake_case_	=				model_name_to_original[model_name]

      # load original model
      print('''Loading original model...''' )
      snake_case_	=				'''cuda''' if torch.cuda.is_available() else '''cpu'''
      snake_case_	,   snake_case_	,   snake_case_	=				load_model_and_preprocess(
          name=SCREAMING_SNAKE_CASE , model_type=SCREAMING_SNAKE_CASE , is_eval=SCREAMING_SNAKE_CASE , device=SCREAMING_SNAKE_CASE )
      original_model.eval()
      print('''Done!''' )

      # update state dict keys
      snake_case_	=				original_model.state_dict()
      snake_case_	=				create_rename_keys(SCREAMING_SNAKE_CASE )
      for src, dest in rename_keys:
            rename_key(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )

      # some keys can be renamed efficiently
      for key, val in state_dict.copy().items():
            snake_case_	=				state_dict.pop(SCREAMING_SNAKE_CASE )
            if key.startswith('''Qformer.bert''' ):
                  snake_case_	=				key.replace('''Qformer.bert''' , '''qformer''' )
            if "attention.self" in key:
                  snake_case_	=				key.replace('''self''' , '''attention''' )
            if "opt_proj" in key:
                  snake_case_	=				key.replace('''opt_proj''' , '''language_projection''' )
            if "t5_proj" in key:
                  snake_case_	=				key.replace('''t5_proj''' , '''language_projection''' )
            if key.startswith('''opt''' ):
                  snake_case_	=				key.replace('''opt''' , '''language''' )
            if key.startswith('''t5''' ):
                  snake_case_	=				key.replace('''t5''' , '''language''' )
            snake_case_	=				val

      # read in qv biases
      read_in_q_v_bias(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )

      snake_case_	,   snake_case_	=				hf_model.load_state_dict(SCREAMING_SNAKE_CASE , strict=SCREAMING_SNAKE_CASE )
      assert len(SCREAMING_SNAKE_CASE ) == 0
      assert unexpected_keys == ["qformer.embeddings.position_ids"]

      snake_case_	=				load_demo_image()
      snake_case_	=				vis_processors['''eval'''](SCREAMING_SNAKE_CASE ).unsqueeze(0 ).to(SCREAMING_SNAKE_CASE )
      snake_case_	=				tokenizer(['''\n'''] , return_tensors='''pt''' ).input_ids.to(SCREAMING_SNAKE_CASE )

      # create processor
      snake_case_	=				BlipImageProcessor(
          size={'''height''': image_size, '''width''': image_size} , image_mean=SCREAMING_SNAKE_CASE , image_std=SCREAMING_SNAKE_CASE )
      snake_case_	=				BlipaProcessor(image_processor=SCREAMING_SNAKE_CASE , tokenizer=SCREAMING_SNAKE_CASE )
      snake_case_	=				processor(images=SCREAMING_SNAKE_CASE , return_tensors='''pt''' ).pixel_values.to(SCREAMING_SNAKE_CASE )

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

      original_model.to(SCREAMING_SNAKE_CASE )
      hf_model.to(SCREAMING_SNAKE_CASE )
      with torch.no_grad():
            if "opt" in model_name:
                  snake_case_	=				original_model({'''image''': original_pixel_values, '''text_input''': ['''''']} ).logits
                  snake_case_	=				hf_model(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE ).logits
            else:
                  snake_case_	=				original_model(
                      {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']} ).logits
                  snake_case_	=				input_ids.masked_fill(input_ids == tokenizer.pad_token_id , -100 )
                  snake_case_	=				hf_model(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , labels=SCREAMING_SNAKE_CASE ).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":
            snake_case_	=				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=SCREAMING_SNAKE_CASE )
            assert torch.allclose(logits[0, :3, :3] , SCREAMING_SNAKE_CASE , atol=1E-4 )
      elif model_name == "blip2-flan-t5-xl-coco":
            snake_case_	=				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=SCREAMING_SNAKE_CASE )
      else:
            # cast to same type
            snake_case_	=				logits.dtype
            assert torch.allclose(original_logits.to(SCREAMING_SNAKE_CASE ) , SCREAMING_SNAKE_CASE , atol=1E-2 )
      print('''Looks ok!''' )

      print('''Generating a caption...''' )
      snake_case_	=				''''''
      snake_case_	=				tokenizer(SCREAMING_SNAKE_CASE , return_tensors='''pt''' ).input_ids.to(SCREAMING_SNAKE_CASE )

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

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

      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      =   argparse.ArgumentParser()
   UpperCAmelCase      =   [
       """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      =   parser.parse_args()

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