
"""simple docstring"""




import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileViTConfig,
    MobileViTForImageClassification,
    MobileViTForSemanticSegmentation,
    MobileViTImageProcessor,
)
from transformers.utils import logging


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



def 							_lowerCamelCase			(			__a     ):
		SCREAMING_SNAKE_CASE_    			=			MobileViTConfig()

		# size of the architecture
		if "mobilevit_s" in mobilevit_name:
				SCREAMING_SNAKE_CASE_    			=			[144, 192, 240]
				SCREAMING_SNAKE_CASE_    			=			[16, 32, 64, 96, 128, 160, 640]
		elif "mobilevit_xs" in mobilevit_name:
				SCREAMING_SNAKE_CASE_    			=			[96, 120, 144]
				SCREAMING_SNAKE_CASE_    			=			[16, 32, 48, 64, 80, 96, 384]
		elif "mobilevit_xxs" in mobilevit_name:
				SCREAMING_SNAKE_CASE_    			=			[64, 80, 96]
				SCREAMING_SNAKE_CASE_    			=			[16, 16, 24, 48, 64, 80, 320]
				SCREAMING_SNAKE_CASE_    			=			0.0_5
				SCREAMING_SNAKE_CASE_    			=			2.0

		if mobilevit_name.startswith('''deeplabv3_'''     ):
				SCREAMING_SNAKE_CASE_    			=			512
				SCREAMING_SNAKE_CASE_    			=			16
				SCREAMING_SNAKE_CASE_    			=			21
				SCREAMING_SNAKE_CASE_    			=			'''pascal-voc-id2label.json'''
		else:
				SCREAMING_SNAKE_CASE_    			=			1_000
				SCREAMING_SNAKE_CASE_    			=			'''imagenet-1k-id2label.json'''

		SCREAMING_SNAKE_CASE_    			=			'''huggingface/label-files'''
		SCREAMING_SNAKE_CASE_    			=			json.load(open(hf_hub_download(__a,				__a,				repo_type='''dataset'''     ),				'''r'''     )     )
		SCREAMING_SNAKE_CASE_    			=			{int(__a     ): v for k, v in idalabel.items()}
		SCREAMING_SNAKE_CASE_    			=			idalabel
		SCREAMING_SNAKE_CASE_    			=			{v: k for k, v in idalabel.items()}

		return config



