
import argparse
import dataclasses
import json
import logging
import os
import shutil
from typing import List, Optional

import datasets
from accelerate import Accelerator
from datasets import load_dataset
from finetuning import finetune
from tqdm.auto import tqdm

import transformers
from transformers import AutoConfig, set_seed
from transformers.trainer_utils import IntervalStrategy


_UpperCAmelCase :			Optional[int]              =	logging.getLogger(__name__)

_UpperCAmelCase :			Optional[Any]              =	"pytorch_model.bin"


@dataclasses.dataclass
class     lowercase :
		__lowercase	:	str     					=    dataclasses.field(
		    metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models."}       )
		__lowercase	:	Optional[str]     					=    dataclasses.field(
		    default=_SCREAMING_SNAKE_CASE  ,       metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co."}  ,       )






@dataclasses.dataclass
class     lowercase :
		__lowercase	:	str     					=    dataclasses.field(metadata={"help": "A csv or a json file containing the training data."}       )
		__lowercase	:	str     					=    dataclasses.field(metadata={"help": "A csv or a json file containing the data to predict on."}       )
		__lowercase	:	Optional[str]     					=    dataclasses.field(
		    default=_SCREAMING_SNAKE_CASE  ,       metadata={"help": "A csv or a json file containing the validation data."}       )
		__lowercase	:	Optional[str]     					=    dataclasses.field(
		    default=_SCREAMING_SNAKE_CASE  ,       metadata={"help": "The name of the task to train on."}  ,       )
		__lowercase	:	Optional[List[str]]     					=    dataclasses.field(
		    default=_SCREAMING_SNAKE_CASE  ,       metadata={"help": "The list of labels for the task."}       )






@dataclasses.dataclass
class     lowercase :
		__lowercase	:	str     					=    dataclasses.field(
		    metadata={"help": "The output directory where the model predictions and checkpoints will be written."}       )
		__lowercase	:	Optional[str]     					=    dataclasses.field(
		    default="accuracy"  ,       metadata={"help": "The evaluation metric used for the task."}       )
		__lowercase	:	Optional[str]     					=    dataclasses.field(
		    default="no"  ,       metadata={
		        "help": "The evaluation strategy to adopt during training. Possible values are: [\"no\", \"step\", \"epoch]"
		    }  ,       )
		__lowercase	:	Optional[int]     					=    dataclasses.field(
		    default=10  ,       metadata={"help": "Number of evaluation calls with no improvement after which training will be stopped."}  ,       )
		__lowercase	:	Optional[float]     					=    dataclasses.field(
		    default=0.0  ,       metadata={
		        "help": "How much the specified evaluation metric must improve to satisfy early stopping conditions."
		    }  ,       )
		__lowercase	:	Optional[bool]     					=    dataclasses.field(
		    default=_SCREAMING_SNAKE_CASE  ,       metadata={"help": "Whether to filter the pseudo-labeled data based on the confidence score."}  ,       )
		__lowercase	:	Optional[bool]     					=    dataclasses.field(
		    default=_SCREAMING_SNAKE_CASE  ,       metadata={"help": "Whether to filter the pseudo-labeled data based on the validation performance."}  ,       )
		__lowercase	:	Optional[bool]     					=    dataclasses.field(
		    default=_SCREAMING_SNAKE_CASE  ,       metadata={"help": "Whether to fine-tune on labeled data after pseudo training."}  ,       )
		__lowercase	:	Optional[float]     					=    dataclasses.field(
		    default=0.0  ,       metadata={"help": "Confidence threshold for pseudo-labeled data filtering."}  ,       )
		__lowercase	:	Optional[int]     					=    dataclasses.field(
		    default=100  ,       metadata={"help": "Number of evaluation calls with no improvement after which training will be stopped."}  ,       )
		__lowercase	:	Optional[int]     					=    dataclasses.field(
		    default=_SCREAMING_SNAKE_CASE  ,       metadata={"help": "Random seed for initialization."}  ,       )






def 				A	( lowercase   ,	lowercase   ,	lowercase   ,	lowercase   ,	lowercase   ,	lowercase			)       ->  Optional[int]:
		'''simple docstring'''

		UpperCamelCase          =					datasets.concatenate_datasets([infer_input, infer_output]   ,	axis=1			)

		if args.do_filter_by_confidence:
				UpperCamelCase          =					dataset.filter(lambda lowercase			: example["probability"] > args.confidence_threshold			)

		if args.do_filter_by_val_performance:
				assert eval_result >= 0.0 and eval_result <= 1.0
				UpperCamelCase          =					int(eval_result * len(lowercase			)			)
				print(lowercase			)
				UpperCamelCase          =					dataset.sort('probability'   ,	reverse=lowercase			)
				UpperCamelCase          =					dataset.select(range(lowercase			)			)

		UpperCamelCase          =					dataset.remove_columns(['label', 'probability']			)
		UpperCamelCase          =					dataset.rename_column('prediction'   ,	'label'			)
		UpperCamelCase          =					dataset.map(lambda lowercase			: {"label": idalabel[example["label"]]}			)
		UpperCamelCase          =					dataset.shuffle(seed=args.seed			)

		UpperCamelCase          =					os.path.join(lowercase   ,	f'''train_pseudo.{args.data_file_extension}'''			)
		if args.data_file_extension == "csv":
				dataset.to_csv(lowercase   ,	index=lowercase			)
		else:
				dataset.to_json(lowercase			)


def 				A	( lowercase   ,	lowercase   ,	lowercase   ,	lowercase   ,	**lowercase			)       ->  Any:
		'''simple docstring'''

		UpperCamelCase          =					Accelerator()
		# Make one log on every process with the configuration for debugging.
		logging.basicConfig(
		    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'   ,	datefmt='%m/%d/%Y %H:%M:%S'   ,	level=logging.INFO   ,	)
		logger.info(accelerator.state			)

		# Setup logging, we only want one process per machine to log things on the
		# screen. accelerator.is_local_main_process is only True for one process per
		# machine.
		logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR			)

		if accelerator.is_local_main_process:
				datasets.utils.logging.set_verbosity_warning()
				transformers.utils.logging.set_verbosity_info()
		else:
				datasets.utils.logging.set_verbosity_error()
				transformers.utils.logging.set_verbosity_error()

		UpperCamelCase          =					STModelArguments(model_name_or_path=lowercase			)
		UpperCamelCase          =					STDataArguments(train_file=lowercase   ,	infer_file=lowercase			)
		UpperCamelCase          =					STTrainingArguments(output_dir=lowercase			)
		UpperCamelCase          =					argparse.Namespace()

		for arg_class in (model_args, data_args, training_args):
				for key, value in vars(lowercase			).items():
						setattr(lowercase   ,	lowercase   ,	lowercase			)

		for key, value in kwargs.items():
				if hasattr(lowercase   ,	lowercase			):
						setattr(lowercase   ,	lowercase   ,	lowercase			)

    # Sanity checks
		UpperCamelCase          =					{}
		UpperCamelCase          =					None

		# You need to provide the training data and the data to predict on
		assert args.train_file is not None
		assert args.infer_file is not None
		UpperCamelCase          =					args.train_file
		UpperCamelCase          =					args.infer_file

		if args.evaluation_strategy != IntervalStrategy.NO.value:
				assert args.eval_file is not None
				UpperCamelCase          =					args.eval_file

		for key in data_files:
				UpperCamelCase          =					data_files[key].split('.'			)[-1]
				assert extension in ["csv", "json"], f'''`{key}_file` should be a csv or a json file.'''
				if args.data_file_extension is None:
						UpperCamelCase          =					extension
				else:
						assert extension == args.data_file_extension, f'''`{key}_file` should be a {args.data_file_extension} file`.'''

		assert (
		    args.eval_metric in datasets.list_metrics()
		), f'''{args.eval_metric} not in the list of supported metrics {datasets.list_metrics()}.'''

		# If passed along, set the training seed now.
		if args.seed is not None:
				set_seed(args.seed			)

		logger.info('Creating the initial data directory for self-training...'			)
		UpperCamelCase          =					f'''{args.output_dir}/self-train_iter-{{}}'''.format
		UpperCamelCase          =					data_dir_format(0			)

		if accelerator.is_main_process:
				if args.output_dir is not None:
						os.makedirs(args.output_dir   ,	exist_ok=lowercase			)
						os.makedirs(lowercase   ,	exist_ok=lowercase			)
		accelerator.wait_for_everyone()

		UpperCamelCase          =					None
		UpperCamelCase          =					None
		UpperCamelCase          =					0
		UpperCamelCase          =					False
		# Show the progress bar
		UpperCamelCase          =					tqdm(range(args.max_selftrain_iterations			)   ,	disable=not accelerator.is_local_main_process			)

		# Self-train
		for iteration in range(0   ,	int(args.max_selftrain_iterations			)			):
				UpperCamelCase          =					data_dir_format(lowercase			)
				assert os.path.exists(lowercase			)

				# Stage 1: initial fine-tuning for iteration = 0 or pseudo-training for
				# iteration > 0
				UpperCamelCase          =					os.path.join(lowercase   ,	'stage-1'			)
				UpperCamelCase          =					{
				    'accelerator': accelerator,
				    'model_name_or_path': args.model_name_or_path,
				    'cache_dir': args.cache_dir,
				    'do_train': True,
				    'train_file': data_files['train'] if iteration == 0 else data_files['train_pseudo'],
				    'do_eval': True if args.eval_file is not None else False,
				    'eval_file': data_files['eval'],
				    'do_predict': True,
				    'infer_file': data_files['infer'],
				    'task_name': args.task_name,
				    'label_list': args.label_list,
				    'output_dir': current_output_dir,
				    'eval_metric': args.eval_metric,
				    'evaluation_strategy': args.evaluation_strategy,
				    'early_stopping_patience': args.early_stopping_patience,
				    'early_stopping_threshold': args.early_stopping_threshold,
				    'seed': args.seed,
				}
				# Add additional training arguments
				for key, value in kwargs.items():
						if key not in arguments_dict and not hasattr(lowercase   ,	lowercase			):
								arguments_dict.update({key: value}			)

				UpperCamelCase          =					os.path.join(lowercase   ,	'best-checkpoint'   ,	lowercase			)
				if os.path.exists(lowercase			):
						logger.info(
						    'Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 1.'   ,	lowercase   ,	lowercase   ,	)
				else:
						logger.info('***** Running self-training: iteration: %d, stage: 1 *****'   ,	lowercase			)
						finetune(**lowercase			)
						accelerator.wait_for_everyone()
						assert os.path.exists(lowercase			)
						logger.info('Self-training job completed: iteration: %d, stage: 1.'   ,	lowercase			)

				if iteration > 0 and args.finetune_on_labeled_data:
						# Stage 2 (optional): fine-tuning on the original labeled data
						UpperCamelCase          =					os.path.join(lowercase   ,	'best-checkpoint'			)
						UpperCamelCase          =					os.path.join(lowercase   ,	'stage-2'			)
						# Update arguments_dict
						UpperCamelCase          =					model_path
						UpperCamelCase          =					data_files['train']
						UpperCamelCase          =					current_output_dir

						UpperCamelCase          =					os.path.join(lowercase   ,	'best-checkpoint'   ,	lowercase			)
						if os.path.exists(lowercase			):
								logger.info(
								    'Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 2.'   ,	lowercase   ,	lowercase   ,	)
						else:
								logger.info('***** Running self-training: iteration: %d, stage: 2 *****'   ,	lowercase			)
								finetune(**lowercase			)
								accelerator.wait_for_everyone()
								assert os.path.exists(lowercase			)
								logger.info('Self-training job completed: iteration: %d, stage: 2.'   ,	lowercase			)

				UpperCamelCase          =					iteration
				UpperCamelCase          =					data_dir_format(iteration + 1			)

				UpperCamelCase          =					AutoConfig.from_pretrained(os.path.join(lowercase   ,	'best-checkpoint'			)			)
				UpperCamelCase          =					config.idalabel
				UpperCamelCase          =					os.path.join(lowercase   ,	'eval_results_best-checkpoint.json'			)
				UpperCamelCase          =					os.path.join(lowercase   ,	'test_results_best-checkpoint.json'			)
				assert os.path.exists(lowercase			)

				with open(lowercase   ,	'r'			) as f:
						UpperCamelCase          =					float(json.load(lowercase			)[args.eval_metric]			)
				UpperCamelCase          =					os.path.join(lowercase   ,	'infer_output_best-checkpoint.csv'			)
				assert os.path.exists(lowercase			)
				# Loading the dataset from local csv or json files.
				UpperCamelCase          =					load_dataset(args.data_file_extension   ,	data_files={'data': data_files['infer']}			)['data']
				UpperCamelCase          =					load_dataset('csv'   ,	data_files={'data': infer_output_file}			)['data']

				if accelerator.is_main_process:
						os.makedirs(lowercase   ,	exist_ok=lowercase			)
						shutil.copy(lowercase   ,	os.path.join(lowercase   ,	f'''eval_results_iter-{iteration}.json'''			)			)
						if os.path.exists(lowercase			):
								shutil.copy(lowercase   ,	os.path.join(lowercase   ,	f'''test_results_iter-{iteration}.json'''			)			)
						create_pseudo_labeled_data(lowercase   ,	lowercase   ,	lowercase   ,	lowercase   ,	lowercase   ,	lowercase			)
				accelerator.wait_for_everyone()

				UpperCamelCase          =					os.path.join(lowercase   ,	f'''train_pseudo.{args.data_file_extension}'''			)

				if args.evaluation_strategy != IntervalStrategy.NO.value:
						UpperCamelCase          =					eval_result

						if best_iteration is None:
								UpperCamelCase          =					new_iteration
								UpperCamelCase          =					new_eval_result
						else:
								if new_eval_result - best_eval_result > args.early_stopping_threshold:
										UpperCamelCase          =					new_iteration
										UpperCamelCase          =					new_eval_result
										UpperCamelCase          =					0
								else:
										if new_eval_result == best_eval_result:
												UpperCamelCase          =					new_iteration
												UpperCamelCase          =					new_eval_result
										early_stopping_patience_counter += 1

								if early_stopping_patience_counter >= args.early_stopping_patience:
										UpperCamelCase          =					True

				progress_bar.update(1			)

				if should_training_stop:
						break

		if best_iteration is not None:
				# Save the best iteration
				logger.info('Best iteration: %d'   ,	lowercase			)
				logger.info('Best evaluation result: %s = %f'   ,	args.eval_metric   ,	lowercase			)
				accelerator.wait_for_everyone()
				if accelerator.is_main_process:
						shutil.copy(
						    os.path.join(lowercase   ,	f'''eval_results_iter-{iteration}.json'''			)   ,	os.path.join(lowercase   ,	'eval_results_best-iteration.json'			)   ,	)
		else:
				# Assume that the last iteration is the best
				logger.info('Best iteration: %d'   ,	args.max_selftrain_iterations - 1			)
				logger.info('Best evaluation result: %s = %f'   ,	args.eval_metric   ,	lowercase			)
				accelerator.wait_for_everyone()
				if accelerator.is_main_process:
						shutil.copy(
						    os.path.join(lowercase   ,	f'''eval_results_iter-{args.max_selftrain_iterations - 1}.json'''			)   ,	os.path.join(lowercase   ,	'eval_results_best-iteration.json'			)   ,	)


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	( )       ->  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			)       ->  Any:
		'''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   ,	lowercase   ,	lowercase			)       ->  Dict:
		'''simple docstring'''

		UpperCamelCase          =					dct.pop(lowercase			)
		UpperCamelCase          =					val


def 				A	( lowercase   ,	lowercase			)       ->  List[str]:
		'''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   ,	lowercase			)       ->  int:
		'''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   ,	lowercase=None   ,	lowercase=False			)       ->  Union[str, Any]:
		'''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__":
					_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 :			List[str]              =	parser.parse_args()

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