



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	       =	logging.getLogger(__name__)

__UpperCAmelCase	       =	'''pytorch_model.bin'''



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






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






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






def 		UpperCamelCase     (       snake_case__	:   Optional[int]	,					snake_case__	:   Dict	,					snake_case__	:   Union[str, Any]	,					snake_case__	:   int	,					snake_case__	:   Tuple	,					snake_case__	:   Dict     )     ->					Optional[Any]:
   UpperCamelCase   :  int						     =     datasets.concatenate_datasets([infer_input, infer_output]	,					axis=1     )

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

   if args.do_filter_by_val_performance:
      assert eval_result >= 0.0 and eval_result <= 1.0
      UpperCamelCase   :  List[str]						     =     int(eval_result * len(snake_case__     )     )
      print(snake_case__     )
      UpperCamelCase   :  Union[str, Any]						     =     dataset.sort('probability'	,					reverse=snake_case__     )
      UpperCamelCase   :  int						     =     dataset.select(range(snake_case__     )     )

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

   UpperCamelCase   :  Dict						     =     os.path.join(snake_case__	,					F"""train_pseudo.{args.data_file_extension}"""     )
   if args.data_file_extension == "csv":
      dataset.to_csv(snake_case__	,					index=snake_case__     )
   else:
      dataset.to_json(snake_case__     )


