


"""simple docstring"""


import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel
class   _UpperCamelCase  :
							'''simple docstring'''







							def __init__(     self    ,   __a    ,   __a=13    ,   __a=7    ,   __a=True    ,   __a=True    ,   __a=False    ,   __a=True    ,   __a=99    ,   __a=32    ,   __a=5    ,   __a=4    ,   __a=37    ,   __a="gelu"    ,   __a=0.1    ,   __a=0.1    ,   __a=5_12    ,   __a=16    ,   __a=2    ,   __a=0.0_2    ,   __a=3    ,   __a=4    ,   __a=None    ,   ):
										__lowerCAmelCase				    =			parent
										__lowerCAmelCase				    =			batch_size
										__lowerCAmelCase				    =			seq_length
										__lowerCAmelCase				    =			is_training
										__lowerCAmelCase				    =			use_input_mask
										__lowerCAmelCase				    =			use_token_type_ids
										__lowerCAmelCase				    =			use_labels
										__lowerCAmelCase				    =			vocab_size
										__lowerCAmelCase				    =			hidden_size
										__lowerCAmelCase				    =			num_hidden_layers
										__lowerCAmelCase				    =			num_attention_heads
										__lowerCAmelCase				    =			intermediate_size
										__lowerCAmelCase				    =			hidden_act
										__lowerCAmelCase				    =			hidden_dropout_prob
										__lowerCAmelCase				    =			attention_probs_dropout_prob
										__lowerCAmelCase				    =			max_position_embeddings
										__lowerCAmelCase				    =			type_vocab_size
										__lowerCAmelCase				    =			type_sequence_label_size
										__lowerCAmelCase				    =			initializer_range
										__lowerCAmelCase				    =			num_labels
										__lowerCAmelCase				    =			num_choices
										__lowerCAmelCase				    =			scope

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			ids_tensor([self.batch_size, self.seq_length]    ,   self.vocab_size	)

										__lowerCAmelCase				    =			None
										if self.use_input_mask:
													__lowerCAmelCase				    =			random_attention_mask([self.batch_size, self.seq_length]	)

										__lowerCAmelCase				    =			None
										if self.use_token_type_ids:
													__lowerCAmelCase				    =			ids_tensor([self.batch_size, self.seq_length]    ,   self.type_vocab_size	)

										__lowerCAmelCase				    =			None
										__lowerCAmelCase				    =			None
										__lowerCAmelCase				    =			None
										if self.use_labels:
													__lowerCAmelCase				    =			ids_tensor([self.batch_size]    ,   self.type_sequence_label_size	)
													__lowerCAmelCase				    =			ids_tensor([self.batch_size, self.seq_length]    ,   self.num_labels	)
													__lowerCAmelCase				    =			ids_tensor([self.batch_size]    ,   self.num_choices	)

										__lowerCAmelCase				    =			self.get_config()

										return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

							def 			snake_case				(     self	):
										return OpenLlamaConfig(
										    vocab_size=self.vocab_size    ,   hidden_size=self.hidden_size    ,   num_hidden_layers=self.num_hidden_layers    ,   num_attention_heads=self.num_attention_heads    ,   intermediate_size=self.intermediate_size    ,   hidden_act=self.hidden_act    ,   hidden_dropout_prob=self.hidden_dropout_prob    ,   attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,   max_position_embeddings=self.max_position_embeddings    ,   type_vocab_size=self.type_vocab_size    ,   is_decoder=__a    ,   initializer_range=self.initializer_range    ,   use_stable_embedding=__a    ,   )

							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a	):
										__lowerCAmelCase				    =			OpenLlamaModel(config=__a	)
										model.to(__a	)
										model.eval()
										__lowerCAmelCase				    =			model(__a    ,   attention_mask=__a	)
										__lowerCAmelCase				    =			model(__a	)
										self.parent.assertEqual(result.last_hidden_state.shape    ,   (self.batch_size, self.seq_length, self.hidden_size)	)

							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   ):
										__lowerCAmelCase				    =			True
										__lowerCAmelCase				    =			OpenLlamaModel(__a	)
										model.to(__a	)
										model.eval()
										__lowerCAmelCase				    =			model(
										    __a    ,   attention_mask=__a    ,   encoder_hidden_states=__a    ,   encoder_attention_mask=__a    ,   )
										__lowerCAmelCase				    =			model(
										    __a    ,   attention_mask=__a    ,   encoder_hidden_states=__a    ,   )
										__lowerCAmelCase				    =			model(__a    ,   attention_mask=__a	)
										self.parent.assertEqual(result.last_hidden_state.shape    ,   (self.batch_size, self.seq_length, self.hidden_size)	)

							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   ):
										__lowerCAmelCase				    =			OpenLlamaForCausalLM(config=__a	)
										model.to(__a	)
										model.eval()
										__lowerCAmelCase				    =			model(__a    ,   attention_mask=__a    ,   labels=__a	)
										self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.seq_length, self.vocab_size)	)

							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   __a    ,   ):
										__lowerCAmelCase				    =			True
										__lowerCAmelCase				    =			True
										__lowerCAmelCase				    =			OpenLlamaForCausalLM(config=__a	)
										model.to(__a	)
										model.eval()

										# first forward pass
										__lowerCAmelCase				    =			model(
										    __a    ,   attention_mask=__a    ,   encoder_hidden_states=__a    ,   encoder_attention_mask=__a    ,   use_cache=__a    ,   )
										__lowerCAmelCase				    =			outputs.past_key_values

										# create hypothetical multiple next token and extent to next_input_ids
										__lowerCAmelCase				    =			ids_tensor((self.batch_size, 3)    ,   config.vocab_size	)
										__lowerCAmelCase				    =			ids_tensor((self.batch_size, 3)    ,   vocab_size=2	)

										# append to next input_ids and
										__lowerCAmelCase				    =			torch.cat([input_ids, next_tokens]    ,   dim=-1	)
										__lowerCAmelCase				    =			torch.cat([input_mask, next_mask]    ,   dim=-1	)

										__lowerCAmelCase				    =			model(
										    __a    ,   attention_mask=__a    ,   encoder_hidden_states=__a    ,   encoder_attention_mask=__a    ,   output_hidden_states=__a    ,   )["hidden_states"][0]
										__lowerCAmelCase				    =			model(
										    __a    ,   attention_mask=__a    ,   encoder_hidden_states=__a    ,   encoder_attention_mask=__a    ,   past_key_values=__a    ,   output_hidden_states=__a    ,   )["hidden_states"][0]

										# select random slice
										__lowerCAmelCase				    =			ids_tensor((1,)    ,   output_from_past.shape[-1]	).item()
										__lowerCAmelCase				    =			output_from_no_past[:, -3:, random_slice_idx].detach()
										__lowerCAmelCase				    =			output_from_past[:, :, random_slice_idx].detach()

										self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]	)

										# test that outputs are equal for slice
										self.parent.assertTrue(torch.allclose(__a    ,   __a    ,   atol=1e-3	)	)



							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			self.prepare_config_and_inputs()
										(
										    (
										    __lowerCAmelCase
										)    ,	(
										    __lowerCAmelCase
										)    ,	(
										    __lowerCAmelCase
										)    ,	(
										    __lowerCAmelCase
										)    ,	(
										    __lowerCAmelCase
										)    ,	(
										    __lowerCAmelCase
										)    ,	(
										    __lowerCAmelCase
										)    ,	
										)				    =			config_and_inputs
										__lowerCAmelCase				    =			{"input_ids": input_ids, "attention_mask": input_mask}
										return config, inputs_dict
