


from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import (
           TFRoFormerForCausalLM,
           TFRoFormerForMaskedLM,
           TFRoFormerForMultipleChoice,
           TFRoFormerForQuestionAnswering,
           TFRoFormerForSequenceClassification,
           TFRoFormerForTokenClassification,
           TFRoFormerModel,
       )
       from transformers.models.roformer.modeling_tf_roformer import (
           TFRoFormerSelfAttention,
           TFRoFormerSinusoidalPositionalEmbedding,
       )





class  lowerCAmelCase__:






    '''simple docstring'''





    def __init__(	self				:   Any       ,   __snake_case				:   Optional[int]       ,   __snake_case				:   Optional[int]=13       ,   __snake_case				:   Optional[int]=7       ,   __snake_case				:   Union[str, Any]=True       ,   __snake_case				:   str=True       ,   __snake_case				:   List[Any]=True       ,   __snake_case				:   Union[str, Any]=True       ,   __snake_case				:   Union[str, Any]=99       ,   __snake_case				:   Tuple=32       ,   __snake_case				:   str=2       ,   __snake_case				:   str=4       ,   __snake_case				:   Union[str, Any]=37       ,   __snake_case				:   Dict="gelu"       ,   __snake_case				:   str=0.1       ,   __snake_case				:   List[Any]=0.1       ,   __snake_case				:   Any=512       ,   __snake_case				:   Any=16       ,   __snake_case				:   str=2       ,   __snake_case				:   List[str]=0.02       ,   __snake_case				:   Dict=3       ,   __snake_case				:   Tuple=4       ,   __snake_case				:   str=None       ,   ):





       '''simple docstring'''
       UpperCAmelCase_							:      Tuple						=   parent
       UpperCAmelCase_							:      Dict						=   13
       UpperCAmelCase_							:      List[Any]						=   7
       UpperCAmelCase_							:      str						=   True
       UpperCAmelCase_							:      str						=   True
       UpperCAmelCase_							:      Union[str, Any]						=   True
       UpperCAmelCase_							:      List[str]						=   True
       UpperCAmelCase_							:      Union[str, Any]						=   99
       UpperCAmelCase_							:      Dict						=   32
       UpperCAmelCase_							:      Tuple						=   2
       UpperCAmelCase_							:      List[str]						=   4
       UpperCAmelCase_							:      int						=   37
       UpperCAmelCase_							:      Optional[Any]						=   '''gelu'''
       UpperCAmelCase_							:      List[str]						=   0.1
       UpperCAmelCase_							:      Dict						=   0.1
       UpperCAmelCase_							:      int						=   512
       UpperCAmelCase_							:      Tuple						=   16
       UpperCAmelCase_							:      Optional[Any]						=   2
       UpperCAmelCase_							:      Optional[Any]						=   0.02
       UpperCAmelCase_							:      List[Any]						=   3
       UpperCAmelCase_							:      List[Any]						=   4
       UpperCAmelCase_							:      Any						=   None





    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   ids_tensor([self.batch_size, self.seq_length]       ,   self.vocab_size	)

       UpperCAmelCase_							:      Union[str, Any]						=   None
       if self.use_input_mask:
          UpperCAmelCase_							:      Tuple						=   random_attention_mask([self.batch_size, self.seq_length]	)

       UpperCAmelCase_							:      Any						=   None
       if self.use_token_type_ids:
          UpperCAmelCase_							:      List[str]						=   ids_tensor([self.batch_size, self.seq_length]       ,   self.type_vocab_size	)

       UpperCAmelCase_							:      List[str]						=   None
       UpperCAmelCase_							:      Union[str, Any]						=   None
       UpperCAmelCase_							:      Union[str, Any]						=   None
       if self.use_labels:
          UpperCAmelCase_							:      List[str]						=   ids_tensor([self.batch_size]       ,   self.type_sequence_label_size	)
          UpperCAmelCase_							:      List[str]						=   ids_tensor([self.batch_size, self.seq_length]       ,   self.num_labels	)
          UpperCAmelCase_							:      Any						=   ids_tensor([self.batch_size]       ,   self.num_choices	)

       UpperCAmelCase_							:      Dict						=   RoFormerConfig(
           vocab_size=self.vocab_size       ,   hidden_size=self.hidden_size       ,   num_hidden_layers=self.num_hidden_layers       ,   num_attention_heads=self.num_attention_heads       ,   intermediate_size=self.intermediate_size       ,   hidden_act=self.hidden_act       ,   hidden_dropout_prob=self.hidden_dropout_prob       ,   attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,   max_position_embeddings=self.max_position_embeddings       ,   type_vocab_size=self.type_vocab_size       ,   initializer_range=self.initializer_range       ,   return_dict=__snake_case       ,   )

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





    def 					_lowerCamelCase (	self				:   List[str]       ,   __snake_case				:   List[Any]       ,   __snake_case				:   Optional[int]       ,   __snake_case				:   Optional[Any]       ,   __snake_case				:   int       ,   __snake_case				:   List[str]       ,   __snake_case				:   Optional[int]       ,   __snake_case				:   int	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   TFRoFormerModel(config=__snake_case	)
       UpperCAmelCase_							:      Tuple						=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}

       UpperCAmelCase_							:      Dict						=   [input_ids, input_mask]
       UpperCAmelCase_							:      Optional[Any]						=   model(__snake_case	)

       UpperCAmelCase_							:      Any						=   model(__snake_case	)

       self.parent.assertEqual(result.last_hidden_state.shape       ,   (self.batch_size, self.seq_length, self.hidden_size)	)





    def 					_lowerCamelCase (	self				:   Union[str, Any]       ,   __snake_case				:   Optional[int]       ,   __snake_case				:   int       ,   __snake_case				:   Optional[int]       ,   __snake_case				:   int       ,   __snake_case				:   Optional[Any]       ,   __snake_case				:   Any       ,   __snake_case				:   Union[str, Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   True
       UpperCAmelCase_							:      Union[str, Any]						=   TFRoFormerForCausalLM(config=__snake_case	)
       UpperCAmelCase_							:      List[str]						=   {
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''token_type_ids''': token_type_ids,
       }
       UpperCAmelCase_							:      Dict						=   model(__snake_case	)['''logits''']
       self.parent.assertListEqual(
           list(prediction_scores.numpy().shape	)       ,   [self.batch_size, self.seq_length, self.vocab_size]	)





    def 					_lowerCamelCase (	self				:   str       ,   __snake_case				:   Tuple       ,   __snake_case				:   Any       ,   __snake_case				:   Dict       ,   __snake_case				:   Union[str, Any]       ,   __snake_case				:   Any       ,   __snake_case				:   Optional[Any]       ,   __snake_case				:   List[str]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   TFRoFormerForMaskedLM(config=__snake_case	)
       UpperCAmelCase_							:      int						=   {
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''token_type_ids''': token_type_ids,
       }
       UpperCAmelCase_							:      Tuple						=   model(__snake_case	)
       self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.seq_length, self.vocab_size)	)





    def 					_lowerCamelCase (	self				:   Optional[int]       ,   __snake_case				:   Tuple       ,   __snake_case				:   Any       ,   __snake_case				:   Optional[int]       ,   __snake_case				:   Any       ,   __snake_case				:   Optional[Any]       ,   __snake_case				:   Union[str, Any]       ,   __snake_case				:   Tuple	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   self.num_labels
       UpperCAmelCase_							:      Any						=   TFRoFormerForSequenceClassification(config=__snake_case	)
       UpperCAmelCase_							:      Optional[Any]						=   {
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''token_type_ids''': token_type_ids,
       }

       UpperCAmelCase_							:      int						=   model(__snake_case	)
       self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.num_labels)	)





    def 					_lowerCamelCase (	self				:   Optional[Any]       ,   __snake_case				:   str       ,   __snake_case				:   Optional[Any]       ,   __snake_case				:   List[Any]       ,   __snake_case				:   int       ,   __snake_case				:   int       ,   __snake_case				:   str       ,   __snake_case				:   Any	):





       '''simple docstring'''
       UpperCAmelCase_							:      Tuple						=   self.num_choices
       UpperCAmelCase_							:      Optional[Any]						=   TFRoFormerForMultipleChoice(config=__snake_case	)
       UpperCAmelCase_							:      Union[str, Any]						=   tf.tile(tf.expand_dims(__snake_case       ,   1	)       ,   (1, self.num_choices, 1)	)
       UpperCAmelCase_							:      str						=   tf.tile(tf.expand_dims(__snake_case       ,   1	)       ,   (1, self.num_choices, 1)	)
       UpperCAmelCase_							:      int						=   tf.tile(tf.expand_dims(__snake_case       ,   1	)       ,   (1, self.num_choices, 1)	)
       UpperCAmelCase_							:      List[str]						=   {
           '''input_ids''': multiple_choice_inputs_ids,
           '''attention_mask''': multiple_choice_input_mask,
           '''token_type_ids''': multiple_choice_token_type_ids,
       }
       UpperCAmelCase_							:      Tuple						=   model(__snake_case	)
       self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.num_choices)	)





    def 					_lowerCamelCase (	self				:   int       ,   __snake_case				:   Any       ,   __snake_case				:   int       ,   __snake_case				:   str       ,   __snake_case				:   List[str]       ,   __snake_case				:   Dict       ,   __snake_case				:   Dict       ,   __snake_case				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   self.num_labels
       UpperCAmelCase_							:      str						=   TFRoFormerForTokenClassification(config=__snake_case	)
       UpperCAmelCase_							:      Union[str, Any]						=   {
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''token_type_ids''': token_type_ids,
       }
       UpperCAmelCase_							:      List[str]						=   model(__snake_case	)
       self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.seq_length, self.num_labels)	)





    def 					_lowerCamelCase (	self				:   str       ,   __snake_case				:   Union[str, Any]       ,   __snake_case				:   int       ,   __snake_case				:   Dict       ,   __snake_case				:   Dict       ,   __snake_case				:   List[Any]       ,   __snake_case				:   List[Any]       ,   __snake_case				:   Optional[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Tuple						=   TFRoFormerForQuestionAnswering(config=__snake_case	)
       UpperCAmelCase_							:      Any						=   {
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''token_type_ids''': token_type_ids,
       }

       UpperCAmelCase_							:      Dict						=   model(__snake_case	)
       self.parent.assertEqual(result.start_logits.shape       ,   (self.batch_size, self.seq_length)	)
       self.parent.assertEqual(result.end_logits.shape       ,   (self.batch_size, self.seq_length)	)





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[int]						=   self.prepare_config_and_inputs()
       (
           (
           UpperCAmelCase_
       )		,      (
           UpperCAmelCase_
       )		,      (
           UpperCAmelCase_
       )		,      (
           UpperCAmelCase_
       )		,      (
           UpperCAmelCase_
       )		,      (
           UpperCAmelCase_
       )		,      (
           UpperCAmelCase_
       )		,      
       )							:      int						=   config_and_inputs
       UpperCAmelCase_							:      Union[str, Any]						=   {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict





@require_tf
class  lowerCAmelCase__(    snake_case__    ,    snake_case__    ,    unittest.TestCase						):






    '''simple docstring'''





    A_  :     Dict			= (
        (
            TFRoFormerModel,
            TFRoFormerForCausalLM,
            TFRoFormerForMaskedLM,
            TFRoFormerForQuestionAnswering,
            TFRoFormerForSequenceClassification,
            TFRoFormerForTokenClassification,
            TFRoFormerForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    A_  :     List[str]			= (
        {
            'feature-extraction': TFRoFormerModel,
            'fill-mask': TFRoFormerForMaskedLM,
            'question-answering': TFRoFormerForQuestionAnswering,
            'text-classification': TFRoFormerForSequenceClassification,
            'text-generation': TFRoFormerForCausalLM,
            'token-classification': TFRoFormerForTokenClassification,
            'zero-shot': TFRoFormerForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )

    A_  :     Tuple			= False
    A_  :     List[str]			= False





    def 					_lowerCamelCase (	self				:   Optional[Any]       ,   __snake_case				:   Optional[Any]       ,   __snake_case				:   Optional[int]       ,   __snake_case				:   str       ,   __snake_case				:   List[Any]       ,   __snake_case				:   str	):





       '''simple docstring'''
       if pipeline_test_casse_name == "TextGenerationPipelineTests":
          return True

       return False





    def 					_lowerCamelCase (	self				:   Tuple	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[Any]						=   TFRoFormerModelTester(self	)
       UpperCAmelCase_							:      Optional[int]						=   ConfigTester(self       ,   config_class=__snake_case       ,   hidden_size=37	)





    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       self.config_tester.run_common_tests()





    def 					_lowerCamelCase (	self				:   int	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*__snake_case	)





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*__snake_case	)





    def 					_lowerCamelCase (	self				:   List[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_lm_head(*__snake_case	)





    def 					_lowerCamelCase (	self				:   int	):





       '''simple docstring'''
       UpperCAmelCase_							:      Dict						=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*__snake_case	)





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      Dict						=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*__snake_case	)





    def 					_lowerCamelCase (	self				:   Union[str, Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[str]						=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*__snake_case	)





    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      Tuple						=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*__snake_case	)





    @slow
    def 					_lowerCamelCase (	self				:   Union[str, Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   TFRoFormerModel.from_pretrained('''junnyu/roformer_chinese_base'''	)
       self.assertIsNotNone(__snake_case	)





@require_tf
class  lowerCAmelCase__(    unittest.TestCase						):






    '''simple docstring'''





    @slow
    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   TFRoFormerForMaskedLM.from_pretrained('''junnyu/roformer_chinese_base'''	)
       UpperCAmelCase_							:      Dict						=   tf.constant([[0, 1, 2, 3, 4, 5]]	)
       UpperCAmelCase_							:      Dict						=   model(__snake_case	)[0]

       # TODO Replace vocab size
       UpperCAmelCase_							:      List[Any]						=   50_000

       UpperCAmelCase_							:      Tuple						=   [1, 6, vocab_size]
       self.assertEqual(output.shape       ,   __snake_case	)

       print(output[:, :3, :3]	)

       # TODO Replace values below with what was printed above.
       UpperCAmelCase_							:      Optional[Any]						=   tf.constant(
           [
               [
                   [-0.12_053_341, -1.0_264_901, 0.29_221_946],
                   [-1.5_133_783, 0.197_433, 0.15_190_607],
                   [-5.0_135_403, -3.900_256, -0.84_038_764],
               ]
           ]	)
       tf.debugging.assert_near(output[:, :3, :3]       ,   __snake_case       ,   atol=1E-4	)





@require_tf
class  lowerCAmelCase__(    unittest.TestCase						):






    '''simple docstring'''





    A_  :     Union[str, Any]			= 1E-4





    def 					_lowerCamelCase (	self				:   List[str]	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[Any]						=   tf.constant([[4, 10]]	)
       UpperCAmelCase_							:      List[str]						=   TFRoFormerSinusoidalPositionalEmbedding(num_positions=6       ,   embedding_dim=6	)

       UpperCAmelCase_							:      Any						=   emba(input_ids.shape	)
       UpperCAmelCase_							:      Optional[int]						=   tf.constant(
           [[0.0_000, 0.0_000, 0.0_000, 1.0_000, 1.0_000, 1.0_000], [0.8_415, 0.0_464, 0.0_022, 0.5_403, 0.9_989, 1.0_000]]	)

       tf.debugging.assert_near(__snake_case       ,   __snake_case       ,   atol=self.tolerance	)





    def 					_lowerCamelCase (	self				:   Union[str, Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Any						=   tf.constant(
           [
               [0.0_000, 0.0_000, 0.0_000, 0.0_000, 0.0_000],
               [0.8_415, 0.8_219, 0.8_020, 0.7_819, 0.7_617],
               [0.9_093, 0.9_364, 0.9_581, 0.9_749, 0.9_870],
           ]	)
       UpperCAmelCase_							:      Tuple						=   TFRoFormerSinusoidalPositionalEmbedding(num_positions=512       ,   embedding_dim=512	)
       emba([2, 16, 512]	)
       UpperCAmelCase_							:      str						=   emba.weight[:3, :5]

       tf.debugging.assert_near(__snake_case       ,   __snake_case       ,   atol=self.tolerance	)





@require_tf
class  lowerCAmelCase__(    unittest.TestCase						):






    '''simple docstring'''





    A_  :     Tuple			= 1E-4





    def 					_lowerCamelCase (	self				:   List[str]	):





       '''simple docstring'''
       # 2,12,16,64
       UpperCAmelCase_							:      str						=   tf.reshape(tf.range(2 * 12 * 16 * 64       ,   dtype=tf.floataa	)       ,   shape=(2, 12, 16, 64)	) / 100

       UpperCAmelCase_							:      Union[str, Any]						=   -tf.reshape(tf.range(2 * 12 * 16 * 64       ,   dtype=tf.floataa	)       ,   shape=(2, 12, 16, 64)	) / 100

       UpperCAmelCase_							:      Dict						=   TFRoFormerSinusoidalPositionalEmbedding(num_positions=32       ,   embedding_dim=64	)
       UpperCAmelCase_							:      Optional[int]						=   embed_positions([2, 16, 768]	)[None, None, :, :]

       UpperCAmelCase_		,      UpperCAmelCase_							:      Dict						=   TFRoFormerSelfAttention.apply_rotary_position_embeddings(
           __snake_case       ,   __snake_case       ,   __snake_case	)

       UpperCAmelCase_							:      Tuple						=   tf.constant(
           [
               [0.0_000, 0.0_100, 0.0_200, 0.0_300, 0.0_400, 0.0_500, 0.0_600, 0.0_700],
               [-0.2_012, 0.8_897, 0.0_263, 0.9_401, 0.2_074, 0.9_463, 0.3_481, 0.9_343],
               [-1.7_057, 0.6_271, -1.2_145, 1.3_897, -0.6_303, 1.7_647, -0.1_173, 1.8_985],
               [-2.1_731, -1.6_397, -2.7_358, 0.2_854, -2.1_840, 1.7_183, -1.3_018, 2.4_871],
               [0.2_717, -3.6_173, -2.9_206, -2.1_988, -3.6_638, 0.3_858, -2.9_155, 2.2_980],
               [3.9_859, -2.1_580, -0.7_984, -4.4_904, -4.1_181, -2.0_252, -4.4_782, 1.1_253],
           ]	)
       UpperCAmelCase_							:      Any						=   tf.constant(
           [
               [0.0_000, -0.0_100, -0.0_200, -0.0_300, -0.0_400, -0.0_500, -0.0_600, -0.0_700],
               [0.2_012, -0.8_897, -0.0_263, -0.9_401, -0.2_074, -0.9_463, -0.3_481, -0.9_343],
               [1.7_057, -0.6_271, 1.2_145, -1.3_897, 0.6_303, -1.7_647, 0.1_173, -1.8_985],
               [2.1_731, 1.6_397, 2.7_358, -0.2_854, 2.1_840, -1.7_183, 1.3_018, -2.4_871],
               [-0.2_717, 3.6_173, 2.9_206, 2.1_988, 3.6_638, -0.3_858, 2.9_155, -2.2_980],
               [-3.9_859, 2.1_580, 0.7_984, 4.4_904, 4.1_181, 2.0_252, 4.4_782, -1.1_253],
           ]	)

       tf.debugging.assert_near(query_layer[0, 0, :6, :8]       ,   __snake_case       ,   atol=self.tolerance	)
       tf.debugging.assert_near(key_layer[0, 0, :6, :8]       ,   __snake_case       ,   atol=self.tolerance	)



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)