def 		UpperCamelCase     (       snake_case__	:   Optional[Any]	,					snake_case__	:   List[Any]	,					snake_case__	:   List[str]	,					snake_case__	:   Tuple	,					**snake_case__	:   Union[str, Any]     )     ->					Any:
   UpperCamelCase   :  List[str]						     =     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   :  List[str]						     =     STModelArguments(model_name_or_path=snake_case__     )
   UpperCamelCase   :  int						     =     STDataArguments(train_file=snake_case__	,					infer_file=snake_case__     )
   UpperCamelCase   :  Dict						     =     STTrainingArguments(output_dir=snake_case__     )
   UpperCamelCase   :  Optional[Any]						     =     argparse.Namespace()

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

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

    # Sanity checks
   UpperCamelCase   :  int						     =     {}
   UpperCamelCase   :  Optional[int]						     =     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   :  int						     =     args.train_file
   UpperCamelCase   :  Optional[Any]						     =     args.infer_file

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

   for key in data_files:
      UpperCamelCase   :  List[str]						     =     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   :  List[Any]						     =     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   :  Dict						     =     F"""{args.output_dir}/self-train_iter-{{}}""".format
   UpperCamelCase   :  Any						     =     data_dir_format(0     )

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

   UpperCamelCase   :  List[Any]						     =     None
   UpperCamelCase   :  Any						     =     None
   UpperCamelCase   :  Tuple						     =     0
   UpperCamelCase   :  str						     =     False
   # Show the progress bar
   UpperCamelCase   :  List[str]						     =     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   :  List[str]						     =     data_dir_format(snake_case__     )
      assert os.path.exists(snake_case__     )

      # Stage 1: initial fine-tuning for iteration = 0 or pseudo-training for
      # iteration > 0
      UpperCamelCase   :  List[Any]						     =     os.path.join(snake_case__	,					'stage-1'     )
      UpperCamelCase   :  Optional[int]						     =     {
          '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(snake_case__	,					snake_case__     ):
            arguments_dict.update({key: value}     )

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

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

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

      UpperCamelCase   :  int						     =     iteration
      UpperCamelCase   :  Any						     =     data_dir_format(iteration + 1     )

      UpperCamelCase   :  Optional[int]						     =     AutoConfig.from_pretrained(os.path.join(snake_case__	,					'best-checkpoint'     )     )
      UpperCamelCase   :  Union[str, Any]						     =     config.idalabel
      UpperCamelCase   :  str						     =     os.path.join(snake_case__	,					'eval_results_best-checkpoint.json'     )
      UpperCamelCase   :  List[Any]						     =     os.path.join(snake_case__	,					'test_results_best-checkpoint.json'     )
      assert os.path.exists(snake_case__     )

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

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

      UpperCamelCase   :  Tuple						     =     os.path.join(snake_case__	,					F"""train_pseudo.{args.data_file_extension}"""     )

      if args.evaluation_strategy != IntervalStrategy.NO.value:
         UpperCamelCase   :  Optional[Any]						     =     eval_result

         if best_iteration is None:
            UpperCamelCase   :  Dict						     =     new_iteration
            UpperCamelCase   :  List[str]						     =     new_eval_result
         else:
            if new_eval_result - best_eval_result > args.early_stopping_threshold:
               UpperCamelCase   :  Tuple						     =     new_iteration
               UpperCamelCase   :  List[Any]						     =     new_eval_result
               UpperCamelCase   :  str						     =     0
            else:
               if new_eval_result == best_eval_result:
                  UpperCamelCase   :  Any						     =     new_iteration
                  UpperCamelCase   :  Tuple						     =     new_eval_result
               early_stopping_patience_counter += 1

            if early_stopping_patience_counter >= args.early_stopping_patience:
               UpperCamelCase   :  Tuple						     =     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'	,					snake_case__     )
      logger.info('Best evaluation result: %s = %f'	,					args.eval_metric	,					snake_case__     )
      accelerator.wait_for_everyone()
      if accelerator.is_main_process:
         shutil.copy(
             os.path.join(snake_case__	,					F"""eval_results_iter-{iteration}.json"""     )	,					os.path.join(snake_case__	,					'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	,					snake_case__     )
      accelerator.wait_for_everyone()
      if accelerator.is_main_process:
         shutil.copy(
             os.path.join(snake_case__	,					F"""eval_results_iter-{args.max_selftrain_iterations - 1}.json"""     )	,					os.path.join(snake_case__	,					'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 		UpperCamelCase     (       )     ->					Tuple:
   UpperCamelCase   :  Any						     =     'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
   UpperCamelCase   :  List[str]						     =     Image.open(requests.get(snake_case__	,					stream=snake_case__     ).raw     ).convert('RGB'     )

   return image
def 		UpperCamelCase     (       snake_case__	:   int     )     ->					List[Any]:
   UpperCamelCase   :  Optional[int]						     =     []
   # 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 		UpperCamelCase     (       snake_case__	:   Optional[Any]	,					snake_case__	:   Dict	,					snake_case__	:   Optional[int]     )     ->					Optional[int]:
   UpperCamelCase   :  Dict						     =     dct.pop(snake_case__     )
   UpperCamelCase   :  str						     =     val
def 		UpperCamelCase     (       snake_case__	:   str	,					snake_case__	:   Union[str, Any]     )     ->					Optional[int]:
   for i in range(config.vision_config.num_hidden_layers     ):
      # read in original q and v biases
      UpperCamelCase   :  Optional[Any]						     =     state_dict.pop(F"""visual_encoder.blocks.{i}.attn.q_bias"""     )
      UpperCamelCase   :  int						     =     state_dict.pop(F"""visual_encoder.blocks.{i}.attn.v_bias"""     )

      # next, set bias in the state dict
      UpperCamelCase   :  int						     =     torch.cat((q_bias, torch.zeros_like(snake_case__	,					requires_grad=snake_case__     ), v_bias)     )
      UpperCamelCase   :  Tuple						     =     qkv_bias
def 		UpperCamelCase     (       snake_case__	:   Tuple	,					snake_case__	:   Optional[Any]     )     ->					Dict:
   UpperCamelCase   :  str						     =     364 if 'coco' in model_name else 224
   UpperCamelCase   :  Union[str, Any]						     =     BlipaVisionConfig(image_size=snake_case__     ).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   :  List[Any]						     =     OPTConfig.from_pretrained('facebook/opt-2.7b'	,					eos_token_id=snake_case__     ).to_dict()
   elif "opt-6.7b" in model_name:
      UpperCamelCase   :  int						     =     OPTConfig.from_pretrained('facebook/opt-6.7b'	,					eos_token_id=snake_case__     ).to_dict()
   elif "t5-xl" in model_name:
      UpperCamelCase   :  List[str]						     =     TaConfig.from_pretrained('google/flan-t5-xl'	,					dense_act_fn='gelu'	,					bos_token_id=1     ).to_dict()
   elif "t5-xxl" in model_name:
      UpperCamelCase   :  int						     =     TaConfig.from_pretrained('google/flan-t5-xxl'	,					dense_act_fn='gelu'	,					bos_token_id=1     ).to_dict()

   UpperCamelCase   :  Any						     =     BlipaConfig(vision_config=snake_case__	,					text_config=snake_case__     )

   return config, image_size
@torch.no_grad()
def 		UpperCamelCase     (       snake_case__	:   int	,					snake_case__	:   Dict=None	,					snake_case__	:   int=False     )     ->					List[Any]:
   UpperCamelCase   :  str						     =     (
       AutoTokenizer.from_pretrained('facebook/opt-2.7b'     )
       if 'opt' in model_name
       else AutoTokenizer.from_pretrained('google/flan-t5-xl'     )
   )
   UpperCamelCase   :  int						     =     tokenizer('\n'	,					add_special_tokens=snake_case__     ).input_ids[0]
   UpperCamelCase     ,					UpperCamelCase   :  Union[str, Any]						     =     get_blipa_config(snake_case__	,					eos_token_id=snake_case__     )

   UpperCamelCase   :  Dict						     =     BlipaForConditionalGeneration(snake_case__     ).eval()

   UpperCamelCase   :  Optional[Any]						     =     {
       '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   :  Optional[Any]						     =     model_name_to_original[model_name]

   # load original model
   print('Loading original model...'     )
   UpperCamelCase   :  List[str]						     =     'cuda' if torch.cuda.is_available() else 'cpu'
   UpperCamelCase     ,					UpperCamelCase     ,					UpperCamelCase   :  Tuple						     =     load_model_and_preprocess(
       name=snake_case__	,					model_type=snake_case__	,					is_eval=snake_case__	,					device=snake_case__     )
   original_model.eval()
   print('Done!'     )

   # update state dict keys
   UpperCamelCase   :  List[Any]						     =     original_model.state_dict()
   UpperCamelCase   :  Tuple						     =     create_rename_keys(snake_case__     )
   for src, dest in rename_keys:
      rename_key(snake_case__	,					snake_case__	,					snake_case__     )

   # some keys can be renamed efficiently
   for key, val in state_dict.copy().items():
      UpperCamelCase   :  Optional[Any]						     =     state_dict.pop(snake_case__     )
      if key.startswith('Qformer.bert'     ):
         UpperCamelCase   :  List[str]						     =     key.replace('Qformer.bert'	,					'qformer'     )
      if "attention.self" in key:
         UpperCamelCase   :  Tuple						     =     key.replace('self'	,					'attention'     )
      if "opt_proj" in key:
         UpperCamelCase   :  Union[str, Any]						     =     key.replace('opt_proj'	,					'language_projection'     )
      if "t5_proj" in key:
         UpperCamelCase   :  Optional[Any]						     =     key.replace('t5_proj'	,					'language_projection'     )
      if key.startswith('opt'     ):
         UpperCamelCase   :  Dict						     =     key.replace('opt'	,					'language'     )
      if key.startswith('t5'     ):
         UpperCamelCase   :  Dict						     =     key.replace('t5'	,					'language'     )
      UpperCamelCase   :  Optional[int]						     =     val

   # read in qv biases
   read_in_q_v_bias(snake_case__	,					snake_case__     )

   UpperCamelCase     ,					UpperCamelCase   :  Any						     =     hf_model.load_state_dict(snake_case__	,					strict=snake_case__     )
   assert len(snake_case__     ) == 0
   assert unexpected_keys == ["qformer.embeddings.position_ids"]

   UpperCamelCase   :  List[str]						     =     load_demo_image()
   UpperCamelCase   :  str						     =     vis_processors['eval'](snake_case__     ).unsqueeze(0     ).to(snake_case__     )
   UpperCamelCase   :  Any						     =     tokenizer(['\n']	,					return_tensors='pt'     ).input_ids.to(snake_case__     )

   # create processor
   UpperCamelCase   :  Optional[Any]						     =     BlipImageProcessor(
       size={'height': image_size, 'width': image_size}	,					image_mean=snake_case__	,					image_std=snake_case__     )
   UpperCamelCase   :  Any						     =     BlipaProcessor(image_processor=snake_case__	,					tokenizer=snake_case__     )
   UpperCamelCase   :  Optional[int]						     =     processor(images=snake_case__	,					return_tensors='pt'     ).pixel_values.to(snake_case__     )

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

   original_model.to(snake_case__     )
   hf_model.to(snake_case__     )
   with torch.no_grad():
      if "opt" in model_name:
         UpperCamelCase   :  Tuple						     =     original_model({'image': original_pixel_values, 'text_input': ['']}     ).logits
         UpperCamelCase   :  str						     =     hf_model(snake_case__	,					snake_case__     ).logits
      else:
         UpperCamelCase   :  Tuple						     =     original_model(
             {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}     ).logits
         UpperCamelCase   :  List[Any]						     =     input_ids.masked_fill(input_ids == tokenizer.pad_token_id	,					-100     )
         UpperCamelCase   :  Optional[int]						     =     hf_model(snake_case__	,					snake_case__	,					labels=snake_case__     ).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   :  List[str]						     =     torch.tensor(
          [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]	,					device=snake_case__     )
      assert torch.allclose(logits[0, :3, :3]	,					snake_case__	,					atol=1E-4     )
   elif model_name == "blip2-flan-t5-xl-coco":
      UpperCamelCase   :  Union[str, Any]						     =     torch.tensor(
          [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]	,					device=snake_case__     )
   else:
      # cast to same type
      UpperCamelCase   :  Optional[int]						     =     logits.dtype
      assert torch.allclose(original_logits.to(snake_case__     )	,					snake_case__	,					atol=1E-2     )
   print('Looks ok!'     )

   print('Generating a caption...'     )
   UpperCamelCase   :  Optional[int]						     =     ''
   UpperCamelCase   :  Union[str, Any]						     =     tokenizer(snake_case__	,					return_tensors='pt'     ).input_ids.to(snake_case__     )

   UpperCamelCase   :  str						     =     original_model.generate({'image': original_pixel_values}     )
   UpperCamelCase   :  str						     =     hf_model.generate(
       snake_case__	,					snake_case__	,					do_sample=snake_case__	,					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:'	,					snake_case__     )
   UpperCamelCase   :  Optional[int]						     =     input_ids.shape[1]
   UpperCamelCase   :  Union[str, Any]						     =     processor.batch_decode(outputs[:, prompt_length:]	,					skip_special_tokens=snake_case__     )
   UpperCamelCase   :  Dict						     =     [text.strip() for text in output_text]
   print('HF generation:'	,					snake_case__     )

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

   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	       =	argparse.ArgumentParser()
    __UpperCAmelCase	       =	[
        '''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	       =	parser.parse_args()

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