def 							_lowerCamelCase			(			__a,				__a=False     ):
		for i in range(1,				6     ):
				if F'layer_{i}.' in name:
						SCREAMING_SNAKE_CASE_    			=			name.replace(F'layer_{i}.',				F'encoder.layer.{i - 1}.'     )

		if "conv_1." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''conv_1.''',				'''conv_stem.'''     )
		if ".block." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.block.''',				'''.'''     )

		if "exp_1x1" in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''exp_1x1''',				'''expand_1x1'''     )
		if "red_1x1" in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''red_1x1''',				'''reduce_1x1'''     )
		if ".local_rep.conv_3x3." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.local_rep.conv_3x3.''',				'''.conv_kxk.'''     )
		if ".local_rep.conv_1x1." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.local_rep.conv_1x1.''',				'''.conv_1x1.'''     )
		if ".norm." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.norm.''',				'''.normalization.'''     )
		if ".conv." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.conv.''',				'''.convolution.'''     )
		if ".conv_proj." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.conv_proj.''',				'''.conv_projection.'''     )

		for i in range(0,				2     ):
				for j in range(0,				4     ):
						if F'.{i}.{j}.' in name:
								SCREAMING_SNAKE_CASE_    			=			name.replace(F'.{i}.{j}.',				F'.{i}.layer.{j}.'     )

		for i in range(2,				6     ):
				for j in range(0,				4     ):
						if F'.{i}.{j}.' in name:
								SCREAMING_SNAKE_CASE_    			=			name.replace(F'.{i}.{j}.',				F'.{i}.'     )
								if "expand_1x1" in name:
										SCREAMING_SNAKE_CASE_    			=			name.replace('''expand_1x1''',				'''downsampling_layer.expand_1x1'''     )
								if "conv_3x3" in name:
										SCREAMING_SNAKE_CASE_    			=			name.replace('''conv_3x3''',				'''downsampling_layer.conv_3x3'''     )
								if "reduce_1x1" in name:
										SCREAMING_SNAKE_CASE_    			=			name.replace('''reduce_1x1''',				'''downsampling_layer.reduce_1x1'''     )

		for i in range(2,				5     ):
				if F'.global_rep.{i}.weight' in name:
						SCREAMING_SNAKE_CASE_    			=			name.replace(F'.global_rep.{i}.weight',				'''.layernorm.weight'''     )
				if F'.global_rep.{i}.bias' in name:
						SCREAMING_SNAKE_CASE_    			=			name.replace(F'.global_rep.{i}.bias',				'''.layernorm.bias'''     )

		if ".global_rep." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.global_rep.''',				'''.transformer.'''     )
		if ".pre_norm_mha.0." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.pre_norm_mha.0.''',				'''.layernorm_before.'''     )
		if ".pre_norm_mha.1.out_proj." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.pre_norm_mha.1.out_proj.''',				'''.attention.output.dense.'''     )
		if ".pre_norm_ffn.0." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.pre_norm_ffn.0.''',				'''.layernorm_after.'''     )
		if ".pre_norm_ffn.1." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.pre_norm_ffn.1.''',				'''.intermediate.dense.'''     )
		if ".pre_norm_ffn.4." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.pre_norm_ffn.4.''',				'''.output.dense.'''     )
		if ".transformer." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.transformer.''',				'''.transformer.layer.'''     )

		if ".aspp_layer." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.aspp_layer.''',				'''.'''     )
		if ".aspp_pool." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''.aspp_pool.''',				'''.'''     )
		if "seg_head." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''seg_head.''',				'''segmentation_head.'''     )
		if "segmentation_head.classifier.classifier." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''segmentation_head.classifier.classifier.''',				'''segmentation_head.classifier.'''     )

		if "classifier.fc." in name:
				SCREAMING_SNAKE_CASE_    			=			name.replace('''classifier.fc.''',				'''classifier.'''     )
		elif (not base_model) and ("segmentation_head." not in name):
				SCREAMING_SNAKE_CASE_    			=			'''mobilevit.''' + name

		return name



def 							_lowerCamelCase			(			__a,				__a,				__a=False     ):
		if base_model:
				SCREAMING_SNAKE_CASE_    			=			''''''
		else:
				SCREAMING_SNAKE_CASE_    			=			'''mobilevit.'''

		for key in orig_state_dict.copy().keys():
				SCREAMING_SNAKE_CASE_    			=			orig_state_dict.pop(__a     )

				if key[:8] == "encoder.":
						SCREAMING_SNAKE_CASE_    			=			key[8:]

				if "qkv" in key:
						SCREAMING_SNAKE_CASE_    			=			key.split('''.'''     )
						SCREAMING_SNAKE_CASE_    			=			int(key_split[0][6:]     ) - 1
						SCREAMING_SNAKE_CASE_    			=			int(key_split[3]     )
						SCREAMING_SNAKE_CASE_    			=			model.get_submodule(F'{model_prefix}encoder.layer.{layer_num}'     )
						SCREAMING_SNAKE_CASE_    			=			layer.transformer.layer[transformer_num].attention.attention.all_head_size
						SCREAMING_SNAKE_CASE_    			=			(
						    F'{model_prefix}encoder.layer.{layer_num}.transformer.layer.{transformer_num}.attention.attention.'
						)
						if "weight" in key:
								SCREAMING_SNAKE_CASE_    			=			val[:dim, :]
								SCREAMING_SNAKE_CASE_    			=			val[dim : dim * 2, :]
								SCREAMING_SNAKE_CASE_    			=			val[-dim:, :]
						else:
								SCREAMING_SNAKE_CASE_    			=			val[:dim]
								SCREAMING_SNAKE_CASE_    			=			val[dim : dim * 2]
								SCREAMING_SNAKE_CASE_    			=			val[-dim:]
				else:
						SCREAMING_SNAKE_CASE_    			=			val

		return orig_state_dict



