


"""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  __snake_case       (   )      ->		Any:







		"""simple docstring"""





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

		return image







def  __snake_case       (   UpperCamelCase__	)      ->		Any:







		"""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.layernorm.weight')	)
		rename_keys.append(('Qformer.bert.embeddings.LayerNorm.bias', 'qformer.layernorm.bias')	)

		# fmt: on
		return rename_keys







def  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__ ,			UpperCamelCase__	)      ->		Tuple:







		"""simple docstring"""





		A					       = dct.pop(UpperCamelCase__	)
		A					       = val







def  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__	)      ->		Dict:







		"""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  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__	)      ->		str:







		"""simple docstring"""





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

		A					       = BlipaConfig(vision_config=UpperCamelCase__ ,			text_config=UpperCamelCase__	)

		return config, image_size







@torch.no_grad()
def  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__=None ,			UpperCamelCase__=False	)      ->		Tuple:







		"""simple docstring"""





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

		A					       = BlipaForConditionalGeneration(UpperCamelCase__	).eval()

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

		A       , A					       = model_name_to_original[model_name]

		# load original model
		print('Loading original model...'	)
		A					       = 'cuda' 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 "opt_proj" in key:
						A					       = key.replace('opt_proj' ,			'language_projection'	)
				if "t5_proj" in key:
						A					       = key.replace('t5_proj' ,			'language_projection'	)
				if key.startswith('opt'	):
						A					       = key.replace('opt' ,			'language'	)
				if key.startswith('t5'	):
						A					       = key.replace('t5' ,			'language'	)
				A					       = val

		# read in qv biases
		read_in_q_v_bias(UpperCamelCase__ ,			UpperCamelCase__	)

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

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

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

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

		original_model.to(UpperCamelCase__	)
		hf_model.to(UpperCamelCase__	)
		with torch.no_grad():
				if "opt" in model_name:
						A					       = original_model({'image': original_pixel_values, 'text_input': ['']}	).logits
						A					       = hf_model(UpperCamelCase__ ,			UpperCamelCase__	).logits
				else:
						A					       = original_model(
						    {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}	).logits
						A					       = input_ids.masked_fill(input_ids == tokenizer.pad_token_id ,			-100	)
						A					       = hf_model(UpperCamelCase__ ,			UpperCamelCase__ ,			labels=UpperCamelCase__	).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":
				A					       = 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=UpperCamelCase__	)
				assert torch.allclose(logits[0, :3, :3] ,			UpperCamelCase__ ,			atol=1E-4	)
		elif model_name == "blip2-flan-t5-xl-coco":
				A					       = 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=UpperCamelCase__	)
		else:
				# cast to same type
				A					       = logits.dtype
				assert torch.allclose(original_logits.to(UpperCamelCase__	) ,			UpperCamelCase__ ,			atol=1E-2	)
		print('Looks ok!'	)

		print('Generating a caption...'	)
		A					       = ''
		A					       = tokenizer(UpperCamelCase__ ,			return_tensors='pt'	).input_ids.to(UpperCamelCase__	)

		A					       = original_model.generate({'image': original_pixel_values}	)
		A					       = hf_model.generate(
		    UpperCamelCase__ ,			UpperCamelCase__ ,			do_sample=UpperCamelCase__ ,			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:' ,			UpperCamelCase__	)
		A					       = input_ids.shape[1]
		A					       = processor.batch_decode(outputs[:, prompt_length:] ,			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'nielsr/{model_name}'	)
				hf_model.push_to_hub(f'nielsr/{model_name}'	)


if __name__ == "__main__":
	UpperCamelCase						:				Optional[int]             =      argparse.ArgumentParser()
	UpperCamelCase						:				str             =      [
	    "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						:				str             =      parser.parse_args()

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









"""simple docstring"""


import argparse
import json
import os

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.deepspeed import DummyOptim, DummyScheduler


UpperCamelCase						:				int             =      16
UpperCamelCase						:				Optional[int]             =      32







def  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__ = 16 ,			UpperCamelCase__ = "bert-base-cased"	)      ->		Optional[int]:







		"""simple docstring"""





		A					       = AutoTokenizer.from_pretrained(UpperCamelCase__	)
		A					       = load_dataset('glue' ,			'mrpc'	)

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

		# Apply the method we just defined to all the examples in all the splits of the dataset
		A					       = datasets.map(
		    UpperCamelCase__ ,			batched=UpperCamelCase__ ,			remove_columns=['idx', 'sentence1', 'sentence2'] ,			load_from_cache_file=UpperCamelCase__	)

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

		def collate_fn(UpperCamelCase__	):
				# On TPU it's best to pad everything to the same length or training will be very slow.
				if accelerator.distributed_type == DistributedType.TPU:
						return tokenizer.pad(UpperCamelCase__ ,			padding='max_length' ,			max_length=128 ,			return_tensors='pt'	)
				return tokenizer.pad(UpperCamelCase__ ,			padding='longest' ,			return_tensors='pt'	)

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

		return train_dataloader, eval_dataloader







def  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__	)      ->		Optional[Any]:







		"""simple docstring"""





		A					       = Accelerator()

		# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
		A					       = config['lr']
		A					       = int(config['num_epochs']	)
		A					       = int(config['seed']	)
		A					       = int(config['batch_size']	)
		A					       = args.model_name_or_path

		set_seed(UpperCamelCase__	)
		A       , A					       = get_dataloaders(UpperCamelCase__ ,			UpperCamelCase__ ,			UpperCamelCase__	)

		# Instantiate the model (we build the model here so that the seed also control new weights initialization)
		A					       = AutoModelForSequenceClassification.from_pretrained(UpperCamelCase__ ,			return_dict=UpperCamelCase__	)

		# Instantiate optimizer
		A					       = (
		    AdamW
		    if accelerator.state.deepspeed_plugin is None
		    or 'optimizer' not in accelerator.state.deepspeed_plugin.deepspeed_config
		    else DummyOptim
		)
		A					       = optimizer_cls(params=model.parameters() ,			lr=UpperCamelCase__	)

		if accelerator.state.deepspeed_plugin is not None:
				A					       = accelerator.state.deepspeed_plugin.deepspeed_config[
				    'gradient_accumulation_steps'
				]
		else:
				A					       = 1
		A					       = (len(UpperCamelCase__	) * num_epochs) // gradient_accumulation_steps

		# Instantiate scheduler
		if (
		    accelerator.state.deepspeed_plugin is None
		    or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
		):
				A					       = get_linear_schedule_with_warmup(
				    optimizer=UpperCamelCase__ ,			num_warmup_steps=0 ,			num_training_steps=UpperCamelCase__ ,			)
		else:
				A					       = DummyScheduler(UpperCamelCase__ ,			total_num_steps=UpperCamelCase__ ,			warmup_num_steps=0	)

		# 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.
		A       , A       , A       , A       , A					       = accelerator.prepare(
		    UpperCamelCase__ ,			UpperCamelCase__ ,			UpperCamelCase__ ,			UpperCamelCase__ ,			UpperCamelCase__	)

		# We need to keep track of how many total steps we have iterated over
		A					       = 0
		# We also need to keep track of the stating epoch so files are named properly
		A					       = 0

		# Now we train the model
		A					       = evaluate.load('glue' ,			'mrpc'	)
		A					       = 0
		A					       = {}
		for epoch in range(UpperCamelCase__ ,			UpperCamelCase__	):
				model.train()
				for step, batch in enumerate(UpperCamelCase__	):
						A					       = model(**UpperCamelCase__	)
						A					       = outputs.loss
						A					       = loss / gradient_accumulation_steps
						accelerator.backward(UpperCamelCase__	)
						if step % gradient_accumulation_steps == 0:
								optimizer.step()
								lr_scheduler.step()
								optimizer.zero_grad()

						overall_step += 1

				model.eval()
				A					       = 0
				for step, batch in enumerate(UpperCamelCase__	):
						# We could avoid this line since we set the accelerator with `device_placement=True`.
						batch.to(accelerator.device	)
						with torch.no_grad():
								A					       = model(**UpperCamelCase__	)
						A					       = outputs.logits.argmax(dim=-1	)
						# It is slightly faster to call this once, than multiple times
						A       , A					       = accelerator.gather(
						    (predictions, batch['labels'])	)  # If we are in a multiprocess environment, the last batch has duplicates
						if accelerator.use_distributed:
								if step == len(UpperCamelCase__	) - 1:
										A					       = predictions[: len(eval_dataloader.dataset	) - samples_seen]
										A					       = references[: len(eval_dataloader.dataset	) - samples_seen]
								else:
										samples_seen += references.shape[0]
						metric.add_batch(
						    predictions=UpperCamelCase__ ,			references=UpperCamelCase__ ,			)

				A					       = metric.compute()
				# Use accelerator.print to print only on the main process.
				accelerator.print(f'epoch {epoch}:' ,			UpperCamelCase__	)
				A					       = eval_metric['accuracy']

				if best_performance < eval_metric["accuracy"]:
						A					       = eval_metric['accuracy']

		if args.performance_lower_bound is not None:
				assert (
				    args.performance_lower_bound <= best_performance
				), f'Best performance metric {best_performance} is lower than the lower bound {args.performance_lower_bound}'

		accelerator.wait_for_everyone()
		if accelerator.is_main_process:
				with open(os.path.join(args.output_dir ,			'all_results.json'	) ,			'w'	) as f:
						json.dump(UpperCamelCase__ ,			UpperCamelCase__	)







def  __snake_case       (   )      ->		Optional[int]:







		"""simple docstring"""





		A					       = argparse.ArgumentParser(description='Simple example of training script tracking peak GPU memory usage.'	)
		parser.add_argument(
		    '--model_name_or_path' ,			type=UpperCamelCase__ ,			default='bert-base-cased' ,			help='Path to pretrained model or model identifier from huggingface.co/models.' ,			required=UpperCamelCase__ ,			)
		parser.add_argument(
		    '--output_dir' ,			type=UpperCamelCase__ ,			default='.' ,			help='Optional save directory where all checkpoint folders will be stored. Default is the current working directory.' ,			)
		parser.add_argument(
		    '--performance_lower_bound' ,			type=UpperCamelCase__ ,			default=UpperCamelCase__ ,			help='Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.' ,			)
		parser.add_argument(
		    '--num_epochs' ,			type=UpperCamelCase__ ,			default=3 ,			help='Number of train epochs.' ,			)
		A					       = parser.parse_args()
		A					       = {'lr': 2E-5, 'num_epochs': args.num_epochs, 'seed': 42, 'batch_size': 16}
		training_function(UpperCamelCase__ ,			UpperCamelCase__	)


if __name__ == "__main__":
	main()





