



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__					  =  logging.get_logger(__name__)


def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__=False					,			lowerCamelCase__=False					,			lowerCamelCase__=False						):

				"""simple docstring"""


				lowercase__					:   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  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__						):

				"""simple docstring"""


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


def  __lowerCamelCase (       lowerCamelCase__						):

				"""simple docstring"""


				lowercase__					:   Optional[int]													=					["head.weight", "head.bias"]
				for k in ignore_keys:
								state_dict.pop(lowerCamelCase__					,			lowerCamelCase__						)


def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__						):

				"""simple docstring"""


				lowercase__					:   int													=					dct.pop(lowerCamelCase__						)
				lowercase__					:   List[str]													=					val


@torch.no_grad()
def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__						):

				"""simple docstring"""


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

				# load state_dict of original model, remove and rename some keys
				lowercase__					:   str													=					torch.hub.load_state_dict_from_url(lowerCamelCase__					,			map_location="cpu"						)["state_dict"]
				lowercase__					:   Dict													=					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__					:   Optional[Any]													=					["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__					:   int													=					model.load_state_dict(lowerCamelCase__					,			strict=lowerCamelCase__						)
								assert missing_keys == ["mlm_score.decoder.bias"]
				else:
								model.load_state_dict(lowerCamelCase__						)

				# Define processor
				lowercase__					:   str													=					ViltImageProcessor(size=384						)
				lowercase__					:   Dict													=					BertTokenizer.from_pretrained("bert-base-uncased"						)
				lowercase__					:   str													=					ViltProcessor(lowerCamelCase__					,			lowerCamelCase__						)

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

				# Verify outputs
				if mlm_model:
								lowercase__					:   List[Any]													=					torch.Size([1, 11, 30_522]						)
								lowercase__					:   int													=					torch.tensor([-12.5061, -12.5123, -12.5174]						)
								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__					:   int													=					outputs.logits[0, 4, :].argmax(-1						).item()
								assert tokenizer.decode([predicted_id]						) == "cats"
				elif vqa_model:
								lowercase__					:   Any													=					torch.Size([1, 3_129]						)
								lowercase__					:   Tuple													=					torch.tensor([-15.9495, -18.1472, -10.3041]						)
								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__					:   Dict													=					outputs.logits.argmax(-1						).item()
								assert model.config.idalabel[predicted_idx] == "2"
				elif nlvr_model:
								lowercase__					:   Union[str, Any]													=					torch.Size([1, 2]						)
								lowercase__					:   str													=					torch.tensor([-2.8721, 2.1291]						)
								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__					  =  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__					  =  parser.parse_args()
				convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









import os
from argparse import ArgumentParser, Namespace

from ..data import SingleSentenceClassificationProcessor as Processor
from ..pipelines import TextClassificationPipeline
from ..utils import is_tf_available, is_torch_available, logging
from . import BaseTransformersCLICommand


if not is_tf_available() and not is_torch_available():
				raise RuntimeError('''At least one of PyTorch or TensorFlow 2.0+ should be installed to use CLI training''')

# TF training parameters
lowerCAmelCase__					  =  False
lowerCAmelCase__					  =  False


def  __lowerCamelCase (       lowerCamelCase__						):

				"""simple docstring"""


				return TrainCommand(lowerCamelCase__						)

class 				snake_case__(_UpperCamelCase      ):

					"""simple docstring"""






					@staticmethod
					def 			snake_case					(							SCREAMING_SNAKE_CASE :     ArgumentParser					):
									lowercase__					:   Optional[int]													=					parser.add_parser("train"				,       help="CLI tool to train a model on a task."					)

									train_parser.add_argument(
									    "--train_data"				,       type=SCREAMING_SNAKE_CASE				,       required=SCREAMING_SNAKE_CASE				,       help="path to train (and optionally evaluation) dataset as a csv with tab separated labels and sentences."				,       )
									train_parser.add_argument(
									    "--column_label"				,       type=SCREAMING_SNAKE_CASE				,       default=0				,       help="Column of the dataset csv file with example labels."					)
									train_parser.add_argument(
									    "--column_text"				,       type=SCREAMING_SNAKE_CASE				,       default=1				,       help="Column of the dataset csv file with example texts."					)
									train_parser.add_argument(
									    "--column_id"				,       type=SCREAMING_SNAKE_CASE				,       default=2				,       help="Column of the dataset csv file with example ids."					)
									train_parser.add_argument(
									    "--skip_first_row"				,       action="store_true"				,       help="Skip the first row of the csv file (headers)."					)

									train_parser.add_argument("--validation_data"				,       type=SCREAMING_SNAKE_CASE				,       default=""				,       help="path to validation dataset."					)
									train_parser.add_argument(
									    "--validation_split"				,       type=SCREAMING_SNAKE_CASE				,       default=0.1				,       help="if validation dataset is not provided, fraction of train dataset to use as validation dataset."				,       )

									train_parser.add_argument("--output"				,       type=SCREAMING_SNAKE_CASE				,       default="./"				,       help="path to saved the trained model."					)

									train_parser.add_argument(
									    "--task"				,       type=SCREAMING_SNAKE_CASE				,       default="text_classification"				,       help="Task to train the model on."					)
									train_parser.add_argument(
									    "--model"				,       type=SCREAMING_SNAKE_CASE				,       default="bert-base-uncased"				,       help="Model's name or path to stored model."					)
									train_parser.add_argument("--train_batch_size"				,       type=SCREAMING_SNAKE_CASE				,       default=32				,       help="Batch size for training."					)
									train_parser.add_argument("--valid_batch_size"				,       type=SCREAMING_SNAKE_CASE				,       default=64				,       help="Batch size for validation."					)
									train_parser.add_argument("--learning_rate"				,       type=SCREAMING_SNAKE_CASE				,       default=3E-5				,       help="Learning rate."					)
									train_parser.add_argument("--adam_epsilon"				,       type=SCREAMING_SNAKE_CASE				,       default=1E-0_8				,       help="Epsilon for Adam optimizer."					)
									train_parser.set_defaults(func=SCREAMING_SNAKE_CASE					)







					def __init__(							self :     int				,       SCREAMING_SNAKE_CASE :     Namespace					):
									lowercase__					:   int													=					logging.get_logger("transformers-cli/training"					)

									lowercase__					:   List[Any]													=					"tf" if is_tf_available() else "torch"

									os.makedirs(args.output				,       exist_ok=SCREAMING_SNAKE_CASE					)
									lowercase__					:   Optional[Any]													=					args.output

									lowercase__					:   Union[str, Any]													=					args.column_label
									lowercase__					:   Optional[int]													=					args.column_text
									lowercase__					:   Optional[int]													=					args.column_id

									self.logger.info(f"""Loading {args.task} pipeline for {args.model}"""					)
									if args.task == "text_classification":
													lowercase__					:   int													=					TextClassificationPipeline.from_pretrained(args.model					)
									elif args.task == "token_classification":
													raise NotImplementedError
									elif args.task == "question_answering":
													raise NotImplementedError

									self.logger.info(f"""Loading dataset from {args.train_data}"""					)
									lowercase__					:   List[str]													=					Processor.create_from_csv(
									    args.train_data				,       column_label=args.column_label				,       column_text=args.column_text				,       column_id=args.column_id				,       skip_first_row=args.skip_first_row				,       )
									lowercase__					:   Union[str, Any]													=					None
									if args.validation_data:
													self.logger.info(f"""Loading validation dataset from {args.validation_data}"""					)
													lowercase__					:   Optional[int]													=					Processor.create_from_csv(
													    args.validation_data				,       column_label=args.column_label				,       column_text=args.column_text				,       column_id=args.column_id				,       skip_first_row=args.skip_first_row				,       )

									lowercase__					:   Dict													=					args.validation_split
									lowercase__					:   List[str]													=					args.train_batch_size
									lowercase__					:   Any													=					args.valid_batch_size
									lowercase__					:   Optional[int]													=					args.learning_rate
									lowercase__					:   int													=					args.adam_epsilon







					def 			snake_case					(							self :     Dict					):
									if self.framework == "tf":
													return self.run_tf()
									return self.run_torch()







					def 			snake_case					(							self :     Union[str, Any]					):
									raise NotImplementedError







					def 			snake_case					(							self :     Union[str, Any]					):
									self.pipeline.fit(
									    self.train_dataset				,       validation_data=self.valid_dataset				,       validation_split=self.validation_split				,       learning_rate=self.learning_rate				,       adam_epsilon=self.adam_epsilon				,       train_batch_size=self.train_batch_size				,       valid_batch_size=self.valid_batch_size				,       )

									# Save trained pipeline
									self.pipeline.save_pretrained(self.output					)




