




'''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
lowerCAmelCase_     						=				pytest.mark.integration


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

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

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







def 						_A      ( UpperCAmelCase						):






   '''simple docstring'''







   @wraps(UpperCAmelCase						)
   def wrapper(self      ,UpperCAmelCase						):
      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						):






   '''simple docstring'''







   @wraps(UpperCAmelCase						)
   def wrapper(self      ,UpperCAmelCase						):
      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						):






   '''simple docstring'''







   @wraps(UpperCAmelCase						)
   def wrapper(self      ,UpperCAmelCase						):
      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      ( ):






   '''simple docstring'''







   A__       =     [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(
    UpperCAmelCase       , UpperCAmelCase       , UpperCAmelCase      )
@local
class 	_snake_case(			parameterized.TestCase      ):
 __snake_case:							Dict						=					{}
 __snake_case:							Tuple						=					None






 @pytest.mark.filterwarnings('ignore:metric_module_factory is deprecated:FutureWarning'    )
 @pytest.mark.filterwarnings('ignore:load_metric is deprecated:FutureWarning'    )
 def 						_UpperCamelCase (self       :  Any    ,		a       :  Dict    )       ->   Dict:






    """simple docstring"""


    A__       =     '[...]'
    A__       =     importlib.import_module(
        datasets.load.metric_module_factory(os.path.join('metrics'    ,		a    )    ).module_path    )
    A__       =     datasets.load.import_main_class(metric_module.__name__    ,		dataset=a    )
    # check parameters
    A__       =     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(a    ,		metric_module.__name__    ):
       with self.use_local_metrics():
          try:
             A__       =     doctest.testmod(a    ,		verbose=a    ,		raise_on_error=a    )
          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 						_UpperCamelCase (self       :  List[str]    ,		a       :  str    )       ->   str:






    """simple docstring"""


    A__       =     '[...]'
    A__       =     importlib.import_module(
        datasets.load.metric_module_factory(os.path.join('metrics'    ,		a    )    ).module_path    )
    # run doctest
    with self.use_local_metrics():
       A__       =     doctest.testmod(a    ,		verbose=a    ,		raise_on_error=a    )
    self.assertEqual(results.failed    ,		0    )
    self.assertGreater(results.attempted    ,		1    )






 @contextmanager
 def 						_UpperCamelCase (self       :  List[Any]    ,		a       :  str    ,		a       :  Union[str, Any]    )       ->   str:






    """simple docstring"""


    if metric_name in self.INTENSIVE_CALLS_PATCHER:
       with self.INTENSIVE_CALLS_PATCHER[metric_name](a    ):
          yield
    else:
       yield






 @contextmanager
 def 						_UpperCamelCase (self       :  List[Any]    )       ->   Tuple:






    """simple docstring"""


    def load_local_metric(a       :  Union[str, Any]    ,		*a       :  Any    ,		**a       :  Union[str, Any]    ):
       return load_metric(os.path.join('metrics'    ,		a    )    ,		*a    ,		**a    )

    with patch('datasets.load_metric'    ) as mock_load_metric:
       A__       =     load_local_metric
       yield






 @classmethod
 def 						_UpperCamelCase (cls       :  Dict    ,		a       :  Dict    )       ->   Dict:






    """simple docstring"""


    def wrapper(a       :  Optional[Any]    ):
       A__       =     contextmanager(a    )
       A__       =     patcher
       return patcher

    return wrapper







@LocalMetricTest.register_intensive_calls_patcher('bleurt'						)
def 						_A      ( UpperCAmelCase						):






   '''simple docstring'''







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

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

   class 	_snake_case(			UpperCAmelCase      ):






    def 						_UpperCamelCase (self       :  List[str]    ,		a       :  List[str]    )       ->   Tuple:






       """simple docstring"""


       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:
      A__       =     MockedPredictor()
      yield







@LocalMetricTest.register_intensive_calls_patcher('bertscore'						)
def 						_A      ( UpperCAmelCase						):






   '''simple docstring'''







   import torch

   def bert_cos_score_idf(UpperCAmelCase      ,UpperCAmelCase      ,*UpperCAmelCase      ,**UpperCAmelCase						):
      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:
      A__       =     bert_cos_score_idf
      yield







@LocalMetricTest.register_intensive_calls_patcher('comet'						)
def 						_A      ( UpperCAmelCase						):






   '''simple docstring'''







   def load_from_checkpoint(UpperCAmelCase						):

      class 	_snake_case:






       def 						_UpperCamelCase (self       :  List[str]    ,		a       :  Tuple    ,		*a       :  Any    ,		**a       :  List[str]    )       ->   Dict:






          """simple docstring"""


          assert len(a    ) == 2
          A__       =     [0.19, 0.92]
          return scores, sum(a    ) / len(a    )




      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:
      A__       =     None
      with patch('comet.load_from_checkpoint'						) as mock_load_from_checkpoint:
         A__       =     load_from_checkpoint
         yield







def 						_A      ( ):






   '''simple docstring'''







   A__       =     load_metric(os.path.join('metrics'      ,'seqeval'						)						)
   A__       =     'ERROR'
   A__       =     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						)








'''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      ( ):






   '''simple docstring'''







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

   return image







