





'''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  a__    (    )		->		Optional[int]:





     """simple docstring"""

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

     return image







def  a__    (    lowercase :					Union[str, Any]				)		->		Optional[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  a__    (    lowercase :					Optional[int], lowercase :					List[Any], lowercase :					List[str]				)		->		List[str]:





     """simple docstring"""

     _UpperCamelCase       =	dct.pop(lowercase				)
     _UpperCamelCase       =	val







def  a__    (    lowercase :					Optional[Any], lowercase :					List[str]				)		->		int:





     """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(lowercase, requires_grad=lowercase				), v_bias)				)
          _UpperCamelCase       =	qkv_bias







def  a__    (    lowercase :					Optional[int], lowercase :					Optional[Any]				)		->		Tuple:





     """simple docstring"""

     _UpperCamelCase       =	364 if '''coco''' in model_name else 224
     _UpperCamelCase       =	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       =	OPTConfig.from_pretrained('''facebook/opt-2.7b''', eos_token_id=lowercase				).to_dict()
     elif "opt-6.7b" in model_name:
          _UpperCamelCase       =	OPTConfig.from_pretrained('''facebook/opt-6.7b''', eos_token_id=lowercase				).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=lowercase, text_config=lowercase				)

     return config, image_size







@torch.no_grad()
def  a__    (    lowercase :					Tuple, lowercase :					Tuple=None, lowercase :					List[str]=False				)		->		Optional[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=lowercase				).input_ids[0]
     _UpperCamelCase   ,    _UpperCamelCase       =	get_blipa_config(lowercase, eos_token_id=lowercase				)

     _UpperCamelCase       =	BlipaForConditionalGeneration(lowercase				).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=lowercase, model_type=lowercase, is_eval=lowercase, device=lowercase				)
     original_model.eval()
     print('''Done!'''				)

     # update state dict keys
     _UpperCamelCase       =	original_model.state_dict()
     _UpperCamelCase       =	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       =	state_dict.pop(lowercase				)
          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(lowercase, lowercase				)

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

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

     # create processor
     _UpperCamelCase       =	BlipImageProcessor(
         size={'''height''': image_size, '''width''': image_size}, image_mean=lowercase, image_std=lowercase				)
     _UpperCamelCase       =	BlipaProcessor(image_processor=lowercase, tokenizer=lowercase				)
     _UpperCamelCase       =	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       =	original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}				).logits
               _UpperCamelCase       =	hf_model(lowercase, lowercase				).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, -100				)
               _UpperCamelCase       =	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       =	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       =	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       =	logits.dtype
          assert torch.allclose(original_logits.to(lowercase				), lowercase, atol=1e-2				)
     print('''Looks ok!'''				)

     print('''Generating a caption...'''				)
     _UpperCamelCase       =	''''''
     _UpperCamelCase       =	tokenizer(lowercase, return_tensors='''pt'''				).input_ids.to(lowercase				)

     _UpperCamelCase       =	original_model.generate({'''image''': original_pixel_values}				)
     _UpperCamelCase       =	hf_model.generate(
         lowercase, lowercase, do_sample=lowercase, 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:''', lowercase				)
     _UpperCamelCase       =	input_ids.shape[1]
     _UpperCamelCase       =	processor.batch_decode(outputs[:, prompt_length:], skip_special_tokens=lowercase				)
     _UpperCamelCase       =	[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__":
   lowercase__						:      List[Any]	=    argparse.ArgumentParser()
   lowercase__						:      Optional[int]	=    [
       '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',
   )

   lowercase__						:      Any	=    parser.parse_args()

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










'''simple docstring'''

import argparse
import os

# New Code #
import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


lowercase__						:      str	=    16
lowercase__						:      int	=    32







def  a__    (    lowercase :					Accelerator, lowercase :					int = 16				)		->		List[str]:





     """simple docstring"""

     _UpperCamelCase       =	AutoTokenizer.from_pretrained('''bert-base-cased'''				)
     _UpperCamelCase       =	load_dataset('''glue''', '''mrpc'''				)

     def tokenize_function(lowercase :					Union[str, Any]				):
          # max_length=None => use the model max length (it's actually the default)
          _UpperCamelCase       =	tokenizer(examples['''sentence1'''], examples['''sentence2'''], truncation=lowercase, max_length=lowercase				)
          return outputs

     # Apply the method we just defined to all the examples in all the splits of the dataset
     # starting with the main process first:
     with accelerator.main_process_first():
          _UpperCamelCase       =	datasets.map(
              lowercase, batched=lowercase, remove_columns=['''idx''', '''sentence1''', '''sentence2'''], )

     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
     _UpperCamelCase       =	tokenized_datasets.rename_column('''label''', '''labels'''				)

     def collate_fn(lowercase :					Any				):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          _UpperCamelCase       =	128 if accelerator.distributed_type == DistributedType.TPU else None
          # When using mixed precision we want round multiples of 8/16
          if accelerator.mixed_precision == "fp8":
               _UpperCamelCase       =	16
          elif accelerator.mixed_precision != "no":
               _UpperCamelCase       =	8
          else:
               _UpperCamelCase       =	None

          return tokenizer.pad(
              lowercase, padding='''longest''', max_length=lowercase, pad_to_multiple_of=lowercase, return_tensors='''pt''', )

     # Instantiate dataloaders.
     _UpperCamelCase       =	DataLoader(
         tokenized_datasets['''train'''], shuffle=lowercase, collate_fn=lowercase, batch_size=lowercase				)
     _UpperCamelCase       =	DataLoader(
         tokenized_datasets['''validation'''], shuffle=lowercase, collate_fn=lowercase, batch_size=lowercase				)

     return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('TESTING_MOCKED_DATALOADERS', None) == "1":
   from accelerate.test_utils.training import mocked_dataloaders

   lowercase__						:      str	=    mocked_dataloaders  # noqa: F811







def  a__    (    lowercase :					List[Any], lowercase :					List[str]				)		->		Any:





     """simple docstring"""

     if os.environ.get('''TESTING_MOCKED_DATALOADERS''', lowercase				) == "1":
          _UpperCamelCase       =	2
     # Initialize accelerator
     _UpperCamelCase       =	Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision				)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     _UpperCamelCase       =	config['''lr''']
     _UpperCamelCase       =	int(config['''num_epochs''']				)
     _UpperCamelCase       =	int(config['''seed''']				)
     _UpperCamelCase       =	int(config['''batch_size''']				)

     _UpperCamelCase       =	evaluate.load('''glue''', '''mrpc'''				)

     # New Code #
     # We now can define an inner training loop function. It should take a batch size as the only parameter,
     # and build the dataloaders in there.
     # It also gets our decorator
     @find_executable_batch_size(starting_batch_size=lowercase				)
     def inner_training_loop(lowercase :					List[Any]				):
          # And now just move everything below under this function
          # We need to bring in the Accelerator object from earlier
          nonlocal accelerator
          # And reset all of its attributes that could hold onto any memory:
          accelerator.free_memory()

          # Then we can declare the model, optimizer, and everything else:
          set_seed(lowercase				)

          # Instantiate the model (we build the model here so that the seed also control new weights initialization)
          _UpperCamelCase       =	AutoModelForSequenceClassification.from_pretrained('''bert-base-cased''', return_dict=lowercase				)

          # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
          # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
          # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
          _UpperCamelCase       =	model.to(accelerator.device				)

          # Instantiate optimizer
          _UpperCamelCase       =	AdamW(params=model.parameters(), lr=lowercase				)
          _UpperCamelCase   ,    _UpperCamelCase       =	get_dataloaders(lowercase, lowercase				)

          # Instantiate scheduler
          _UpperCamelCase       =	get_linear_schedule_with_warmup(
              optimizer=lowercase, num_warmup_steps=100, num_training_steps=(len(lowercase				) * num_epochs), )

          # Prepare everything
          # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
          # prepare method.
          _UpperCamelCase   ,    _UpperCamelCase   ,    _UpperCamelCase   ,    _UpperCamelCase   ,    _UpperCamelCase       =	accelerator.prepare(
              lowercase, lowercase, lowercase, lowercase, lowercase				)

          # Now we train the model
          for epoch in range(lowercase				):
               model.train()
               for step, batch in enumerate(lowercase				):
                    # We could avoid this line since we set the accelerator with `device_placement=True`.
                    batch.to(accelerator.device				)
                    _UpperCamelCase       =	model(**lowercase				)
                    _UpperCamelCase       =	outputs.loss
                    accelerator.backward(lowercase				)
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()

               model.eval()
               for step, batch in enumerate(lowercase				):
                    # We could avoid this line since we set the accelerator with `device_placement=True`.
                    batch.to(accelerator.device				)
                    with torch.no_grad():
                         _UpperCamelCase       =	model(**lowercase				)
                    _UpperCamelCase       =	outputs.logits.argmax(dim=-1				)
                    _UpperCamelCase   ,    _UpperCamelCase       =	accelerator.gather_for_metrics((predictions, batch['''labels'''])				)
                    metric.add_batch(
                        predictions=lowercase, references=lowercase, )

               _UpperCamelCase       =	metric.compute()
               # Use accelerator.print to print only on the main process.
               accelerator.print(F"""epoch {epoch}:""", lowercase				)

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
     inner_training_loop()







def  a__    (    )		->		str:





     """simple docstring"""

     _UpperCamelCase       =	argparse.ArgumentParser(description='''Simple example of training script.'''				)
     parser.add_argument(
         '''--mixed_precision''', type=lowercase, default=lowercase, choices=['''no''', '''fp16''', '''bf16''', '''fp8'''], help='''Whether to use mixed precision. Choose'''
         '''between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'''
         '''and an Nvidia Ampere GPU.''', )
     parser.add_argument('''--cpu''', action='''store_true''', help='''If passed, will train on the CPU.'''				)
     _UpperCamelCase       =	parser.parse_args()
     _UpperCamelCase       =	{'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
     training_function(lowercase, lowercase				)


if __name__ == "__main__":
   main()



