


"""simple docstring"""






import argparse
from collections import OrderedDict
from pathlib import Path

import torch

from transformers import (
    VisualBertConfig,
    VisualBertForMultipleChoice,
    VisualBertForPreTraining,
    VisualBertForQuestionAnswering,
    VisualBertForVisualReasoning,
)
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCamelCase     =  logging.get_logger(__name__)

_lowerCamelCase     =  [
    ('bert.bert', 'visual_bert'),
    ('bert.cls', 'cls'),
    ('bert.classifier', 'cls'),
    ('token_type_embeddings_visual', 'visual_token_type_embeddings'),
    ('position_embeddings_visual', 'visual_position_embeddings'),
    ('projection', 'visual_projection'),
]

_lowerCamelCase     =  [
    'nlvr2_coco_pre_trained.th',
    'nlvr2_fine_tuned.th',
    'nlvr2_pre_trained.th',
    'vcr_coco_pre_train.th',
    'vcr_fine_tune.th',
    'vcr_pre_train.th',
    'vqa_coco_pre_trained.th',
    'vqa_fine_tuned.th',
    'vqa_pre_trained.th',
]


def __lowercase     (      lowerCamelCase_  :      Union[str, Any]  ):
				SCREAMING_SNAKE_CASE__     							=				torch.load(snake_case__						,					map_location="cpu"  )
				return sd


def __lowercase     (      lowerCamelCase_  :      List[str]						,					lowerCamelCase_  :      Optional[Any]						,					lowerCamelCase_  :      int=rename_keys_prefix  ):
				SCREAMING_SNAKE_CASE__     							=				OrderedDict()
				SCREAMING_SNAKE_CASE__     							=				torch.arange(config.max_position_embeddings  ).expand((1, -1)  )
				# detector_d = OrderedDict()
				for key in d:
								if "detector" in key:
												# detector_d[key.replace('detector.','')] = d[key]
												continue
								SCREAMING_SNAKE_CASE__     							=				key
								for name_pair in rename_keys_prefix:
												SCREAMING_SNAKE_CASE__     							=				new_key.replace(name_pair[0]						,					name_pair[1]  )
								SCREAMING_SNAKE_CASE__     							=				d[key]
								if key == "bert.cls.predictions.decoder.weight":
												# Old bert code didn't have `decoder.bias`, but was added separately
												SCREAMING_SNAKE_CASE__     							=				new_d["cls.predictions.bias"]
				return new_d


@torch.no_grad()
def __lowercase     (      lowerCamelCase_  :      Tuple						,					lowerCamelCase_  :      Tuple  ):
				assert (
				    checkpoint_path.split("/"  )[-1] in ACCEPTABLE_CHECKPOINTS
				), F'''The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}.'''

				# Get Config
				if "pre" in checkpoint_path:
								SCREAMING_SNAKE_CASE__     							=				"pretraining"
								if "vcr" in checkpoint_path:
												SCREAMING_SNAKE_CASE__     							=				{"visual_embedding_dim": 512}
								elif "vqa_advanced" in checkpoint_path:
												SCREAMING_SNAKE_CASE__     							=				{"visual_embedding_dim": 2048}
								elif "vqa" in checkpoint_path:
												SCREAMING_SNAKE_CASE__     							=				{"visual_embedding_dim": 2048}
								elif "nlvr" in checkpoint_path:
												SCREAMING_SNAKE_CASE__     							=				{"visual_embedding_dim": 1024}
								else:
												raise NotImplementedError(F'''No implementation found for `{checkpoint_path}`.'''  )
				else:
								if "vcr" in checkpoint_path:
												SCREAMING_SNAKE_CASE__     							=				{"visual_embedding_dim": 512}
												SCREAMING_SNAKE_CASE__     							=				"multichoice"
								elif "vqa_advanced" in checkpoint_path:
												SCREAMING_SNAKE_CASE__     							=				{"visual_embedding_dim": 2048}
												SCREAMING_SNAKE_CASE__     							=				"vqa_advanced"
								elif "vqa" in checkpoint_path:
												SCREAMING_SNAKE_CASE__     							=				{"visual_embedding_dim": 2048, "num_labels": 3129}
												SCREAMING_SNAKE_CASE__     							=				"vqa"
								elif "nlvr" in checkpoint_path:
												SCREAMING_SNAKE_CASE__     							=				{
												    "visual_embedding_dim": 1024,
												    "num_labels": 2,
												}
												SCREAMING_SNAKE_CASE__     							=				"nlvr"

				SCREAMING_SNAKE_CASE__     							=				VisualBertConfig(**snake_case__  )

				# Load State Dict
				SCREAMING_SNAKE_CASE__     							=				load_state_dict(snake_case__  )

				SCREAMING_SNAKE_CASE__     							=				get_new_dict(snake_case__						,					snake_case__  )

				if model_type == "pretraining":
								SCREAMING_SNAKE_CASE__     							=				VisualBertForPreTraining(snake_case__  )
				elif model_type == "vqa":
								SCREAMING_SNAKE_CASE__     							=				VisualBertForQuestionAnswering(snake_case__  )
				elif model_type == "nlvr":
								SCREAMING_SNAKE_CASE__     							=				VisualBertForVisualReasoning(snake_case__  )
				elif model_type == "multichoice":
								SCREAMING_SNAKE_CASE__     							=				VisualBertForMultipleChoice(snake_case__  )

				model.load_state_dict(snake_case__  )
				# Save Checkpoints
				Path(snake_case__  ).mkdir(exist_ok=snake_case__  )
				model.save_pretrained(snake_case__  )


