



"""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()
__snake_case							:     str   =					logging.get_logger(__name__)


def     _lowercase       (						__snake_case      ,__snake_case=False      ,__snake_case=False      ,__snake_case=False			)		->     Union[str, Any]:
			__lowerCAmelCase   :						Optional[Any]				= []
			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     _lowercase       (						__snake_case      ,__snake_case			)		->     Optional[Any]:
			for i in range(config.num_hidden_layers			):
						__lowerCAmelCase   :						List[str]				= "vilt."
						# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
						__lowerCAmelCase   :						Optional[int]				= state_dict.pop(F"""transformer.blocks.{i}.attn.qkv.weight"""			)
						__lowerCAmelCase   :						Optional[Any]				= state_dict.pop(F"""transformer.blocks.{i}.attn.qkv.bias"""			)
						# next, add query, keys and values (in that order) to the state dict
						__lowerCAmelCase   :						List[str]				= in_proj_weight[
						    : config.hidden_size, :
						]
						__lowerCAmelCase   :						List[Any]				= in_proj_bias[: config.hidden_size]
						__lowerCAmelCase   :						List[Any]				= in_proj_weight[
						    config.hidden_size : config.hidden_size * 2, :
						]
						__lowerCAmelCase   :						Dict				= in_proj_bias[
						    config.hidden_size : config.hidden_size * 2
						]
						__lowerCAmelCase   :						str				= in_proj_weight[
						    -config.hidden_size :, :
						]
						__lowerCAmelCase   :						List[Any]				= in_proj_bias[-config.hidden_size :]


def     _lowercase       (						__snake_case			)		->     List[Any]:
			__lowerCAmelCase   :						Dict				= ["head.weight", "head.bias"]
			for k in ignore_keys:
						state_dict.pop(__snake_case      ,__snake_case			)


def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case			)		->     List[Any]:
			__lowerCAmelCase   :						int				= dct.pop(__snake_case			)
			__lowerCAmelCase   :						int				= val


@torch.no_grad()
def     _lowercase       (						__snake_case      ,__snake_case			)		->     List[str]:
			__lowerCAmelCase   :						List[str]				= ViltConfig(image_size=384      ,patch_size=32      ,tie_word_embeddings=__snake_case			)
			__lowerCAmelCase   :						str				= False
			__lowerCAmelCase   :						List[str]				= False
			__lowerCAmelCase   :						List[str]				= False
			__lowerCAmelCase   :						Any				= False
			if "vqa" in checkpoint_url:
						__lowerCAmelCase   :						str				= True
						__lowerCAmelCase   :						Union[str, Any]				= 3_129
						__lowerCAmelCase   :						str				= "huggingface/label-files"
						__lowerCAmelCase   :						Tuple				= "vqa2-id2label.json"
						__lowerCAmelCase   :						Any				= json.load(open(hf_hub_download(__snake_case      ,__snake_case      ,repo_type="dataset"			)      ,"r"			)			)
						__lowerCAmelCase   :						Tuple				= {int(__snake_case			): v for k, v in idalabel.items()}
						__lowerCAmelCase   :						Union[str, Any]				= idalabel
						__lowerCAmelCase   :						Optional[int]				= {v: k for k, v in idalabel.items()}
						__lowerCAmelCase   :						Tuple				= ViltForQuestionAnswering(__snake_case			)
			elif "nlvr" in checkpoint_url:
						__lowerCAmelCase   :						Dict				= True
						__lowerCAmelCase   :						Tuple				= 2
						__lowerCAmelCase   :						List[Any]				= {0: "False", 1: "True"}
						__lowerCAmelCase   :						List[Any]				= {v: k for k, v in config.idalabel.items()}
						__lowerCAmelCase   :						Any				= 3
						__lowerCAmelCase   :						str				= ViltForImagesAndTextClassification(__snake_case			)
			elif "irtr" in checkpoint_url:
						__lowerCAmelCase   :						str				= True
						__lowerCAmelCase   :						List[str]				= ViltForImageAndTextRetrieval(__snake_case			)
			elif "mlm_itm" in checkpoint_url:
						__lowerCAmelCase   :						str				= True
						__lowerCAmelCase   :						Tuple				= ViltForMaskedLM(__snake_case			)
			else:
						raise ValueError("Unknown model type"			)

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

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

			# Define processor
			__lowerCAmelCase   :						Dict				= ViltImageProcessor(size=384			)
			__lowerCAmelCase   :						List[Any]				= BertTokenizer.from_pretrained("bert-base-uncased"			)
			__lowerCAmelCase   :						List[Any]				= ViltProcessor(__snake_case      ,__snake_case			)

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

			# Verify outputs
			if mlm_model:
						__lowerCAmelCase   :						int				= torch.Size([1, 11, 30_522]			)
						__lowerCAmelCase   :						int				= torch.tensor([-12.5061, -12.5123, -12.5174]			)
						assert outputs.logits.shape == expected_shape
						assert torch.allclose(outputs.logits[0, 0, :3]      ,__snake_case      ,atol=1e-4			)

						# verify masked token prediction equals "cats"
						__lowerCAmelCase   :						Optional[int]				= outputs.logits[0, 4, :].argmax(-1			).item()
						assert tokenizer.decode([predicted_id]			) == "cats"
			elif vqa_model:
						__lowerCAmelCase   :						Optional[Any]				= torch.Size([1, 3_129]			)
						__lowerCAmelCase   :						str				= torch.tensor([-15.9495, -18.1472, -10.3041]			)
						assert torch.allclose(outputs.logits[0, :3]      ,__snake_case      ,atol=1e-4			)
						assert outputs.logits.shape == expected_shape
						assert torch.allclose(outputs.logits[0, 0, :3]      ,__snake_case      ,atol=1e-4			)

						# verify vqa prediction equals "2"
						__lowerCAmelCase   :						Dict				= outputs.logits.argmax(-1			).item()
						assert model.config.idalabel[predicted_idx] == "2"
			elif nlvr_model:
						__lowerCAmelCase   :						List[Any]				= torch.Size([1, 2]			)
						__lowerCAmelCase   :						Any				= torch.tensor([-2.8721, 2.1291]			)
						assert torch.allclose(outputs.logits[0, :3]      ,__snake_case      ,atol=1e-4			)
						assert outputs.logits.shape == expected_shape

			Path(__snake_case			).mkdir(exist_ok=__snake_case			)
			print(F"""Saving model and processor to {pytorch_dump_folder_path}"""			)
			model.save_pretrained(__snake_case			)
			processor.save_pretrained(__snake_case			)


if __name__ == "__main__":
	__snake_case							:     int   =					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.'
	)

	__snake_case							:     Tuple   =					parser.parse_args()
	convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




"""simple docstring"""






import pyarrow.parquet as pq
import pytest

from datasets import Audio, Dataset, DatasetDict, Features, NamedSplit, Sequence, Value, config
from datasets.features.image import Image
from datasets.io.parquet import ParquetDatasetReader, ParquetDatasetWriter, get_writer_batch_size

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases


def     _lowercase       (						__snake_case      ,__snake_case			)		->     Dict:
			assert isinstance(__snake_case      ,__snake_case			)
			assert dataset.num_rows == 4
			assert dataset.num_columns == 3
			assert dataset.column_names == ["col_1", "col_2", "col_3"]
			for feature, expected_dtype in expected_features.items():
						assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("keep_in_memory"      ,[False, True]			)
def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case			)		->     Any:
			__lowerCAmelCase   :						Dict				= tmp_path / "cache"
			__lowerCAmelCase   :						Union[str, Any]				= {"col_1": "string", "col_2": "int64", "col_3": "float64"}
			with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
						__lowerCAmelCase   :						Optional[Any]				= ParquetDatasetReader(__snake_case      ,cache_dir=__snake_case      ,keep_in_memory=__snake_case			).read()
			_check_parquet_dataset(__snake_case      ,__snake_case			)


@pytest.mark.parametrize(
    "features"      ,[
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ]      ,)
def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case			)		->     str:
			__lowerCAmelCase   :						Tuple				= tmp_path / "cache"
			__lowerCAmelCase   :						Dict				= {"col_1": "string", "col_2": "int64", "col_3": "float64"}
			__lowerCAmelCase   :						List[Any]				= features.copy() if features else default_expected_features
			__lowerCAmelCase   :						Dict				= (
			    Features({feature: Value(__snake_case			) for feature, dtype in features.items()}			) if features is not None else None
			)
			__lowerCAmelCase   :						Optional[int]				= ParquetDatasetReader(__snake_case      ,features=__snake_case      ,cache_dir=__snake_case			).read()
			_check_parquet_dataset(__snake_case      ,__snake_case			)


@pytest.mark.parametrize("split"      ,[None, NamedSplit("train"			), "train", "test"]			)
def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case			)		->     Any:
			__lowerCAmelCase   :						List[str]				= tmp_path / "cache"
			__lowerCAmelCase   :						List[str]				= {"col_1": "string", "col_2": "int64", "col_3": "float64"}
			__lowerCAmelCase   :						List[Any]				= ParquetDatasetReader(__snake_case      ,cache_dir=__snake_case      ,split=__snake_case			).read()
			_check_parquet_dataset(__snake_case      ,__snake_case			)
			assert dataset.split == split if split else "train"


@pytest.mark.parametrize("path_type"      ,[str, list]			)
def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case			)		->     Dict:
			if issubclass(__snake_case      ,__snake_case			):
						__lowerCAmelCase   :						List[Any]				= parquet_path
			elif issubclass(__snake_case      ,__snake_case			):
						__lowerCAmelCase   :						List[Any]				= [parquet_path]
			__lowerCAmelCase   :						str				= tmp_path / "cache"
			__lowerCAmelCase   :						List[str]				= {"col_1": "string", "col_2": "int64", "col_3": "float64"}
			__lowerCAmelCase   :						Tuple				= ParquetDatasetReader(__snake_case      ,cache_dir=__snake_case			).read()
			_check_parquet_dataset(__snake_case      ,__snake_case			)


def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case=("train",)			)		->     int:
			assert isinstance(__snake_case      ,__snake_case			)
			for split in splits:
						__lowerCAmelCase   :						str				= dataset_dict[split]
						assert dataset.num_rows == 4
						assert dataset.num_columns == 3
						assert dataset.column_names == ["col_1", "col_2", "col_3"]
						for feature, expected_dtype in expected_features.items():
									assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("keep_in_memory"      ,[False, True]			)
def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case			)		->     str:
			__lowerCAmelCase   :						Any				= tmp_path / "cache"
			__lowerCAmelCase   :						int				= {"col_1": "string", "col_2": "int64", "col_3": "float64"}
			with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
						__lowerCAmelCase   :						str				= ParquetDatasetReader(
						    {"train": parquet_path}      ,cache_dir=__snake_case      ,keep_in_memory=__snake_case			).read()
			_check_parquet_datasetdict(__snake_case      ,__snake_case			)


@pytest.mark.parametrize(
    "features"      ,[
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ]      ,)
def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case			)		->     Dict:
			__lowerCAmelCase   :						List[Any]				= tmp_path / "cache"
			__lowerCAmelCase   :						Dict				= {"col_1": "string", "col_2": "int64", "col_3": "float64"}
			__lowerCAmelCase   :						str				= features.copy() if features else default_expected_features
			__lowerCAmelCase   :						str				= (
			    Features({feature: Value(__snake_case			) for feature, dtype in features.items()}			) if features is not None else None
			)
			__lowerCAmelCase   :						str				= ParquetDatasetReader({"train": parquet_path}      ,features=__snake_case      ,cache_dir=__snake_case			).read()
			_check_parquet_datasetdict(__snake_case      ,__snake_case			)


@pytest.mark.parametrize("split"      ,[None, NamedSplit("train"			), "train", "test"]			)
def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case			)		->     str:
			if split:
						__lowerCAmelCase   :						Optional[int]				= {split: parquet_path}
			else:
						__lowerCAmelCase   :						str				= "train"
						__lowerCAmelCase   :						Optional[Any]				= {"train": parquet_path, "test": parquet_path}
			__lowerCAmelCase   :						Tuple				= tmp_path / "cache"
			__lowerCAmelCase   :						int				= {"col_1": "string", "col_2": "int64", "col_3": "float64"}
			__lowerCAmelCase   :						Union[str, Any]				= ParquetDatasetReader(__snake_case      ,cache_dir=__snake_case			).read()
			_check_parquet_datasetdict(__snake_case      ,__snake_case      ,splits=list(path.keys()			)			)
			assert all(dataset[split].split == split for split in path.keys()			)


def     _lowercase       (						__snake_case      ,__snake_case			)		->     Union[str, Any]:
			__lowerCAmelCase   :						int				= ParquetDatasetWriter(__snake_case      ,tmp_path / "foo.parquet"			)
			assert writer.write() > 0
			__lowerCAmelCase   :						str				= pq.ParquetFile(tmp_path / "foo.parquet"			)
			__lowerCAmelCase   :						int				= pf.read()
			assert dataset.data.table == output_table


def     _lowercase       (						__snake_case      ,__snake_case			)		->     int:
			__lowerCAmelCase   :						int				= str(shared_datadir / "test_image_rgb.jpg"			)
			__lowerCAmelCase   :						Any				= {"image": [image_path]}
			__lowerCAmelCase   :						List[Any]				= Features({"image": Image()}			)
			__lowerCAmelCase   :						Tuple				= Dataset.from_dict(__snake_case      ,features=__snake_case			)
			__lowerCAmelCase   :						Dict				= ParquetDatasetWriter(__snake_case      ,tmp_path / "foo.parquet"			)
			assert writer.write() > 0

			__lowerCAmelCase   :						List[Any]				= Dataset.from_parquet(str(tmp_path / "foo.parquet"			)			)
			assert dataset.features == reloaded_dataset.features

			__lowerCAmelCase   :						List[Any]				= ParquetDatasetReader(str(tmp_path / "foo.parquet"			)      ,streaming=__snake_case			).read()
			assert dataset.features == reloaded_iterable_dataset.features





@pytest.mark.parametrize(
    "feature, expected"      ,[
        (Features({"foo": Value("int32"			)}			), None),
        (Features({"image": Image(), "foo": Value("int32"			)}			), config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS),
        (Features({"nested": Sequence(Audio()			)}			), config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS),
    ]      ,)
def     _lowercase       (						__snake_case      ,__snake_case			)		->     List[Any]:
			assert get_writer_batch_size(__snake_case			) == expected