def 							_lowerCamelCase			(			):
		SCREAMING_SNAKE_CASE_    			=			'''http://images.cocodataset.org/val2017/000000039769.jpg'''
		SCREAMING_SNAKE_CASE_    			=			Image.open(requests.get(__a,				stream=__a     ).raw     )
		return im



@torch.no_grad()
def 							_lowerCamelCase			(			__a,				__a,				__a,				__a=False     ):
		SCREAMING_SNAKE_CASE_    			=			get_mobilevit_config(__a     )

		# load original state_dict
		SCREAMING_SNAKE_CASE_    			=			torch.load(__a,				map_location='''cpu'''     )

		# load 🤗 model
		if mobilevit_name.startswith('''deeplabv3_'''     ):
				SCREAMING_SNAKE_CASE_    			=			MobileViTForSemanticSegmentation(__a     ).eval()
		else:
				SCREAMING_SNAKE_CASE_    			=			MobileViTForImageClassification(__a     ).eval()

		SCREAMING_SNAKE_CASE_    			=			convert_state_dict(__a,				__a     )
		model.load_state_dict(__a     )

		# Check outputs on an image, prepared by MobileViTImageProcessor
		SCREAMING_SNAKE_CASE_    			=			MobileViTImageProcessor(crop_size=config.image_size,				size=config.image_size + 32     )
		SCREAMING_SNAKE_CASE_    			=			image_processor(images=prepare_img(),				return_tensors='''pt'''     )
		SCREAMING_SNAKE_CASE_    			=			model(**__a     )
		SCREAMING_SNAKE_CASE_    			=			outputs.logits

		if mobilevit_name.startswith('''deeplabv3_'''     ):
				assert logits.shape == (1, 21, 32, 32)

				if mobilevit_name == "deeplabv3_mobilevit_s":
						SCREAMING_SNAKE_CASE_    			=			torch.tensor(
						    [
						        [[6.2_0_6_5, 6.1_2_9_2, 6.2_0_7_0], [6.1_0_7_9, 6.1_2_5_4, 6.1_7_4_7], [6.0_0_4_2, 6.1_0_7_1, 6.1_0_3_4]],
						        [[-6.9_2_5_3, -6.8_6_5_3, -7.0_3_9_8], [-7.3_2_1_8, -7.3_9_8_3, -7.3_6_7_0], [-7.1_9_6_1, -7.2_4_8_2, -7.1_5_6_9]],
						        [[-4.4_7_2_3, -4.4_3_4_8, -4.3_7_6_9], [-5.3_6_2_9, -5.4_6_3_2, -5.4_5_9_8], [-5.1_5_8_7, -5.3_4_0_2, -5.5_0_5_9]],
						    ]     )
				elif mobilevit_name == "deeplabv3_mobilevit_xs":
						SCREAMING_SNAKE_CASE_    			=			torch.tensor(
						    [
						        [[5.4_4_4_9, 5.5_7_3_3, 5.6_3_1_4], [5.1_8_1_5, 5.3_9_3_0, 5.5_9_6_3], [5.1_6_5_6, 5.4_3_3_3, 5.4_8_5_3]],
						        [[-9.4_4_2_3, -9.7_7_6_6, -9.6_7_1_4], [-9.1_5_8_1, -9.5_7_2_0, -9.5_5_1_9], [-9.1_0_0_6, -9.6_4_5_8, -9.5_7_0_3]],
						        [[-7.7_7_2_1, -7.3_7_1_6, -7.1_5_8_3], [-8.4_5_9_9, -8.0_6_2_4, -7.7_9_4_4], [-8.4_1_7_2, -7.8_3_6_6, -7.5_0_2_5]],
						    ]     )
				elif mobilevit_name == "deeplabv3_mobilevit_xxs":
						SCREAMING_SNAKE_CASE_    			=			torch.tensor(
						    [
						        [[6.9_8_1_1, 6.9_7_4_3, 7.3_1_2_3], [7.1_7_7_7, 7.1_9_3_1, 7.3_9_3_8], [7.5_6_3_3, 7.8_0_5_0, 7.8_9_0_1]],
						        [[-1_0.5_5_3_6, -1_0.2_3_3_2, -1_0.2_9_2_4], [-1_0.2_3_3_6, -9.8_6_2_4, -9.5_9_6_4], [-1_0.8_8_4_0, -1_0.8_1_5_8, -1_0.6_6_5_9]],
						        [[-3.4_9_3_8, -3.0_6_3_1, -2.8_6_2_0], [-3.4_2_0_5, -2.8_1_3_5, -2.6_8_7_5], [-3.4_1_7_9, -2.7_9_4_5, -2.8_7_5_0]],
						    ]     )
				else:
						raise ValueError(F'Unknown mobilevit_name: {mobilevit_name}'     )

				assert torch.allclose(logits[0, :3, :3, :3],				__a,				atol=1E-4     )
		else:
				assert logits.shape == (1, 1_000)

				if mobilevit_name == "mobilevit_s":
						SCREAMING_SNAKE_CASE_    			=			torch.tensor([-0.9_8_6_6, 0.2_3_9_2, -1.1_2_4_1]     )
				elif mobilevit_name == "mobilevit_xs":
						SCREAMING_SNAKE_CASE_    			=			torch.tensor([-2.4_7_6_1, -0.9_3_9_9, -1.9_5_8_7]     )
				elif mobilevit_name == "mobilevit_xxs":
						SCREAMING_SNAKE_CASE_    			=			torch.tensor([-1.9_3_6_4, -1.2_3_2_7, -0.4_6_5_3]     )
				else:
						raise ValueError(F'Unknown mobilevit_name: {mobilevit_name}'     )

				assert torch.allclose(logits[0, :3],				__a,				atol=1E-4     )

		Path(__a     ).mkdir(exist_ok=__a     )
		print(F'Saving model {mobilevit_name} to {pytorch_dump_folder_path}'     )
		model.save_pretrained(__a     )
		print(F'Saving image processor to {pytorch_dump_folder_path}'     )
		image_processor.save_pretrained(__a     )

		if push_to_hub:
				SCREAMING_SNAKE_CASE_    			=			{
				    '''mobilevit_s''': '''mobilevit-small''',
				    '''mobilevit_xs''': '''mobilevit-x-small''',
				    '''mobilevit_xxs''': '''mobilevit-xx-small''',
				    '''deeplabv3_mobilevit_s''': '''deeplabv3-mobilevit-small''',
				    '''deeplabv3_mobilevit_xs''': '''deeplabv3-mobilevit-x-small''',
				    '''deeplabv3_mobilevit_xxs''': '''deeplabv3-mobilevit-xx-small''',
				}

				print('''Pushing to the hub...'''     )
				SCREAMING_SNAKE_CASE_    			=			model_mapping[mobilevit_name]
				image_processor.push_to_hub(__a,				organization='''apple'''     )
				model.push_to_hub(__a,				organization='''apple'''     )