if __name__ == "__main__":
		_lowerCamelCase     =  argparse.ArgumentParser()
		# Required parameters
		parser.add_argument('orig_checkpoint_path', type=str, help='A path to .th on local filesystem.')
		parser.add_argument('pytorch_dump_folder_path', type=str, help='Path to the output PyTorch model.')
		_lowerCamelCase     =  parser.parse_args()
		convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)




"""simple docstring"""






import json
import os
import tempfile

import transformers

import datasets
from utils import generate_example_dataset, get_duration


_lowerCamelCase     =  500000

_lowerCamelCase     ,     _lowerCamelCase     =  os.path.split(__file__)
_lowerCamelCase     =  os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))


@get_duration
def __lowercase     (      lowerCamelCase_  :      datasets.Dataset						,					**lowerCamelCase_  :      Optional[int]  ):
				SCREAMING_SNAKE_CASE__     							=				dataset.map(**lowerCamelCase_  )


@get_duration
def __lowercase     (      lowerCamelCase_  :      datasets.Dataset						,					**lowerCamelCase_  :      Optional[int]  ):
				SCREAMING_SNAKE_CASE__     							=				dataset.filter(**lowerCamelCase_  )


def __lowercase     (      ):
				SCREAMING_SNAKE_CASE__     							=				{"num examples": SPEED_TEST_N_EXAMPLES}
				with tempfile.TemporaryDirectory() as tmp_dir:
								SCREAMING_SNAKE_CASE__     							=				datasets.Features({"text": datasets.Value("string"  ), "numbers": datasets.Value("float32"  )}  )
								SCREAMING_SNAKE_CASE__     							=				generate_example_dataset(
								    os.path.join(lowerCamelCase_						,					"dataset.arrow"  )						,					lowerCamelCase_						,					num_examples=lowerCamelCase_  )

								SCREAMING_SNAKE_CASE__     							=				transformers.AutoTokenizer.from_pretrained("bert-base-cased"						,					use_fast=lowerCamelCase_  )

								def tokenize(lowerCamelCase_  :      Union[str, Any]  ):
												return tokenizer(examples["text"]  )

								SCREAMING_SNAKE_CASE__     							=				map(lowerCamelCase_  )

								SCREAMING_SNAKE_CASE__     							=				map(lowerCamelCase_						,					batched=lowerCamelCase_  )

								SCREAMING_SNAKE_CASE__     							=				map(lowerCamelCase_						,					function=lambda lowerCamelCase_  : None						,					batched=lowerCamelCase_  )

								with dataset.formatted_as(type="numpy"  ):
												SCREAMING_SNAKE_CASE__     							=				map(lowerCamelCase_						,					function=lambda lowerCamelCase_  : None						,					batched=lowerCamelCase_  )

								with dataset.formatted_as(type="pandas"  ):
												SCREAMING_SNAKE_CASE__     							=				map(lowerCamelCase_						,					function=lambda lowerCamelCase_  : None						,					batched=lowerCamelCase_  )

								with dataset.formatted_as(type="torch"						,					columns="numbers"  ):
												SCREAMING_SNAKE_CASE__     							=				map(lowerCamelCase_						,					function=lambda lowerCamelCase_  : None						,					batched=lowerCamelCase_  )

								with dataset.formatted_as(type="tensorflow"						,					columns="numbers"  ):
												SCREAMING_SNAKE_CASE__     							=				map(lowerCamelCase_						,					function=lambda lowerCamelCase_  : None						,					batched=lowerCamelCase_  )

								SCREAMING_SNAKE_CASE__     							=				map(lowerCamelCase_						,					function=lowerCamelCase_						,					batched=lowerCamelCase_  )

								SCREAMING_SNAKE_CASE__     							=				filter(lowerCamelCase_  )

								# Activate later when tokenizer support batched inputs
								# with dataset.formatted_as(type='numpy'):
								#     times[func.__name__ + " fast-tokenizer batched numpy"] = func(dataset, function=tokenize, batched=True)

				with open(lowerCamelCase_						,					"wb"  ) as f:
								f.write(json.dumps(lowerCamelCase_  ).encode("utf-8"  )  )


if __name__ == "__main__":  # useful to run the profiler
	benchmark_map_filter()
