





'''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 (
    BertTokenizer,
    ViltConfig,
    ViltForImageAndTextRetrieval,
    ViltForImagesAndTextClassification,
    ViltForMaskedLM,
    ViltForQuestionAnswering,
    ViltImageProcessor,
    ViltProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__lowerCamelCase : Union[str, Any]											=			logging.get_logger(__name__)







def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_=False		,							lowerCAmelCase_=False		,							lowerCAmelCase_=False							):

				"""simple docstring"""






				lowercase         = []
				for i in range(config.num_hidden_layers							):
								# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
								rename_keys.append((f'transformer.blocks.{i}.norm1.weight', f'vilt.encoder.layer.{i}.layernorm_before.weight')							)
								rename_keys.append((f'transformer.blocks.{i}.norm1.bias', f'vilt.encoder.layer.{i}.layernorm_before.bias')							)
								rename_keys.append(
								    (f'transformer.blocks.{i}.attn.proj.weight', f'vilt.encoder.layer.{i}.attention.output.dense.weight')							)
								rename_keys.append(
								    (f'transformer.blocks.{i}.attn.proj.bias', f'vilt.encoder.layer.{i}.attention.output.dense.bias')							)
								rename_keys.append((f'transformer.blocks.{i}.norm2.weight', f'vilt.encoder.layer.{i}.layernorm_after.weight')							)
								rename_keys.append((f'transformer.blocks.{i}.norm2.bias', f'vilt.encoder.layer.{i}.layernorm_after.bias')							)
								rename_keys.append(
								    (f'transformer.blocks.{i}.mlp.fc1.weight', f'vilt.encoder.layer.{i}.intermediate.dense.weight')							)
								rename_keys.append((f'transformer.blocks.{i}.mlp.fc1.bias', f'vilt.encoder.layer.{i}.intermediate.dense.bias')							)
								rename_keys.append((f'transformer.blocks.{i}.mlp.fc2.weight', f'vilt.encoder.layer.{i}.output.dense.weight')							)
								rename_keys.append((f'transformer.blocks.{i}.mlp.fc2.bias', f'vilt.encoder.layer.{i}.output.dense.bias')							)

				# embeddings
				rename_keys.extend(
				    [
				        # text embeddings
				        ("text_embeddings.word_embeddings.weight", "vilt.embeddings.text_embeddings.word_embeddings.weight"),
				        (
				            "text_embeddings.position_embeddings.weight",
				            "vilt.embeddings.text_embeddings.position_embeddings.weight",
				        ),
				        ("text_embeddings.position_ids", "vilt.embeddings.text_embeddings.position_ids"),
				        (
				            "text_embeddings.token_type_embeddings.weight",
				            "vilt.embeddings.text_embeddings.token_type_embeddings.weight",
				        ),
				        ("text_embeddings.LayerNorm.weight", "vilt.embeddings.text_embeddings.LayerNorm.weight"),
				        ("text_embeddings.LayerNorm.bias", "vilt.embeddings.text_embeddings.LayerNorm.bias"),
				        # patch embeddings
				        ("transformer.cls_token", "vilt.embeddings.cls_token"),
				        ("transformer.patch_embed.proj.weight", "vilt.embeddings.patch_embeddings.projection.weight"),
				        ("transformer.patch_embed.proj.bias", "vilt.embeddings.patch_embeddings.projection.bias"),
				        ("transformer.pos_embed", "vilt.embeddings.position_embeddings"),
				        # token type embeddings
				        ("token_type_embeddings.weight", "vilt.embeddings.token_type_embeddings.weight"),
				    ]							)

				# final layernorm + pooler
				rename_keys.extend(
				    [
				        ("transformer.norm.weight", "vilt.layernorm.weight"),
				        ("transformer.norm.bias", "vilt.layernorm.bias"),
				        ("pooler.dense.weight", "vilt.pooler.dense.weight"),
				        ("pooler.dense.bias", "vilt.pooler.dense.bias"),
				    ]							)

				# classifier head(s)
				if vqa_model:
								# classification head
								rename_keys.extend(
								    [
								        ("vqa_classifier.0.weight", "classifier.0.weight"),
								        ("vqa_classifier.0.bias", "classifier.0.bias"),
								        ("vqa_classifier.1.weight", "classifier.1.weight"),
								        ("vqa_classifier.1.bias", "classifier.1.bias"),
								        ("vqa_classifier.3.weight", "classifier.3.weight"),
								        ("vqa_classifier.3.bias", "classifier.3.bias"),
								    ]							)
				elif nlvr_model:
								# classification head
								rename_keys.extend(
								    [
								        ("nlvr2_classifier.0.weight", "classifier.0.weight"),
								        ("nlvr2_classifier.0.bias", "classifier.0.bias"),
								        ("nlvr2_classifier.1.weight", "classifier.1.weight"),
								        ("nlvr2_classifier.1.bias", "classifier.1.bias"),
								        ("nlvr2_classifier.3.weight", "classifier.3.weight"),
								        ("nlvr2_classifier.3.bias", "classifier.3.bias"),
								    ]							)
				else:
								pass

				return rename_keys







def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_							):

				"""simple docstring"""






				for i in range(config.num_hidden_layers							):
								lowercase         = "vilt."
								# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
								lowercase         = state_dict.pop(f'transformer.blocks.{i}.attn.qkv.weight'							)
								lowercase         = state_dict.pop(f'transformer.blocks.{i}.attn.qkv.bias'							)
								# next, add query, keys and values (in that order) to the state dict
								lowercase         = in_proj_weight[
								    : config.hidden_size, :
								]
								lowercase         = in_proj_bias[: config.hidden_size]
								lowercase         = in_proj_weight[
								    config.hidden_size : config.hidden_size * 2, :
								]
								lowercase         = in_proj_bias[
								    config.hidden_size : config.hidden_size * 2
								]
								lowercase         = in_proj_weight[
								    -config.hidden_size :, :
								]
								lowercase         = in_proj_bias[-config.hidden_size :]







def 							UpperCAmelCase_    (	lowerCAmelCase_							):

				"""simple docstring"""






				lowercase         = ["head.weight", "head.bias"]
				for k in ignore_keys:
								state_dict.pop(lowerCAmelCase_		,							lowerCAmelCase_							)







def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_		,							lowerCAmelCase_							):

				"""simple docstring"""






				lowercase         = dct.pop(lowerCAmelCase_							)
				lowercase         = val







@torch.no_grad()
def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_							):

				"""simple docstring"""






				lowercase         = ViltConfig(image_size=384		,							patch_size=32		,							tie_word_embeddings=lowerCAmelCase_							)
				lowercase         = False
				lowercase         = False
				lowercase         = False
				lowercase         = False
				if "vqa" in checkpoint_url:
								lowercase         = True
								lowercase         = 3129
								lowercase         = "huggingface/label-files"
								lowercase         = "vqa2-id2label.json"
								lowercase         = json.load(open(hf_hub_download(lowerCAmelCase_		,							lowerCAmelCase_		,							repo_type="dataset"							)		,							"r"							)							)
								lowercase         = {int(lowerCAmelCase_							): v for k, v in idalabel.items()}
								lowercase         = idalabel
								lowercase         = {v: k for k, v in idalabel.items()}
								lowercase         = ViltForQuestionAnswering(lowerCAmelCase_							)
				elif "nlvr" in checkpoint_url:
								lowercase         = True
								lowercase         = 2
								lowercase         = {0: "False", 1: "True"}
								lowercase         = {v: k for k, v in config.idalabel.items()}
								lowercase         = 3
								lowercase         = ViltForImagesAndTextClassification(lowerCAmelCase_							)
				elif "irtr" in checkpoint_url:
								lowercase         = True
								lowercase         = ViltForImageAndTextRetrieval(lowerCAmelCase_							)
				elif "mlm_itm" in checkpoint_url:
								lowercase         = True
								lowercase         = ViltForMaskedLM(lowerCAmelCase_							)
				else:
								raise ValueError("Unknown model type"							)

				# load state_dict of original model, remove and rename some keys
				lowercase         = torch.hub.load_state_dict_from_url(lowerCAmelCase_		,							map_location="cpu"							)["state_dict"]
				lowercase         = create_rename_keys(lowerCAmelCase_		,							lowerCAmelCase_		,							lowerCAmelCase_		,							lowerCAmelCase_							)
				for src, dest in rename_keys:
								rename_key(lowerCAmelCase_		,							lowerCAmelCase_		,							lowerCAmelCase_							)
				read_in_q_k_v(lowerCAmelCase_		,							lowerCAmelCase_							)
				if mlm_model or irtr_model:
								lowercase         = ["itm_score.fc.weight", "itm_score.fc.bias"]
								for k in ignore_keys:
												state_dict.pop(lowerCAmelCase_		,							lowerCAmelCase_							)

    # load state dict into HuggingFace model
				model.eval()
				if mlm_model:
								lowercase   ,							lowercase         = model.load_state_dict(lowerCAmelCase_		,							strict=lowerCAmelCase_							)
								assert missing_keys == ["mlm_score.decoder.bias"]
				else:
								model.load_state_dict(lowerCAmelCase_							)

				# Define processor
				lowercase         = ViltImageProcessor(size=384							)
				lowercase         = BertTokenizer.from_pretrained("bert-base-uncased"							)
				lowercase         = ViltProcessor(lowerCAmelCase_		,							lowerCAmelCase_							)

				# Forward pass on example inputs (image + text)
				if nlvr_model:
								lowercase         = Image.open(requests.get("https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg"		,							stream=lowerCAmelCase_							).raw							)
								lowercase         = Image.open(requests.get("https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg"		,							stream=lowerCAmelCase_							).raw							)
								lowercase         = (
								    "The left image contains twice the number of dogs as the right image, and at least two dogs in total are"
								    " standing."
								)
								lowercase         = processor(lowerCAmelCase_		,							lowerCAmelCase_		,							return_tensors="pt"							)
								lowercase         = processor(lowerCAmelCase_		,							lowerCAmelCase_		,							return_tensors="pt"							)
								lowercase         = model(
								    input_ids=encoding_a.input_ids		,							pixel_values=encoding_a.pixel_values		,							pixel_values_a=encoding_a.pixel_values		,							)
				else:
								lowercase         = Image.open(requests.get("http://images.cocodataset.org/val2017/000000039769.jpg"		,							stream=lowerCAmelCase_							).raw							)
								if mlm_model:
												lowercase         = "a bunch of [MASK] laying on a [MASK]."
								else:
												lowercase         = "How many cats are there?"
								lowercase         = processor(lowerCAmelCase_		,							lowerCAmelCase_		,							return_tensors="pt"							)
								lowercase         = model(**lowerCAmelCase_							)

				# Verify outputs
				if mlm_model:
								lowercase         = torch.Size([1, 11, 3_0522]							)
								lowercase         = torch.tensor([-12.50_61, -12.51_23, -12.51_74]							)
								assert outputs.logits.shape == expected_shape
								assert torch.allclose(outputs.logits[0, 0, :3]		,							lowerCAmelCase_		,							atol=1E-4							)

								# verify masked token prediction equals "cats"
								lowercase         = outputs.logits[0, 4, :].argmax(-1							).item()
								assert tokenizer.decode([predicted_id]							) == "cats"
				elif vqa_model:
								lowercase         = torch.Size([1, 3129]							)
								lowercase         = torch.tensor([-15.94_95, -18.14_72, -10.30_41]							)
								assert torch.allclose(outputs.logits[0, :3]		,							lowerCAmelCase_		,							atol=1E-4							)
								assert outputs.logits.shape == expected_shape
								assert torch.allclose(outputs.logits[0, 0, :3]		,							lowerCAmelCase_		,							atol=1E-4							)

								# verify vqa prediction equals "2"
								lowercase         = outputs.logits.argmax(-1							).item()
								assert model.config.idalabel[predicted_idx] == "2"
				elif nlvr_model:
								lowercase         = torch.Size([1, 2]							)
								lowercase         = torch.tensor([-2.87_21, 2.12_91]							)
								assert torch.allclose(outputs.logits[0, :3]		,							lowerCAmelCase_		,							atol=1E-4							)
								assert outputs.logits.shape == expected_shape

				Path(lowerCAmelCase_							).mkdir(exist_ok=lowerCAmelCase_							)
				print(f'Saving model and processor to {pytorch_dump_folder_path}'							)
				model.save_pretrained(lowerCAmelCase_							)
				processor.save_pretrained(lowerCAmelCase_							)


if __name__ == "__main__":
			__lowerCamelCase : Dict											=			argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    "--checkpoint_url",
			    default="https://github.com/dandelin/ViLT/releases/download/200k/vilt_200k_mlm_itm.ckpt",
			    type=str,
			    help="URL of the checkpoint you'd like to convert.",
			)
			parser.add_argument(
			    "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
			)

			__lowerCamelCase : Union[str, Any]											=			parser.parse_args()
			convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)











'''simple docstring'''

import datasets

from .evaluate import evaluate


__lowerCamelCase : Optional[Any]											=			"\\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 : Any											=			"\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 : Union[str, Any]											=			"\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       UpperCAmelCase (						datasets.Metric  ):




		def 				UpperCAmelCase__   (self      :	Union[str, Any]				)					-> Union[str, Any]:
						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 				UpperCAmelCase__   (self      :	List[str]							,			A__      :	List[Any]							,			A__      :	List[str]				)					-> str:
						lowercase         = {prediction["id"]: prediction["prediction_text"] for prediction in predictions}
						lowercase         = [
						    {
						        "paragraphs": [
						            {
						                "qas": [
						                    {
						                        "answers": [{"text": answer_text} for answer_text in ref["answers"]["text"]],
						                        "id": ref["id"],
						                    }
						                    for ref in references
						                ]
						            }
						        ]
						    }
						]
						lowercase         = evaluate(dataset=A__							,			predictions=A__				)
						return score




