





"""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       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	torch.load(snake_case__  ,	map_location="cpu"    )
		return sd


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE=rename_keys_prefix    ):




		"""simple docstring"""


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


@torch.no_grad()
def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""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:
				UpperCamelCase   =	"pretraining"
				if "vcr" in checkpoint_path:
						UpperCamelCase   =	{"visual_embedding_dim": 512}
				elif "vqa_advanced" in checkpoint_path:
						UpperCamelCase   =	{"visual_embedding_dim": 2_048}
				elif "vqa" in checkpoint_path:
						UpperCamelCase   =	{"visual_embedding_dim": 2_048}
				elif "nlvr" in checkpoint_path:
						UpperCamelCase   =	{"visual_embedding_dim": 1_024}
				else:
						raise NotImplementedError(F"No implementation found for `{checkpoint_path}`."    )
		else:
				if "vcr" in checkpoint_path:
						UpperCamelCase   =	{"visual_embedding_dim": 512}
						UpperCamelCase   =	"multichoice"
				elif "vqa_advanced" in checkpoint_path:
						UpperCamelCase   =	{"visual_embedding_dim": 2_048}
						UpperCamelCase   =	"vqa_advanced"
				elif "vqa" in checkpoint_path:
						UpperCamelCase   =	{"visual_embedding_dim": 2_048, "num_labels": 3_129}
						UpperCamelCase   =	"vqa"
				elif "nlvr" in checkpoint_path:
						UpperCamelCase   =	{
						    "visual_embedding_dim": 1_024,
						    "num_labels": 2,
						}
						UpperCamelCase   =	"nlvr"

		UpperCamelCase   =	VisualBertConfig(**snake_case__    )

		# Load State Dict
		UpperCamelCase   =	load_state_dict(snake_case__    )

		UpperCamelCase   =	get_new_dict(snake_case__  ,	snake_case__    )

		if model_type == "pretraining":
				UpperCamelCase   =	VisualBertForPreTraining(snake_case__    )
		elif model_type == "vqa":
				UpperCamelCase   =	VisualBertForQuestionAnswering(snake_case__    )
		elif model_type == "nlvr":
				UpperCamelCase   =	VisualBertForVisualReasoning(snake_case__    )
		elif model_type == "multichoice":
				UpperCamelCase   =	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 unittest

from huggingface_hub import hf_hub_download

from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VideoMAEFeatureExtractor
from transformers.pipelines import VideoClassificationPipeline, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_decord,
    require_tf,
    require_torch,
    require_torch_or_tf,
    require_vision,
)

from .test_pipelines_common import ANY





@is_pipeline_test
@require_torch_or_tf
@require_vision
@require_decord
class        _lowerCamelCase      (       unittest.TestCase ):
						UpperCAmelCase_ =			MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING


						def     snake_case_   (self   ,							__a   ,							__a   ,							__a		) ->   str:
								UpperCamelCase   =	hf_hub_download(
								    repo_id="nateraw/video-demo"   ,							filename="archery.mp4"   ,							repo_type="dataset"		)
								UpperCamelCase   =	VideoClassificationPipeline(model=__a   ,							image_processor=__a   ,							top_k=2		)
								UpperCamelCase   =	[
								    example_video_filepath,
								    "https://huggingface.co/datasets/nateraw/video-demo/resolve/main/archery.mp4",
								]
								return video_classifier, examples


						def     snake_case_   (self   ,							__a   ,							__a		) ->   str:
								for example in examples:
										UpperCamelCase   =	video_classifier(__a		)

										self.assertEqual(
										    __a   ,							[
										        {"score": ANY(__a		), "label": ANY(__a		)},
										        {"score": ANY(__a		), "label": ANY(__a		)},
										    ]   ,							)


						@require_torch
						def     snake_case_   (self		) ->   int:
								UpperCamelCase   =	"hf-internal-testing/tiny-random-VideoMAEForVideoClassification"
								UpperCamelCase   =	VideoMAEFeatureExtractor(
								    size={"shortest_edge": 10}   ,							crop_size={"height": 10, "width": 10}		)
								UpperCamelCase   =	pipeline(
								    "video-classification"   ,							model=__a   ,							feature_extractor=__a   ,							frame_sampling_rate=4		)

								UpperCamelCase   =	hf_hub_download(repo_id="nateraw/video-demo"   ,							filename="archery.mp4"   ,							repo_type="dataset"		)
								UpperCamelCase   =	video_classifier(__a   ,							top_k=2		)
								self.assertEqual(
								    nested_simplify(__a   ,							decimals=4		)   ,							[{"score": 0.5199, "label": "LABEL_0"}, {"score": 0.4801, "label": "LABEL_1"}]   ,							)

								UpperCamelCase   =	video_classifier(
								    [
								        video_file_path,
								        video_file_path,
								    ]   ,							top_k=2   ,							)
								self.assertEqual(
								    nested_simplify(__a   ,							decimals=4		)   ,							[
								        [{"score": 0.5199, "label": "LABEL_0"}, {"score": 0.4801, "label": "LABEL_1"}],
								        [{"score": 0.5199, "label": "LABEL_0"}, {"score": 0.4801, "label": "LABEL_1"}],
								    ]   ,							)




						@require_tf
						def     snake_case_   (self		) ->   Optional[Any]:
								pass






