






from __future__ import annotations

import unittest

from transformers import XGLMConfig, XGLMTokenizer, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
							import tensorflow as tf

							from transformers.models.xglm.modeling_tf_xglm import (
							    TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
							    TFXGLMForCausalLM,
							    TFXGLMModel,
							)




@require_tf
class 				__A     :





				"""simple docstring"""







				UpperCamelCase__ :      int			 =XGLMConfig
				UpperCamelCase__ :      Optional[Any]			 ={}
				UpperCamelCase__ :      List[str]			 ="""gelu"""


				def __init__( self							,  lowerCamelCase__							,  lowerCamelCase__=14							,  lowerCamelCase__=7							,  lowerCamelCase__=True							,  lowerCamelCase__=True							,  lowerCamelCase__=True							,  lowerCamelCase__=99							,  lowerCamelCase__=32							,  lowerCamelCase__=2							,  lowerCamelCase__=4							,  lowerCamelCase__=37							,  lowerCamelCase__="gelu"							,  lowerCamelCase__=0.1							,  lowerCamelCase__=0.1							,  lowerCamelCase__=512							,  lowerCamelCase__=0.02							,  ):




									"""simple docstring"""

									__UpperCamelCase      :					Tuple       		=parent
									__UpperCamelCase      :					List[str]       		=batch_size
									__UpperCamelCase      :					str       		=seq_length
									__UpperCamelCase      :					Dict       		=is_training
									__UpperCamelCase      :					Tuple       		=use_input_mask
									__UpperCamelCase      :					List[Any]       		=use_labels
									__UpperCamelCase      :					Any       		=vocab_size
									__UpperCamelCase      :					List[Any]       		=d_model
									__UpperCamelCase      :					Optional[int]       		=num_hidden_layers
									__UpperCamelCase      :					List[str]       		=num_attention_heads
									__UpperCamelCase      :					Optional[int]       		=ffn_dim
									__UpperCamelCase      :					str       		=activation_function
									__UpperCamelCase      :					Any       		=activation_dropout
									__UpperCamelCase      :					Optional[int]       		=attention_dropout
									__UpperCamelCase      :					Optional[int]       		=max_position_embeddings
									__UpperCamelCase      :					Any       		=initializer_range
									__UpperCamelCase      :					Dict       		=None
									__UpperCamelCase      :					Optional[int]       		=0
									__UpperCamelCase      :					Optional[Any]       		=2
									__UpperCamelCase      :					str       		=1


				def       __lowercase							( self    ):




									"""simple docstring"""

									return XGLMConfig.from_pretrained('facebook/xglm-564M'    )


				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					List[Any]       		=tf.clip_by_value(
									    ids_tensor([self.batch_size, self.seq_length]							,  self.vocab_size    )							,  clip_value_min=0							,  clip_value_max=3    )

									__UpperCamelCase      :					Union[str, Any]       		=None
									if self.use_input_mask:
														__UpperCamelCase      :					Dict       		=random_attention_mask([self.batch_size, self.seq_length]    )

									__UpperCamelCase      :					Any       		=self.get_config()

									__UpperCamelCase      :					Optional[Any]       		=floats_tensor([self.num_hidden_layers, self.num_attention_heads]							,  2    )

									return (
									    config,
									    input_ids,
									    input_mask,
									    head_mask,
									)


				def       __lowercase							( self    ):




									"""simple docstring"""

									return XGLMConfig(
									    vocab_size=self.vocab_size							,  d_model=self.hidden_size							,  num_layers=self.num_hidden_layers							,  attention_heads=self.num_attention_heads							,  ffn_dim=self.ffn_dim							,  activation_function=self.activation_function							,  activation_dropout=self.activation_dropout							,  attention_dropout=self.attention_dropout							,  max_position_embeddings=self.max_position_embeddings							,  initializer_range=self.initializer_range							,  use_cache=lowerCamelCase__							,  bos_token_id=self.bos_token_id							,  eos_token_id=self.eos_token_id							,  pad_token_id=self.pad_token_id							,  return_dict=lowerCamelCase__							,  )







				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					List[str]       		=self.prepare_config_and_inputs()

									(
									    (
									    __UpperCamelCase
									)		,				(
									    __UpperCamelCase
									)		,				(
									    __UpperCamelCase
									)		,				(
									    __UpperCamelCase
									)		,				
									)      :					int       		=config_and_inputs

									__UpperCamelCase      :					Optional[Any]       		={
									    'input_ids': input_ids,
									    'head_mask': head_mask,
									}

									return config, inputs_dict