@require_torch
class   _UpperCamelCase  (   lowerCAmelCase__       ,lowerCAmelCase__       ,lowerCAmelCase__       ,unittest.TestCase		):
							'''simple docstring'''







							__UpperCAmelCase						:    Union[str, Any]				      =(
							    (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
							)
							__UpperCAmelCase						:    List[str]				      =(OpenLlamaForCausalLM,) if is_torch_available() else ()
							__UpperCAmelCase						:    Any				      =(
							    {
							        """feature-extraction""": OpenLlamaModel,
							        """text-classification""": OpenLlamaForSequenceClassification,
							        """text-generation""": OpenLlamaForCausalLM,
							        """zero-shot""": OpenLlamaForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							__UpperCAmelCase						:    int				      =False
							__UpperCAmelCase						:    List[Any]				      =False

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			OpenLlamaModelTester(self	)
										__lowerCAmelCase				    =			ConfigTester(self    ,   config_class=__a    ,   hidden_size=37	)

							def 			snake_case				(     self	):
										self.config_tester.run_common_tests()

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__a	)

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
													__lowerCAmelCase				    =			type
													self.model_tester.create_and_check_model(*__a	)

							def 			snake_case				(     self	):
										__lowerCAmelCase    ,	__lowerCAmelCase				    =			self.model_tester.prepare_config_and_inputs_for_common()
										__lowerCAmelCase				    =			3
										__lowerCAmelCase				    =			input_dict["input_ids"]
										__lowerCAmelCase				    =			input_ids.ne(1	).to(__a	)
										__lowerCAmelCase				    =			ids_tensor([self.model_tester.batch_size]    ,   self.model_tester.type_sequence_label_size	)
										__lowerCAmelCase				    =			OpenLlamaForSequenceClassification(__a	)
										model.to(__a	)
										model.eval()
										__lowerCAmelCase				    =			model(__a    ,   attention_mask=__a    ,   labels=__a	)
										self.assertEqual(result.logits.shape    ,   (self.model_tester.batch_size, self.model_tester.num_labels)	)

							def 			snake_case				(     self	):
										__lowerCAmelCase    ,	__lowerCAmelCase				    =			self.model_tester.prepare_config_and_inputs_for_common()
										__lowerCAmelCase				    =			3
										__lowerCAmelCase				    =			"single_label_classification"
										__lowerCAmelCase				    =			input_dict["input_ids"]
										__lowerCAmelCase				    =			input_ids.ne(1	).to(__a	)
										__lowerCAmelCase				    =			ids_tensor([self.model_tester.batch_size]    ,   self.model_tester.type_sequence_label_size	)
										__lowerCAmelCase				    =			OpenLlamaForSequenceClassification(__a	)
										model.to(__a	)
										model.eval()
										__lowerCAmelCase				    =			model(__a    ,   attention_mask=__a    ,   labels=__a	)
										self.assertEqual(result.logits.shape    ,   (self.model_tester.batch_size, self.model_tester.num_labels)	)

							def 			snake_case				(     self	):
										__lowerCAmelCase    ,	__lowerCAmelCase				    =			self.model_tester.prepare_config_and_inputs_for_common()
										__lowerCAmelCase				    =			3
										__lowerCAmelCase				    =			"multi_label_classification"
										__lowerCAmelCase				    =			input_dict["input_ids"]
										__lowerCAmelCase				    =			input_ids.ne(1	).to(__a	)
										__lowerCAmelCase				    =			ids_tensor(
										    [self.model_tester.batch_size, config.num_labels]    ,   self.model_tester.type_sequence_label_size	).to(torch.float	)
										__lowerCAmelCase				    =			OpenLlamaForSequenceClassification(__a	)
										model.to(__a	)
										model.eval()
										__lowerCAmelCase				    =			model(__a    ,   attention_mask=__a    ,   labels=__a	)
										self.assertEqual(result.logits.shape    ,   (self.model_tester.batch_size, self.model_tester.num_labels)	)

							@unittest.skip("Open-Llama buffers include complex numbers, which breaks this test"	)
							def 			snake_case				(     self	):
										pass



							@parameterized.expand([("linear",), ("dynamic",)]	)
							def 			snake_case				(     self    ,   __a	):
										__lowerCAmelCase    ,	__lowerCAmelCase				    =			self.model_tester.prepare_config_and_inputs_for_common()
										__lowerCAmelCase				    =			ids_tensor([1, 10]    ,   config.vocab_size	)
										__lowerCAmelCase				    =			ids_tensor([1, int(config.max_position_embeddings * 1.5	)]    ,   config.vocab_size	)

										set_seed(42	)  # Fixed seed at init time so the two models get the same random weights
										__lowerCAmelCase				    =			OpenLlamaModel(__a	)
										original_model.to(__a	)
										original_model.eval()
										__lowerCAmelCase				    =			original_model(__a	).last_hidden_state
										__lowerCAmelCase				    =			original_model(__a	).last_hidden_state

										set_seed(42	)  # Fixed seed at init time so the two models get the same random weights
										__lowerCAmelCase				    =			{"type": scaling_type, "factor": 1_0.0}
										__lowerCAmelCase				    =			OpenLlamaModel(__a	)
										scaled_model.to(__a	)
										scaled_model.eval()
										__lowerCAmelCase				    =			scaled_model(__a	).last_hidden_state
										__lowerCAmelCase				    =			scaled_model(__a	).last_hidden_state

										# Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
										# maximum sequence length, so the outputs for the short input should match.
										if scaling_type == "dynamic":
													self.assertTrue(torch.allclose(__a    ,   __a    ,   atol=1e-5	)	)
										else:
													self.assertFalse(torch.allclose(__a    ,   __a    ,   atol=1e-5	)	)

										# The output should be different for long inputs
										self.assertFalse(torch.allclose(__a    ,   __a    ,   atol=1e-5	)	)








"""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)




