'''simple docstring'''





import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD





def      lowerCAmelCase_    ( )   ->				int:


     '''simple docstring'''




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

     return image





def      lowerCAmelCase_    ( snake_case_	:			Any      )   ->				int:


     '''simple docstring'''




     UpperCAmelCase_	   =       []
     # 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_    ( snake_case_	:			Optional[int] ,   snake_case_	:			int ,   snake_case_	:			int      )   ->				Dict:


     '''simple docstring'''




     UpperCAmelCase_	   =       dct.pop(snake_case_      )
     UpperCAmelCase_	   =       val





def      lowerCAmelCase_    ( snake_case_	:			Optional[Any] ,   snake_case_	:			Optional[Any]      )   ->				List[Any]:


     '''simple docstring'''




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

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





def      lowerCAmelCase_    ( snake_case_	:			List[str] ,   snake_case_	:			Optional[int]      )   ->				Dict:


     '''simple docstring'''




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

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

     return config, image_size





@torch.no_grad()
def      lowerCAmelCase_    ( snake_case_	:			int ,   snake_case_	:			Optional[int]=None ,   snake_case_	:			Optional[Any]=False      )   ->				int:


     '''simple docstring'''




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

     UpperCAmelCase_	   =       BlipaForConditionalGeneration(snake_case_      ).eval()

     UpperCAmelCase_	   =       {
         "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_	   =       model_name_to_original[model_name]

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

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

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

     # read in qv biases
     read_in_q_v_bias(snake_case_ ,   snake_case_      )

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

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

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

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

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

     print("Generating a caption..."      )
     UpperCAmelCase_	   =       ""
     UpperCAmelCase_	   =       tokenizer(snake_case_ ,   return_tensors="pt"      ).input_ids.to(snake_case_      )

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

     if pytorch_dump_folder_path is not None:
          processor.save_pretrained(snake_case_      )
          hf_model.save_pretrained(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__":
    SCREAMING_SNAKE_CASE_:					Any  =argparse.ArgumentParser()
    SCREAMING_SNAKE_CASE_:					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',
    )

    SCREAMING_SNAKE_CASE_:					List[str]  =parser.parse_args()

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





'''simple docstring'''





from typing import List, Union

import numpy as np

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, logging
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, ChunkPipeline


SCREAMING_SNAKE_CASE_:					Dict  =logging.get_logger(__name__)




class 	__A				(			UpperCamelCase__     ):



    def  _lowercase						(self   :   Dict       ,      __a   :   Any			):
         if isinstance(__a       ,      __a			):
              UpperCAmelCase_	   =       [label.strip() for label in labels.split(","			) if label.strip()]
         return labels






    def __call__(self   :   Union[str, Any]       ,      __a   :   Optional[int]       ,      __a   :   Optional[int]       ,      __a   :   int			):
         if len(__a			) == 0 or len(__a			) == 0:
              raise ValueError("You must include at least one label and at least one sequence."			)
         if hypothesis_template.format(labels[0]			) == hypothesis_template:
              raise ValueError(
                  (
                      "The provided hypothesis_template \"{}\" was not able to be formatted with the target labels. "
                      "Make sure the passed template includes formatting syntax such as {{}} where the label should go."
                  ).format(__a			)			)

         if isinstance(__a       ,      __a			):
              UpperCAmelCase_	   =       [sequences]

         UpperCAmelCase_	   =       []
         for sequence in sequences:
              sequence_pairs.extend([[sequence, hypothesis_template.format(__a			)] for label in labels]			)

         return sequence_pairs, sequences




@add_end_docstrings(UpperCamelCase__     )
class 	__A				(			UpperCamelCase__     ):



    def __init__(self   :   Optional[Any]       ,      __a   :   Union[str, Any]=ZeroShotClassificationArgumentHandler()       ,      *__a   :   Optional[int]       ,      **__a   :   List[str]			):
         UpperCAmelCase_	   =       args_parser
         super().__init__(*__a       ,      **__a			)
         if self.entailment_id == -1:
              logger.warning(
                  "Failed to determine 'entailment' label id from the label2id mapping in the model config. Setting to "
                  "-1. Define a descriptive label2id mapping in the model config to ensure correct outputs."			)



    @property
    def  _lowercase						(self   :   str			):
         for label, ind in self.model.config.labelaid.items():
              if label.lower().startswith("entail"			):
                   return ind
         return -1



    def  _lowercase						(self   :   Any       ,      __a   :   Any       ,      __a   :   int=True       ,      __a   :   Dict=True       ,      __a   :   Any=TruncationStrategy.ONLY_FIRST       ,      **__a   :   Tuple			):
         UpperCAmelCase_	   =       self.framework
         if self.tokenizer.pad_token is None:
              # Override for tokenizers not supporting padding
              logger.error(
                  "Tokenizer was not supporting padding necessary for zero-shot, attempting to use "
                  " `pad_token=eos_token`"			)
              UpperCAmelCase_	   =       self.tokenizer.eos_token
         try:
              UpperCAmelCase_	   =       self.tokenizer(
                  __a       ,      add_special_tokens=__a       ,      return_tensors=__a       ,      padding=__a       ,      truncation=__a       ,      )
         except Exception as e:
              if "too short" in str(__a			):
                   # tokenizers might yell that we want to truncate
                   # to a value that is not even reached by the input.
                   # In that case we don't want to truncate.
                   # It seems there's not a really better way to catch that
                   # exception.

                   UpperCAmelCase_	   =       self.tokenizer(
                       __a       ,      add_special_tokens=__a       ,      return_tensors=__a       ,      padding=__a       ,      truncation=TruncationStrategy.DO_NOT_TRUNCATE       ,      )
              else:
                   raise e

         return inputs



    def  _lowercase						(self   :   List[str]       ,      **__a   :   Tuple			):
         if kwargs.get("multi_class"       ,      __a			) is not None:
              UpperCAmelCase_	   =       kwargs["multi_class"]
              logger.warning(
                  "The `multi_class` argument has been deprecated and renamed to `multi_label`. "
                  "`multi_class` will be removed in a future version of Transformers."			)
         UpperCAmelCase_	   =       {}
         if "candidate_labels" in kwargs:
              UpperCAmelCase_	   =       self._args_parser._parse_labels(kwargs["candidate_labels"]			)
         if "hypothesis_template" in kwargs:
              UpperCAmelCase_	   =       kwargs["hypothesis_template"]

         UpperCAmelCase_	   =       {}
         if "multi_label" in kwargs:
              UpperCAmelCase_	   =       kwargs["multi_label"]
         return preprocess_params, {}, postprocess_params



    def __call__(self   :   Tuple       ,      __a   :   Union[str, List[str]]       ,      *__a   :   Optional[Any]       ,      **__a   :   Tuple       ,      ):
         if len(__a			) == 0:
              pass
         elif len(__a			) == 1 and "candidate_labels" not in kwargs:
              UpperCAmelCase_	   =       args[0]
         else:
              raise ValueError(f"""Unable to understand extra arguments {args}"""			)

         return super().__call__(__a       ,      **__a			)



    def  _lowercase						(self   :   Optional[int]       ,      __a   :   Optional[Any]       ,      __a   :   List[str]=None       ,      __a   :   Any="This example is {}."			):
         UpperCAmelCase_   ,			UpperCAmelCase_	   =       self._args_parser(__a       ,      __a       ,      __a			)

         for i, (candidate_label, sequence_pair) in enumerate(zip(__a       ,      __a			)			):
              UpperCAmelCase_	   =       self._parse_and_tokenize([sequence_pair]			)

              yield {
                  "candidate_label": candidate_label,
                  "sequence": sequences[0],
                  "is_last": i == len(__a			) - 1,
                  **model_input,
              }



    def  _lowercase						(self   :   List[str]       ,      __a   :   Any			):
         UpperCAmelCase_	   =       inputs["candidate_label"]
         UpperCAmelCase_	   =       inputs["sequence"]
         UpperCAmelCase_	   =       {k: inputs[k] for k in self.tokenizer.model_input_names}
         UpperCAmelCase_	   =       self.model(**__a			)

         UpperCAmelCase_	   =       {
             "candidate_label": candidate_label,
             "sequence": sequence,
             "is_last": inputs["is_last"],
             **outputs,
         }
         return model_outputs






    def  _lowercase						(self   :   Optional[Any]       ,      __a   :   List[str]       ,      __a   :   Tuple=False			):
         UpperCAmelCase_	   =       [outputs["candidate_label"] for outputs in model_outputs]
         UpperCAmelCase_	   =       [outputs["sequence"] for outputs in model_outputs]
         UpperCAmelCase_	   =       np.concatenate([output["logits"].numpy() for output in model_outputs]			)
         UpperCAmelCase_	   =       logits.shape[0]
         UpperCAmelCase_	   =       len(__a			)
         UpperCAmelCase_	   =       N // n
         UpperCAmelCase_	   =       logits.reshape((num_sequences, n, -1)			)

         if multi_label or len(__a			) == 1:
              # softmax over the entailment vs. contradiction dim for each label independently
              UpperCAmelCase_	   =       self.entailment_id
              UpperCAmelCase_	   =       -1 if entailment_id == 0 else 0
              UpperCAmelCase_	   =       reshaped_outputs[..., [contradiction_id, entailment_id]]
              UpperCAmelCase_	   =       np.exp(__a			) / np.exp(__a			).sum(-1       ,      keepdims=__a			)
              UpperCAmelCase_	   =       scores[..., 1]
         else:
              # softmax the "entailment" logits over all candidate labels
              UpperCAmelCase_	   =       reshaped_outputs[..., self.entailment_id]
              UpperCAmelCase_	   =       np.exp(__a			) / np.exp(__a			).sum(-1       ,      keepdims=__a			)

         UpperCAmelCase_	   =       list(reversed(scores[0].argsort()			)			)
         return {
             "sequence": sequences[0],
             "labels": [candidate_labels[i] for i in top_inds],
             "scores": scores[0, top_inds].tolist(),
         }