if __name__ == "__main__":
						lowerCAmelCase__				=					argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--mobilevit_name',
						    default='mobilevit_s',
						    type=str,
						    help=(
						        'Name of the MobileViT model you\'d like to convert. Should be one of \'mobilevit_s\', \'mobilevit_xs\','
						        ' \'mobilevit_xxs\', \'deeplabv3_mobilevit_s\', \'deeplabv3_mobilevit_xs\', \'deeplabv3_mobilevit_xxs\'.'
						    ),
						)
						parser.add_argument(
						    '--checkpoint_path', required=True, type=str, help='Path to the original state dict (.pt file).'
						)
						parser.add_argument(
						    '--pytorch_dump_folder_path', required=True, type=str, help='Path to the output PyTorch model directory.'
						)
						parser.add_argument(
						    '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
						)

						lowerCAmelCase__				=					parser.parse_args()
						convert_movilevit_checkpoint(
						    args.mobilevit_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
						)

"""simple docstring"""




import datasets

from .evaluate import evaluate


lowerCAmelCase__				=					'\\n@inproceedings{Rajpurkar2016SQuAD10,\n  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},\n  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},\n  booktitle={EMNLP},\n  year={2016}\n}\n'

lowerCAmelCase__				=					'\nThis metric wrap the official scoring script for version 1 of the Stanford Question Answering Dataset (SQuAD).\n\nStanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by\ncrowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,\nfrom the corresponding reading passage, or the question might be unanswerable.\n'

