

"""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_float32 (there's also the fix_lavis branch)
# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
# same for Vicuna-13b
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipImageProcessor,
    InstructBlipConfig,
    InstructBlipForConditionalGeneration,
    InstructBlipProcessor,
    InstructBlipQFormerConfig,
    InstructBlipVisionConfig,
    LlamaConfig,
    LlamaTokenizerFast,
    TaConfig,
    TaTokenizerFast,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD




def   _A  (  )     ->						Optional[Any]:
   '''simple docstring'''


   __lowercase								=						"https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"
   __lowercase								=						Image.open(requests.get(UpperCamelCase_,      stream=UpperCamelCase_).raw).convert("RGB")

   return image




def   _A  (  UpperCamelCase_					:	List[Any])     ->						Optional[Any]:
   '''simple docstring'''


   __lowercase								=						[]
   # 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.embeddings.layernorm.weight"))
   rename_keys.append(("Qformer.bert.embeddings.LayerNorm.bias", "qformer.embeddings.layernorm.bias"))

   # fmt: on
   return rename_keys




def   _A  (  UpperCamelCase_					:	Optional[Any],      UpperCamelCase_					:	Any,      UpperCamelCase_					:	List[Any])     ->						Any:
   '''simple docstring'''


   __lowercase								=						dct.pop(UpperCamelCase_)
   __lowercase								=						val




def   _A  (  UpperCamelCase_					:	Union[str, Any],      UpperCamelCase_					:	int)     ->						int:
   '''simple docstring'''


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

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




def   _A  (  UpperCamelCase_					:	List[Any])     ->						List[str]:
   '''simple docstring'''


   __lowercase								=						364 if "coco" in model_name else 224
   __lowercase								=						InstructBlipVisionConfig(image_size=UpperCamelCase_).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 "t5-xl" in model_name:
      __lowercase								=						TaConfig.from_pretrained("google/flan-t5-xl",      dense_act_fn="gelu",      bos_token_id=1).to_dict()
   elif "t5-xxl" in model_name:
      __lowercase								=						TaConfig.from_pretrained("google/flan-t5-xxl",      dense_act_fn="gelu",      bos_token_id=1).to_dict()
   elif "vicuna-7b" in model_name:
      __lowercase								=						LlamaConfig.from_pretrained("decapoda-research/llama-7b-hf",      vocab_size=32001).to_dict()
   elif "vicuna-13b" in model_name:
      __lowercase								=						LlamaConfig.from_pretrained("decapoda-research/llama-13b-hf",      vocab_size=32001).to_dict()
   else:
      raise ValueError("Model name not supported")

   # the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
   __lowercase								=						InstructBlipQFormerConfig(vocab_size=30523).to_dict()
   __lowercase								=						InstructBlipConfig(vision_config=UpperCamelCase_,      text_config=UpperCamelCase_,      qformer_config=UpperCamelCase_)

   return config, image_size




@torch.no_grad()
def   _A  (  UpperCamelCase_					:	List[str],      UpperCamelCase_					:	int=None,      UpperCamelCase_					:	Any=False)     ->						Tuple:
   '''simple docstring'''


   __lowercase								=						AutoTokenizer.from_pretrained("bert-base-uncased",      truncation_side="left")
   qformer_tokenizer.add_special_tokens({"bos_token": "[DEC]"})

   if "t5" in model_name:
      __lowercase								=						TaTokenizerFast.from_pretrained("google/flan-t5-xl",      truncation_side="left")
   elif "vicuna" in model_name:
      # the following was used in the original implementation:
      # tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
      # tokenizer.add_special_tokens({"pad_token": "[PAD]"})
      # tokenizer.add_special_tokens({"bos_token": "</s>"})
      # tokenizer.add_special_tokens({"eos_token": "</s>"})
      # tokenizer.add_special_tokens({"unk_token": "</s>"})
      __lowercase								=						LlamaTokenizerFast.from_pretrained(
          "huggyllama/llama-7b",      truncation_side="left",      bos_token="</s>",      unk_token="</s>")
      tokenizer.add_special_tokens({"pad_token": "[PAD]"})

   __lowercase		,__lowercase								=						get_blipa_config(UpperCamelCase_)
   __lowercase								=						InstructBlipForConditionalGeneration(UpperCamelCase_).eval()

   __lowercase								=						{
       "instructblip-vicuna-7b": ("blip2_vicuna_instruct", "vicuna7b"),
       "instructblip-vicuna-13b": ("blip2_vicuna_instruct", "vicuna13b"),
       "instructblip-flan-t5-xl": ("blip2_t5_instruct", "flant5xl"),
       "instructblip-flan-t5-xxl": ("blip2_t5_instruct", "flant5xxl"),
   }

   __lowercase		,__lowercase								=						model_name_to_original[model_name]

   # load original model
   print("Loading original model...")
   __lowercase								=						"cuda:1" if torch.cuda.is_available() else "cpu"
   __lowercase								=						"cuda:2" if torch.cuda.is_available() else "cpu"
   __lowercase		,__lowercase		,__lowercase								=						load_model_and_preprocess(
       name=UpperCamelCase_,      model_type=UpperCamelCase_,      is_eval=UpperCamelCase_,      device=UpperCamelCase_)
   original_model.eval()
   print("Done!")

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

   # some keys can be renamed efficiently
   for key, val in state_dict.copy().items():
      __lowercase								=						state_dict.pop(UpperCamelCase_)
      if key.startswith("Qformer.bert"):
         __lowercase								=						key.replace("Qformer.bert",      "qformer")
      if "attention.self" in key:
         __lowercase								=						key.replace("self",      "attention")
      if "llm_proj" in key:
         __lowercase								=						key.replace("llm_proj",      "language_projection")
      if "t5_proj" in key:
         __lowercase								=						key.replace("t5_proj",      "language_projection")
      if key.startswith("llm_model"):
         __lowercase								=						key.replace("llm_model",      "language_model")
      if key.startswith("t5"):
         __lowercase								=						key.replace("t5",      "language")
      __lowercase								=						val

   # read in qv biases
   read_in_q_v_bias(UpperCamelCase_,      UpperCamelCase_)

   # note: weights get loaded in torch.float32 by default
   hf_model.load_state_dict(UpperCamelCase_,      strict=UpperCamelCase_)

   __lowercase								=						load_demo_image()
   __lowercase								=						"What is unusual about this image?"

   # create processor
   __lowercase								=						BlipImageProcessor(
       size={"height": image_size, "width": image_size},      image_mean=UpperCamelCase_,      image_std=UpperCamelCase_)
   __lowercase								=						InstructBlipProcessor(
       image_processor=UpperCamelCase_,      tokenizer=UpperCamelCase_,      qformer_tokenizer=UpperCamelCase_,      )
   __lowercase								=						processor(images=UpperCamelCase_,      text=UpperCamelCase_,      return_tensors="pt").to(UpperCamelCase_)

   # make sure processor creates exact same pixel values
   __lowercase								=						vis_processors["eval"](UpperCamelCase_).unsqueeze(0).to(UpperCamelCase_)
   __lowercase								=						inputs.pixel_values
   assert torch.allclose(original_pixel_values.to(pixel_values.device),      UpperCamelCase_)

   original_model.to(UpperCamelCase_)
   hf_model.to(UpperCamelCase_)
   with torch.no_grad():
      if "vicuna" in model_name:
         __lowercase								=						original_model({"image": original_pixel_values, "text_input": [prompt]}).logits
         __lowercase								=						hf_model(**UpperCamelCase_).logits
      else:
         __lowercase								=						original_model(
             {"image": original_pixel_values, "text_input": [prompt], "text_output": ["\n"]}).logits
         __lowercase								=						tokenizer("\n",      return_tensors="pt").input_ids.to(UpperCamelCase_)
         __lowercase								=						label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id,      -100)
         __lowercase								=						hf_model(**UpperCamelCase_,      labels=UpperCamelCase_).logits

   print("First values of original logits:",      original_logits[0, :3, :3])
   print("First values of HF logits:",      logits[0, :3, :3])

   # assert values
   assert original_logits.shape == logits.shape
   __lowercase								=						1E-4 if "vicuna" in model_name else 1E-5
   assert torch.allclose(original_logits.to(logits.device),      UpperCamelCase_,      atol=UpperCamelCase_)
   print("Looks ok!")

   print("Generating with original model...")
   __lowercase								=						original_model.generate({"image": original_pixel_values, "prompt": prompt},      num_beams=5)

   # important: we need to cast the weights of the HF model to the appropriate type
   print("Generating with HF model...")
   __lowercase								=						hf_model.generate(
       **UpperCamelCase_,      do_sample=UpperCamelCase_,      num_beams=5,      max_length=256,      min_length=1,      top_p=0.9,      repetition_penalty=1.5,      length_penalty=1.0,      temperature=1,      )
   if "vicuna" in model_name:
      # convert output id 0 to 2 (eos_token_id)
      # TODO add this in the generate method?
      __lowercase								=						2
   print("Original generation:",      UpperCamelCase_)
   __lowercase								=						processor.batch_decode(UpperCamelCase_,      skip_special_tokens=UpperCamelCase_)
   __lowercase								=						[text.strip() for text in output_text]
   print("HF generation:",      UpperCamelCase_)

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

   if push_to_hub:
      processor.push_to_hub(F"""Salesforce/{model_name}""")
      hf_model.push_to_hub(F"""Salesforce/{model_name}""")


if __name__ == "__main__":
   _a							=							argparse.ArgumentParser()
   _a							=							[
       'instructblip-vicuna-7b',
       'instructblip-vicuna-13b',
       'instructblip-flan-t5-xl',
       'instructblip-flan-t5-xxl',
   ]
   parser.add_argument(
       '--model_name',
       default='instructblip-flan-t5-xl',
       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',
   )

   _a							=							parser.parse_args()

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







"""simple docstring"""





import doctest
import glob
import importlib
import inspect
import os
import re
from contextlib import contextmanager
from functools import wraps
from unittest.mock import patch

import numpy as np
import pytest
from absl.testing import parameterized

import datasets
from datasets import load_metric

from .utils import for_all_test_methods, local, slow


# mark all tests as integration
_a							=							pytest.mark.integration


_a							=							{'comet'}
_a							=							importlib.util.find_spec('fairseq') is not None

_a							=							{'code_eval'}
_a							=							os.name == 'nt'

_a							=							{'bertscore', 'frugalscore', 'perplexity'}
_a							=							importlib.util.find_spec('transformers') is not None




def   _A  (  UpperCamelCase_					:	Dict)     ->						Any:
   '''simple docstring'''


   @wraps(UpperCamelCase_)
   def wrapper(self					:	Dict,      UpperCamelCase_					:	Dict):
      if not _has_fairseq and metric_name in REQUIRE_FAIRSEQ:
         self.skipTest("\"test requires Fairseq\"")
      else:
         test_case(self,      UpperCamelCase_)

   return wrapper




def   _A  (  UpperCamelCase_					:	Dict)     ->						int:
   '''simple docstring'''


   @wraps(UpperCamelCase_)
   def wrapper(self					:	int,      UpperCamelCase_					:	str):
      if not _has_transformers and metric_name in REQUIRE_TRANSFORMERS:
         self.skipTest("\"test requires transformers\"")
      else:
         test_case(self,      UpperCamelCase_)

   return wrapper




def   _A  (  UpperCamelCase_					:	Tuple)     ->						str:
   '''simple docstring'''


   @wraps(UpperCamelCase_)
   def wrapper(self					:	Optional[int],      UpperCamelCase_					:	Optional[Any]):
      if _on_windows and metric_name in UNSUPPORTED_ON_WINDOWS:
         self.skipTest("\"test not supported on Windows\"")
      else:
         test_case(self,      UpperCamelCase_)

   return wrapper




def   _A  (  )     ->						str:
   '''simple docstring'''


   __lowercase								=						[metric_dir.split(os.sep)[-2] for metric_dir in glob.glob("./metrics/*/")]
   return [{"testcase_name": x, "metric_name": x} for x in metrics if x != "gleu"]  # gleu is unfinished




@parameterized.named_parameters(get_local_metric_names()       )
@for_all_test_methods(
    lowercase					,lowercase					,lowercase       )
@local
class  _lowerCAmelCase  (     parameterized.TestCase       ):







 """simple docstring"""







 __UpperCAmelCase     : Optional[int]       		=     {}
 __UpperCAmelCase     : Tuple       		=     None


 @pytest.mark.filterwarnings("ignore:metric_module_factory is deprecated:FutureWarning"     )
 @pytest.mark.filterwarnings("ignore:load_metric is deprecated:FutureWarning"     )
 def  _lowercase			(      self							:      Dict, UpperCAmelCase__							:      int     ):
    __lowercase								=						"[...]"
    __lowercase								=						importlib.import_module(
        datasets.load.metric_module_factory(os.path.join("metrics", UpperCAmelCase__     )     ).module_path     )
    __lowercase								=						datasets.load.import_main_class(metric_module.__name__, dataset=UpperCAmelCase__     )
    # check parameters
    __lowercase								=						inspect.signature(metric._compute     ).parameters
    self.assertTrue(all(p.kind != p.VAR_KEYWORD for p in parameters.values()     )     )  # no **kwargs
    # run doctest
    with self.patch_intensive_calls(UpperCAmelCase__, metric_module.__name__     ):
       with self.use_local_metrics():
          try:
             __lowercase								=						doctest.testmod(UpperCAmelCase__, verbose=UpperCAmelCase__, raise_on_error=UpperCAmelCase__     )
          except doctest.UnexpectedException as e:
             raise e.exc_info[1]  # raise the exception that doctest caught
    self.assertEqual(results.failed, 0     )
    self.assertGreater(results.attempted, 1     )


 @slow
 def  _lowercase			(      self							:      List[Any], UpperCAmelCase__							:      Optional[Any]     ):
    __lowercase								=						"[...]"
    __lowercase								=						importlib.import_module(
        datasets.load.metric_module_factory(os.path.join("metrics", UpperCAmelCase__     )     ).module_path     )
    # run doctest
    with self.use_local_metrics():
       __lowercase								=						doctest.testmod(UpperCAmelCase__, verbose=UpperCAmelCase__, raise_on_error=UpperCAmelCase__     )
    self.assertEqual(results.failed, 0     )
    self.assertGreater(results.attempted, 1     )


 @contextmanager
 def  _lowercase			(      self							:      List[Any], UpperCAmelCase__							:      Any, UpperCAmelCase__							:      Tuple     ):
    if metric_name in self.INTENSIVE_CALLS_PATCHER:
       with self.INTENSIVE_CALLS_PATCHER[metric_name](UpperCAmelCase__     ):
          yield
    else:
       yield


 @contextmanager
 def  _lowercase			(      self							:      List[Any]     ):
    def load_local_metric(UpperCAmelCase__							:      Any, *UpperCAmelCase__							:      List[Any], **UpperCAmelCase__							:      Any     ):
       return load_metric(os.path.join("metrics", UpperCAmelCase__     ), *UpperCAmelCase__, **UpperCAmelCase__     )

    with patch("datasets.load_metric"     ) as mock_load_metric:
       __lowercase								=						load_local_metric
       yield



 @classmethod
 def  _lowercase			(      cls							:      Optional[Any], UpperCAmelCase__							:      List[Any]     ):
    def wrapper(UpperCAmelCase__							:      Tuple     ):
       __lowercase								=						contextmanager(UpperCAmelCase__     )
       __lowercase								=						patcher
       return patcher

    return wrapper







@LocalMetricTest.register_intensive_calls_patcher("bleurt")
def   _A  (  UpperCamelCase_					:	Any)     ->						Optional[Any]:
   '''simple docstring'''


   import tensorflow.compat.va as tf
   from bleurt.score import Predictor

   tf.flags.DEFINE_string("sv",      "",      "")  # handle pytest cli flags




   class  _lowerCAmelCase  (     lowercase       ):







    """simple docstring"""







    def  _lowercase			(      self							:      Tuple, UpperCAmelCase__							:      Tuple     ):
       assert len(input_dict["input_ids"]     ) == 2
       return np.array([1.03, 1.04]     )

    # mock predict_fn which is supposed to do a forward pass with a bleurt model
   with patch("bleurt.score._create_predictor") as mock_create_predictor:
      __lowercase								=						MockedPredictor()
      yield




@LocalMetricTest.register_intensive_calls_patcher("bertscore")
def   _A  (  UpperCamelCase_					:	Tuple)     ->						int:
   '''simple docstring'''


   import torch

   def bert_cos_score_idf(UpperCamelCase_					:	Tuple,      UpperCamelCase_					:	str,      *UpperCamelCase_					:	Optional[Any],      **UpperCamelCase_					:	Dict):
      return torch.tensor([[1.0, 1.0, 1.0]] * len(UpperCamelCase_))

   # mock get_model which is supposed to do download a bert model
   # mock bert_cos_score_idf which is supposed to do a forward pass with a bert model
   with patch("bert_score.scorer.get_model"), patch(
       "bert_score.scorer.bert_cos_score_idf") as mock_bert_cos_score_idf:
      __lowercase								=						bert_cos_score_idf
      yield




@LocalMetricTest.register_intensive_calls_patcher("comet")
def   _A  (  UpperCamelCase_					:	Tuple)     ->						List[Any]:
   '''simple docstring'''


   def load_from_checkpoint(UpperCamelCase_					:	Tuple):




      class  _lowerCAmelCase  :







       """simple docstring"""







       def  _lowercase			(      self							:      str, UpperCAmelCase__							:      int, *UpperCAmelCase__							:      Dict, **UpperCAmelCase__							:      Dict     ):
          assert len(UpperCAmelCase__     ) == 2
          __lowercase								=						[0.19, 0.92]
          return scores, sum(UpperCAmelCase__     ) / len(UpperCAmelCase__     )







      return Model()

   # mock load_from_checkpoint which is supposed to do download a bert model
   # mock load_from_checkpoint which is supposed to do download a bert model
   with patch("comet.download_model") as mock_download_model:
      __lowercase								=						None
      with patch("comet.load_from_checkpoint") as mock_load_from_checkpoint:
         __lowercase								=						load_from_checkpoint
         yield




def   _A  (  )     ->						Tuple:
   '''simple docstring'''


   __lowercase								=						load_metric(os.path.join("metrics",      "seqeval"))
   __lowercase								=						"ERROR"
   __lowercase								=						F"""Scheme should be one of [IOB1, IOB2, IOE1, IOE2, IOBES, BILOU], got {wrong_scheme}"""
   with pytest.raises(UpperCamelCase_,      match=re.escape(UpperCamelCase_)):
      metric.compute(predictions=[],      references=[],      scheme=UpperCamelCase_)




