




"""simple docstring"""

from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
		from PIL import Image

		from ..image_utils import load_image

if is_tf_available():
		import tensorflow as tf

		from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
		from ..tf_utils import stable_softmax

if is_torch_available():
		from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING

_lowerCamelCase  :			Optional[Any]        =       logging.get_logger(__name__)

@add_end_docstrings(__UpperCAmelCase)
class        lowercase      (    __UpperCAmelCase):







				def __init__(							self   :					Optional[Any]      ,     *_lowerCamelCase   :					List[Any]      ,     **_lowerCamelCase   :					Union[str, Any]       ):



					"""simple docstring"""

					super().__init__(*_lowerCamelCase      ,     **_lowerCamelCase       )
					requires_backends(self      ,     '''vision'''       )
					self.check_model_type(
					    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
					    if self.framework == '''tf'''
					    else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING       )







				def 					a_  (							self   :					List[str]      ,     _lowerCamelCase   :					Dict=None       ):



					"""simple docstring"""

					A_  :				List[Any]	   =		{}
					if top_k is not None:
						A_  :				int	   =		top_k
					return {}, {}, postprocess_params







				def __call__(							self   :					Union[str, Any]      ,     _lowerCamelCase   :					Union[str, List[str], "Image.Image", List["Image.Image"]]      ,     **_lowerCamelCase   :					Dict       ):



					"""simple docstring"""

					return super().__call__(_lowerCamelCase      ,     **_lowerCamelCase       )







				def 					a_  (							self   :					Any      ,     _lowerCamelCase   :					Tuple       ):



					"""simple docstring"""

					A_  :				Dict	   =		load_image(_lowerCamelCase       )
					A_  :				Union[str, Any]	   =		self.image_processor(images=_lowerCamelCase      ,     return_tensors=self.framework       )
					return model_inputs







				def 					a_  (							self   :					Optional[int]      ,     _lowerCamelCase   :					Union[str, Any]       ):



					"""simple docstring"""

					A_  :				int	   =		self.model(**_lowerCamelCase       )
					return model_outputs







				def 					a_  (							self   :					int      ,     _lowerCamelCase   :					Optional[int]      ,     _lowerCamelCase   :					Optional[int]=5       ):



					"""simple docstring"""

					if top_k > self.model.config.num_labels:
						A_  :				Optional[Any]	   =		self.model.config.num_labels

					if self.framework == "pt":
						A_  :				List[str]	   =		model_outputs.logits.softmax(-1       )[0]
						A_   , A_  :				Tuple	   =		probs.topk(_lowerCamelCase       )
					elif self.framework == "tf":
						A_  :				List[str]	   =		stable_softmax(model_outputs.logits      ,     axis=-1       )[0]
						A_  :				Optional[int]	   =		tf.math.top_k(_lowerCamelCase      ,     k=_lowerCamelCase       )
						A_   , A_  :				Dict	   =		topk.values.numpy(), topk.indices.numpy()
					else:
						raise ValueError(F"""Unsupported framework: {self.framework}"""       )

					A_  :				int	   =		scores.tolist()
					A_  :				Dict	   =		ids.tolist()
					return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(_lowerCamelCase      ,     _lowerCamelCase       )]











"""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  :			Optional[Any]        =       logging.get_logger(__name__)

_lowerCamelCase  :			Any        =       [
    ('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  :			int        =       [
    '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_      (		_UpperCAmelCase	):




	"""simple docstring"""






	A_  :				Union[str, Any]	   =		torch.load(_UpperCAmelCase  ,      map_location='''cpu'''	)
	return sd

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase  ,      _UpperCAmelCase=rename_keys_prefix	):




	"""simple docstring"""






	A_  :				Any	   =		OrderedDict()
	A_  :				Union[str, Any]	   =		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
		A_  :				List[Any]	   =		key
		for name_pair in rename_keys_prefix:
			A_  :				Any	   =		new_key.replace(name_pair[0]  ,      name_pair[1]	)
		A_  :				Dict	   =		d[key]
		if key == "bert.cls.predictions.decoder.weight":
			# Old bert code didn't have `decoder.bias`, but was added separately
			A_  :				Any	   =		new_d['''cls.predictions.bias''']
	return new_d

@torch.no_grad()
def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase	):




	"""simple docstring"""






	assert (
	    checkpoint_path.split('''/'''	)[-1] in ACCEPTABLE_CHECKPOINTS
	), f"""The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}."""

	# Get Config
	if "pre" in checkpoint_path:
		A_  :				List[Any]	   =		'''pretraining'''
		if "vcr" in checkpoint_path:
			A_  :				Optional[Any]	   =		{'''visual_embedding_dim''': 512}
		elif "vqa_advanced" in checkpoint_path:
			A_  :				Tuple	   =		{'''visual_embedding_dim''': 2048}
		elif "vqa" in checkpoint_path:
			A_  :				str	   =		{'''visual_embedding_dim''': 2048}
		elif "nlvr" in checkpoint_path:
			A_  :				List[Any]	   =		{'''visual_embedding_dim''': 1024}
		else:
			raise NotImplementedError(f"""No implementation found for `{checkpoint_path}`."""	)
	else:
		if "vcr" in checkpoint_path:
			A_  :				Any	   =		{'''visual_embedding_dim''': 512}
			A_  :				int	   =		'''multichoice'''
		elif "vqa_advanced" in checkpoint_path:
			A_  :				Any	   =		{'''visual_embedding_dim''': 2048}
			A_  :				Any	   =		'''vqa_advanced'''
		elif "vqa" in checkpoint_path:
			A_  :				int	   =		{'''visual_embedding_dim''': 2048, '''num_labels''': 3129}
			A_  :				Any	   =		'''vqa'''
		elif "nlvr" in checkpoint_path:
			A_  :				List[Any]	   =		{
			    '''visual_embedding_dim''': 1024,
			    '''num_labels''': 2,
			}
			A_  :				List[Any]	   =		'''nlvr'''

	A_  :				Dict	   =		VisualBertConfig(**_UpperCAmelCase	)

	# Load State Dict
	A_  :				List[str]	   =		load_state_dict(_UpperCAmelCase	)

	A_  :				List[str]	   =		get_new_dict(_UpperCAmelCase  ,      _UpperCAmelCase	)

	if model_type == "pretraining":
		A_  :				List[Any]	   =		VisualBertForPreTraining(_UpperCAmelCase	)
	elif model_type == "vqa":
		A_  :				Optional[Any]	   =		VisualBertForQuestionAnswering(_UpperCAmelCase	)
	elif model_type == "nlvr":
		A_  :				Any	   =		VisualBertForVisualReasoning(_UpperCAmelCase	)
	elif model_type == "multichoice":
		A_  :				Optional[int]	   =		VisualBertForMultipleChoice(_UpperCAmelCase	)

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


if __name__ == "__main__":
		_lowerCamelCase  :			Dict        =       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  :			Optional[int]        =       parser.parse_args()
		convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)