def 						_A      ( UpperCAmelCase						):






   '''simple docstring'''







   A__       =     []
   # 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      ,UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   A__       =     dct.pop(UpperCAmelCase						)
   A__       =     val







def 						_A      ( UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







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

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







def 						_A      ( UpperCAmelCase						):






   '''simple docstring'''







   A__       =     364 if 'coco' in model_name else 224
   A__       =     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:
      A__       =     TaConfig.from_pretrained('google/flan-t5-xl'      ,dense_act_fn='gelu'      ,bos_token_id=1						).to_dict()
   elif "t5-xxl" in model_name:
      A__       =     TaConfig.from_pretrained('google/flan-t5-xxl'      ,dense_act_fn='gelu'      ,bos_token_id=1						).to_dict()
   elif "vicuna-7b" in model_name:
      A__       =     LlamaConfig.from_pretrained('decapoda-research/llama-7b-hf'      ,vocab_size=32001						).to_dict()
   elif "vicuna-13b" in model_name:
      A__       =     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
   A__       =     InstructBlipQFormerConfig(vocab_size=30523						).to_dict()
   A__       =     InstructBlipConfig(vision_config=UpperCAmelCase      ,text_config=UpperCAmelCase      ,qformer_config=UpperCAmelCase						)

   return config, image_size







@torch.no_grad()
def 						_A      ( UpperCAmelCase      ,UpperCAmelCase=None      ,UpperCAmelCase=False						):






   '''simple docstring'''







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

   if "t5" in model_name:
      A__       =     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>"})
      A__       =     LlamaTokenizerFast.from_pretrained(
          'huggyllama/llama-7b'      ,truncation_side='left'      ,bos_token='</s>'      ,unk_token='</s>'						)
      tokenizer.add_special_tokens({'pad_token': '[PAD]'}						)

   A__  ,			A__       =     get_blipa_config(UpperCAmelCase						)
   A__       =     InstructBlipForConditionalGeneration(UpperCAmelCase						).eval()

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

   A__  ,			A__       =     model_name_to_original[model_name]

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

   # update state dict keys
   A__       =     original_model.state_dict()
   A__       =     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():
      A__       =     state_dict.pop(UpperCAmelCase						)
      if key.startswith('Qformer.bert'						):
         A__       =     key.replace('Qformer.bert'      ,'qformer'						)
      if "attention.self" in key:
         A__       =     key.replace('self'      ,'attention'						)
      if "llm_proj" in key:
         A__       =     key.replace('llm_proj'      ,'language_projection'						)
      if "t5_proj" in key:
         A__       =     key.replace('t5_proj'      ,'language_projection'						)
      if key.startswith('llm_model'						):
         A__       =     key.replace('llm_model'      ,'language_model'						)
      if key.startswith('t5'						):
         A__       =     key.replace('t5'      ,'language'						)
      A__       =     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						)

   A__       =     load_demo_image()
   A__       =     'What is unusual about this image?'

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

   # make sure processor creates exact same pixel values
   A__       =     vis_processors['eval'](UpperCAmelCase						).unsqueeze(0						).to(UpperCAmelCase						)
   A__       =     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:
         A__       =     original_model({'image': original_pixel_values, 'text_input': [prompt]}						).logits
         A__       =     hf_model(**UpperCAmelCase						).logits
      else:
         A__       =     original_model(
             {'image': original_pixel_values, 'text_input': [prompt], 'text_output': ['\n']}						).logits
         A__       =     tokenizer('\n'      ,return_tensors='pt'						).input_ids.to(UpperCAmelCase						)
         A__       =     label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id      ,-100						)
         A__       =     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
   A__       =     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...'						)
   A__       =     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...'						)
   A__       =     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?
      A__       =     2
   print('Original generation:'      ,UpperCAmelCase						)
   A__       =     processor.batch_decode(UpperCAmelCase      ,skip_special_tokens=UpperCAmelCase						)
   A__       =     [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__":
    lowerCAmelCase_     						=				argparse.ArgumentParser()
    lowerCAmelCase_     						=				[
        '''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''',
    )

    lowerCAmelCase_     						=				parser.parse_args()

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