lowerCAmelCase__				=					'\nComputes SQuAD scores (F1 and EM).\nArgs:\n    predictions: List of question-answers dictionaries with the following key-values:\n        - \'id\': id of the question-answer pair as given in the references (see below)\n        - \'prediction_text\': the text of the answer\n    references: List of question-answers dictionaries with the following key-values:\n        - \'id\': id of the question-answer pair (see above),\n        - \'answers\': a Dict in the SQuAD dataset format\n            {\n                \'text\': list of possible texts for the answer, as a list of strings\n                \'answer_start\': list of start positions for the answer, as a list of ints\n            }\n            Note that answer_start values are not taken into account to compute the metric.\nReturns:\n    \'exact_match\': Exact match (the normalized answer exactly match the gold answer)\n    \'f1\': The F-score of predicted tokens versus the gold answer\nExamples:\n\n    >>> predictions = [{\'prediction_text\': \'1976\', \'id\': \'56e10a3be3433e1400422b22\'}]\n    >>> references = [{\'answers\': {\'answer_start\': [97], \'text\': [\'1976\']}, \'id\': \'56e10a3be3433e1400422b22\'}]\n    >>> squad_metric = datasets.load_metric("squad")\n    >>> results = squad_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'exact_match\': 100.0, \'f1\': 100.0}\n'





@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION							,      _KWARGS_DESCRIPTION )
class 			snake_case  (				datasets.Metric ):

							def _lowercase       (self     ):

									"""simple docstring"""







									return datasets.MetricInfo(
									    description=_DESCRIPTION							,							citation=_CITATION							,							inputs_description=_KWARGS_DESCRIPTION							,							features=datasets.Features(
									        {
									            '''predictions''': {'''id''': datasets.Value('''string'''     ), '''prediction_text''': datasets.Value('''string'''     )},
									            '''references''': {
									                '''id''': datasets.Value('''string'''     ),
									                '''answers''': datasets.features.Sequence(
									                    {
									                        '''text''': datasets.Value('''string'''     ),
									                        '''answer_start''': datasets.Value('''int32'''     ),
									                    }     ),
									            },
									        }     )							,							codebase_urls=['''https://rajpurkar.github.io/SQuAD-explorer/''']							,							reference_urls=['''https://rajpurkar.github.io/SQuAD-explorer/''']							,							)







							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			{prediction['''id''']: prediction['''prediction_text'''] for prediction in predictions}
									SCREAMING_SNAKE_CASE_    			=			[
									    {
									        '''paragraphs''': [
									            {
									                '''qas''': [
									                    {
									                        '''answers''': [{'''text''': answer_text} for answer_text in ref['''answers''']['''text''']],
									                        '''id''': ref['''id'''],
									                    }
									                    for ref in references
									                ]
									            }
									        ]
									    }
									]
									SCREAMING_SNAKE_CASE_    			=			evaluate(dataset=SCREAMING_SNAKE_CASE_							,							predictions=SCREAMING_SNAKE_CASE_     )
									return score