@require_tf
class 				__A     (			a      ,     a      ,     unittest.TestCase ):





				"""simple docstring"""







				UpperCamelCase__ :      Union[str, Any]			 =(TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
				UpperCamelCase__ :      str			 =(TFXGLMForCausalLM,) if is_tf_available() else ()
				UpperCamelCase__ :      Optional[Any]			 =(
				    {"""feature-extraction""": TFXGLMModel, """text-generation""": TFXGLMForCausalLM} if is_tf_available() else {}
				)
				UpperCamelCase__ :      Tuple			 =False
				UpperCamelCase__ :      Tuple			 =False
				UpperCamelCase__ :      Optional[Any]			 =False


				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Tuple       		=TFXGLMModelTester(self    )
									__UpperCamelCase      :					Dict       		=ConfigTester(self							,  config_class=lowerCamelCase__							,  n_embd=37    )


				def       __lowercase							( self    ):




									"""simple docstring"""

									self.config_tester.run_common_tests()


				@slow
				def       __lowercase							( self    ):




									"""simple docstring"""

									for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														__UpperCamelCase      :					Optional[Any]       		=TFXGLMModel.from_pretrained(lowerCamelCase__    )
														self.assertIsNotNone(lowerCamelCase__    )







				@unittest.skip(reason='Currently, model embeddings are going to undergo a major refactor.'    )
				def       __lowercase							( self    ):




									"""simple docstring"""

									super().test_resize_token_embeddings()




@require_tf
class 				__A     (			unittest.TestCase ):





				"""simple docstring"""







				@slow
				def       __lowercase							( self							,  lowerCamelCase__=True    ):




									"""simple docstring"""

									__UpperCamelCase      :					int       		=TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M'    )
									__UpperCamelCase      :					List[str]       		=tf.convert_to_tensor([[2, 268, 9865]]							,  dtype=tf.intaa    )  # The dog
									# </s> The dog is a very friendly dog. He is very affectionate and loves to play with other
									# fmt: off
									__UpperCamelCase      :					str       		=[2, 268, 9865, 67, 11, 1988, 57252, 9865, 5, 984, 67, 1988, 213838, 1658, 53, 70446, 33, 6657, 278, 1581]
									# fmt: on
									__UpperCamelCase      :					Optional[Any]       		=model.generate(lowerCamelCase__							,  do_sample=lowerCamelCase__							,  num_beams=1    )
									if verify_outputs:
														self.assertListEqual(output_ids[0].numpy().tolist()							,  lowerCamelCase__    )


				@slow
				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					List[str]       		=XGLMTokenizer.from_pretrained('facebook/xglm-564M'    )
									__UpperCamelCase      :					Union[str, Any]       		=TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M'    )

									tf.random.set_seed(0    )
									__UpperCamelCase      :					str       		=tokenizer('Today is a nice day and'							,  return_tensors='tf'    )
									__UpperCamelCase      :					Union[str, Any]       		=tokenized.input_ids
									# forces the generation to happen on CPU, to avoid GPU-related quirks (and assure same output regardless of the available devices)
									with tf.device(':/CPU:0'    ):
														__UpperCamelCase      :					Any       		=model.generate(lowerCamelCase__							,  do_sample=lowerCamelCase__							,  seed=[7, 0]    )
									__UpperCamelCase      :					Tuple       		=tokenizer.decode(output_ids[0]							,  skip_special_tokens=lowerCamelCase__    )

									__UpperCamelCase      :					List[Any]       		=(
									    'Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due'
									)
									self.assertEqual(lowerCamelCase__							,  lowerCamelCase__    )







				@slow
				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Tuple       		=TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M'    )
									__UpperCamelCase      :					Optional[Any]       		=XGLMTokenizer.from_pretrained('facebook/xglm-564M'    )

									__UpperCamelCase      :					Optional[Any]       		='left'

									# use different length sentences to test batching
									__UpperCamelCase      :					Optional[int]       		=[
									    'This is an extremelly long sentence that only exists to test the ability of the model to cope with '
									    'left-padding, such as in batched generation. The output for the sequence below should be the same '
									    'regardless of whether left padding is applied or not. When',
									    'Hello, my dog is a little',
									]

									__UpperCamelCase      :					List[Any]       		=tokenizer(lowerCamelCase__							,  return_tensors='tf'							,  padding=lowerCamelCase__    )
									__UpperCamelCase      :					Union[str, Any]       		=inputs['input_ids']

									__UpperCamelCase      :					Dict       		=model.generate(input_ids=lowerCamelCase__							,  attention_mask=inputs['attention_mask']							,  max_new_tokens=12    )

									__UpperCamelCase      :					List[Any]       		=tokenizer(sentences[0]							,  return_tensors='tf'    ).input_ids
									__UpperCamelCase      :					Dict       		=model.generate(input_ids=lowerCamelCase__							,  max_new_tokens=12    )

									__UpperCamelCase      :					Any       		=tokenizer(sentences[1]							,  return_tensors='tf'    ).input_ids
									__UpperCamelCase      :					Optional[Any]       		=model.generate(input_ids=lowerCamelCase__							,  max_new_tokens=12    )

									__UpperCamelCase      :					Optional[int]       		=tokenizer.batch_decode(lowerCamelCase__							,  skip_special_tokens=lowerCamelCase__    )
									__UpperCamelCase      :					Union[str, Any]       		=tokenizer.decode(output_non_padded[0]							,  skip_special_tokens=lowerCamelCase__    )
									__UpperCamelCase      :					int       		=tokenizer.decode(output_padded[0]							,  skip_special_tokens=lowerCamelCase__    )

									__UpperCamelCase      :					Any       		=[
									    'This is an extremelly long sentence that only exists to test the ability of the model to cope with '
									    'left-padding, such as in batched generation. The output for the sequence below should be the same '
									    'regardless of whether left padding is applied or not. When left padding is applied, the sequence will be '
									    'a single',
									    'Hello, my dog is a little bit of a shy one, but he is very friendly',
									]
									self.assertListEqual(lowerCamelCase__							,  lowerCamelCase__    )
									self.assertListEqual(lowerCamelCase__							,  [non_padded_sentence, padded_sentence]    )





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

A     :   Tuple			    =			[
    ("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"),
]

A     :   Optional[Any]			    =			[
    "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 _lowerCamelCase	(		_UpperCamelCase      ):

			'''simple docstring'''




			__lowerCAmelCase				    =			torch.load(_UpperCamelCase    , map_location="cpu"      )
			return sd


def _lowerCamelCase	(		_UpperCamelCase    , _UpperCamelCase    , _UpperCamelCase=rename_keys_prefix      ):

			'''simple docstring'''




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


@torch.no_grad()
def _lowerCamelCase	(		_UpperCamelCase    , _UpperCamelCase      ):

			'''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:
						__lowerCAmelCase				    =			"pretraining"
						if "vcr" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 512}
						elif "vqa_advanced" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 2048}
						elif "vqa" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 2048}
						elif "nlvr" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 1024}
						else:
									raise NotImplementedError(f"No implementation found for `{checkpoint_path}`."      )
			else:
						if "vcr" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 512}
									__lowerCAmelCase				    =			"multichoice"
						elif "vqa_advanced" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 2048}
									__lowerCAmelCase				    =			"vqa_advanced"
						elif "vqa" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 2048, "num_labels": 3129}
									__lowerCAmelCase				    =			"vqa"
						elif "nlvr" in checkpoint_path:
									__lowerCAmelCase				    =			{
									    "visual_embedding_dim": 1024,
									    "num_labels": 2,
									}
									__lowerCAmelCase				    =			"nlvr"

			__lowerCAmelCase				    =			VisualBertConfig(**_UpperCamelCase      )

			# Load State Dict
			__lowerCAmelCase				    =			load_state_dict(_UpperCamelCase      )

			__lowerCAmelCase				    =			get_new_dict(_UpperCamelCase    , _UpperCamelCase      )

			if model_type == "pretraining":
						__lowerCAmelCase				    =			VisualBertForPreTraining(_UpperCamelCase      )
			elif model_type == "vqa":
						__lowerCAmelCase				    =			VisualBertForQuestionAnswering(_UpperCamelCase      )
			elif model_type == "nlvr":
						__lowerCAmelCase				    =			VisualBertForVisualReasoning(_UpperCamelCase      )
			elif model_type == "multichoice":
						__lowerCAmelCase				    =			VisualBertForMultipleChoice(_UpperCamelCase      )

			model.load_state_dict(_UpperCamelCase      )
			# Save Checkpoints
			Path(_UpperCamelCase      ).mkdir(exist_ok=_UpperCamelCase      )
			model.save_pretrained(_UpperCamelCase      )


if __name__ == "__main__":
	A     :   Union[str, Any]			    =			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.")
	A     :   Optional[int]			    =			parser.parse_args()
	convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)




