





import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from ...utils import logging
from ..auto import CONFIG_MAPPING


UpperCAmelCase_ 							=       logging.get_logger(__name__)

UpperCAmelCase_ 							=       {
    """Salesforce/instruct-blip-flan-t5""": """https://huggingface.co/Salesforce/instruct-blip-flan-t5/resolve/main/config.json""",
}
class      UpperCamelCase_						(					_lowerCamelCase   ):
					lowerCAmelCase_							=		'''instructblip_vision_model'''




					def __init__(						self			,  lowerCAmelCase_=1408			,  lowerCAmelCase_=6144			,  lowerCAmelCase_=39			,  lowerCAmelCase_=16			,  lowerCAmelCase_=224			,  lowerCAmelCase_=14			,  lowerCAmelCase_="gelu"			,  lowerCAmelCase_=1E-6			,  lowerCAmelCase_=0.0			,  lowerCAmelCase_=1E-10			,  lowerCAmelCase_=True			,  **lowerCAmelCase_			,  )						->							Tuple:
										super().__init__(**lowerCAmelCase_					)

										_snake_case													=			hidden_size
										_snake_case													=			intermediate_size
										_snake_case													=			num_hidden_layers
										_snake_case													=			num_attention_heads
										_snake_case													=			patch_size
										_snake_case													=			image_size
										_snake_case													=			initializer_range
										_snake_case													=			attention_dropout
										_snake_case													=			layer_norm_eps
										_snake_case													=			hidden_act
										_snake_case													=			qkv_bias






					@classmethod
					def 		lowerCAmelCase     (						cls			,  lowerCAmelCase_			,  **lowerCAmelCase_					)						->							"PretrainedConfig":
										cls._set_token_in_kwargs(lowerCAmelCase_					)

										_snake_case							,       _snake_case													=			cls.get_config_dict(lowerCAmelCase_			,  **lowerCAmelCase_					)

										# get the vision config dict if we are loading from InstructBlipConfig
										if config_dict.get('model_type'					) == "instructblip":
															_snake_case													=			config_dict['vision_config']

										if "model_type" in config_dict and hasattr(cls			,  'model_type'					) and config_dict["model_type"] != cls.model_type:
															logger.warning(
															    F'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
															    F'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''					)

										return cls.from_dict(lowerCAmelCase_			,  **lowerCAmelCase_					)





class      UpperCamelCase_						(					_lowerCamelCase   ):
					lowerCAmelCase_							=		'''instructblip_qformer'''




					def __init__(						self			,  lowerCAmelCase_=3_0522			,  lowerCAmelCase_=768			,  lowerCAmelCase_=12			,  lowerCAmelCase_=12			,  lowerCAmelCase_=3072			,  lowerCAmelCase_="gelu"			,  lowerCAmelCase_=0.1			,  lowerCAmelCase_=0.1			,  lowerCAmelCase_=512			,  lowerCAmelCase_=0.02			,  lowerCAmelCase_=1E-12			,  lowerCAmelCase_=0			,  lowerCAmelCase_="absolute"			,  lowerCAmelCase_=2			,  lowerCAmelCase_=1408			,  **lowerCAmelCase_			,  )						->							List[str]:
										super().__init__(pad_token_id=lowerCAmelCase_			,  **lowerCAmelCase_					)

										_snake_case													=			vocab_size
										_snake_case													=			hidden_size
										_snake_case													=			num_hidden_layers
										_snake_case													=			num_attention_heads
										_snake_case													=			hidden_act
										_snake_case													=			intermediate_size
										_snake_case													=			hidden_dropout_prob
										_snake_case													=			attention_probs_dropout_prob
										_snake_case													=			max_position_embeddings
										_snake_case													=			initializer_range
										_snake_case													=			layer_norm_eps
										_snake_case													=			position_embedding_type
										_snake_case													=			cross_attention_frequency
										_snake_case													=			encoder_hidden_size






					@classmethod
					def 		lowerCAmelCase     (						cls			,  lowerCAmelCase_			,  **lowerCAmelCase_					)						->							"PretrainedConfig":
										cls._set_token_in_kwargs(lowerCAmelCase_					)

										_snake_case							,       _snake_case													=			cls.get_config_dict(lowerCAmelCase_			,  **lowerCAmelCase_					)

										# get the qformer config dict if we are loading from InstructBlipConfig
										if config_dict.get('model_type'					) == "instructblip":
															_snake_case													=			config_dict['qformer_config']

										if "model_type" in config_dict and hasattr(cls			,  'model_type'					) and config_dict["model_type"] != cls.model_type:
															logger.warning(
															    F'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
															    F'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''					)

										return cls.from_dict(lowerCAmelCase_			,  **lowerCAmelCase_					)





class      UpperCamelCase_						(					_lowerCamelCase   ):
					lowerCAmelCase_							=		'''instructblip'''
					lowerCAmelCase_							=		True




					def __init__(						self			,  lowerCAmelCase_=None			,  lowerCAmelCase_=None			,  lowerCAmelCase_=None			,  lowerCAmelCase_=32			,  **lowerCAmelCase_					)						->							str:
										super().__init__(**lowerCAmelCase_					)

										if vision_config is None:
															_snake_case													=			{}
															logger.info('vision_config is None. initializing the InstructBlipVisionConfig with default values.'					)

										if qformer_config is None:
															_snake_case													=			{}
															logger.info('qformer_config is None. Initializing the InstructBlipQFormerConfig with default values.'					)

										if text_config is None:
															_snake_case													=			{}
															logger.info('text_config is None. Initializing the text config with default values (`OPTConfig`).'					)

										_snake_case													=			InstructBlipVisionConfig(**lowerCAmelCase_					)
										_snake_case													=			InstructBlipQFormerConfig(**lowerCAmelCase_					)
										_snake_case													=			text_config['model_type'] if 'model_type' in text_config else 'opt'
										_snake_case													=			CONFIG_MAPPING[text_model_type](**lowerCAmelCase_					)

										_snake_case													=			self.text_config.tie_word_embeddings
										_snake_case													=			self.text_config.is_encoder_decoder

										_snake_case													=			num_query_tokens
										_snake_case													=			self.vision_config.hidden_size
										_snake_case													=			self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
										_snake_case													=			1.0
										_snake_case													=			0.02




					@classmethod
					def 		lowerCAmelCase     (						cls			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  **lowerCAmelCase_			,  )						->							Tuple:
										return cls(
										    vision_config=vision_config.to_dict()			,  qformer_config=qformer_config.to_dict()			,  text_config=text_config.to_dict()			,  **lowerCAmelCase_			,  )






					def 		lowerCAmelCase     (						self					)						->							Optional[int]:
										_snake_case													=			copy.deepcopy(self.__dict__					)
										_snake_case													=			self.vision_config.to_dict()
										_snake_case													=			self.qformer_config.to_dict()
										_snake_case													=			self.text_config.to_dict()
										_snake_case													=			self.__class__.model_type
										return output










import argparse

from torch import nn

# transformers_old should correspond to branch `save_old_prophetnet_model_structure` here
# original prophetnet_checkpoints are saved under `patrickvonplaten/..._old` respectively
from transformers_old.modeling_prophetnet import (
    ProphetNetForConditionalGeneration as ProphetNetForConditionalGenerationOld,
)
from transformers_old.modeling_xlm_prophetnet import (
    XLMProphetNetForConditionalGeneration as XLMProphetNetForConditionalGenerationOld,
)

from transformers import ProphetNetForConditionalGeneration, XLMProphetNetForConditionalGeneration, logging


UpperCAmelCase_ 							=       logging.get_logger(__name__)
logging.set_verbosity_info()





def   lowerCamelCase__       (       UpperCamelCase__     :    str    ,      UpperCamelCase__     :    str  )     ->	Union[str, Any]:



					'''simple docstring'''




					if "xprophetnet" in prophetnet_checkpoint_path:
										_snake_case													=			XLMProphetNetForConditionalGenerationOld.from_pretrained(UpperCamelCase__  )
										_snake_case							,       _snake_case													=			XLMProphetNetForConditionalGeneration.from_pretrained(
										    UpperCamelCase__    ,      output_loading_info=UpperCamelCase__  )
					else:
										_snake_case													=			ProphetNetForConditionalGenerationOld.from_pretrained(UpperCamelCase__  )
										_snake_case							,       _snake_case													=			ProphetNetForConditionalGeneration.from_pretrained(
										    UpperCamelCase__    ,      output_loading_info=UpperCamelCase__  )

					_snake_case													=			['key_proj', 'value_proj', 'query_proj']

					_snake_case													=			{
					    'self_attn': 'ngram_self_attn',
					    'cross_attn': 'encoder_attn',
					    'cross_attn_layer_norm': 'encoder_attn_layer_norm',
					    'feed_forward_layer_norm': 'final_layer_norm',
					    'feed_forward': '',
					    'intermediate': 'fc1',
					    'output': 'fc2',
					    'key_proj': 'k_proj',
					    'query_proj': 'q_proj',
					    'value_proj': 'v_proj',
					    'word_embeddings': 'embed_tokens',
					    'embeddings_layer_norm': 'emb_layer_norm',
					    'relative_pos_embeddings': 'relative_linear',
					    'ngram_embeddings': 'ngram_input_embed',
					    'position_embeddings': 'embed_positions',
					}

					for key in loading_info["missing_keys"]:
										_snake_case													=			key.split('.'  )

										if attributes[0] == "lm_head":
															_snake_case													=			prophet
															_snake_case													=			prophet_old
										else:
															_snake_case													=			prophet.prophetnet
															_snake_case													=			prophet_old.model

										_snake_case													=			False
										for attribute in attributes:
															if attribute in mapping:
																				_snake_case													=			mapping[attribute]
																				if not hasattr(UpperCamelCase__    ,      UpperCamelCase__  ) and len(UpperCamelCase__  ) > 0:
																									_snake_case													=			attribute
															elif hasattr(UpperCamelCase__    ,      UpperCamelCase__  ):
																				_snake_case													=			attribute

															if attribute == "weight":
																				assert old_model.weight.shape == model.weight.shape, "Shapes have to match!"
																				_snake_case													=			old_model.weight
																				logger.info(F'''{attribute} is initialized.'''  )
																				_snake_case													=			True
																				break
															elif attribute == "bias":
																				assert old_model.bias.shape == model.bias.shape, "Shapes have to match!"
																				_snake_case													=			old_model.bias
																				logger.info(F'''{attribute} is initialized'''  )
																				_snake_case													=			True
																				break
															elif attribute in special_keys and hasattr(UpperCamelCase__    ,      'in_proj_weight'  ):
																				_snake_case													=			old_model.in_proj_weight.shape[0] // 3
																				_snake_case													=			getattr(UpperCamelCase__    ,      UpperCamelCase__  )
																				param.weight.shape == old_model.in_proj_weight[:embed_dim, :].shape, "Shapes have to match"
																				param.bias.shape == old_model.in_proj_bias[:embed_dim].shape, "Shapes have to match"
																				if attribute == "query_proj":
																									_snake_case													=			nn.Parameter(old_model.in_proj_weight[:embed_dim, :]  )
																									_snake_case													=			nn.Parameter(old_model.in_proj_bias[:embed_dim]  )

																				elif attribute == "key_proj":
																									_snake_case													=			nn.Parameter(old_model.in_proj_weight[embed_dim : 2 * embed_dim, :]  )
																									_snake_case													=			nn.Parameter(old_model.in_proj_bias[embed_dim : 2 * embed_dim]  )
																				elif attribute == "value_proj":
																									_snake_case													=			nn.Parameter(old_model.in_proj_weight[2 * embed_dim :, :]  )
																									_snake_case													=			nn.Parameter(old_model.in_proj_bias[2 * embed_dim :]  )
																				_snake_case													=			True
																				break
															elif attribute == "position_embeddings":
																				assert (
																				    model.position_embeddings.weight.shape[-1] == old_model.embed_positions.weight.shape[-1]
																				), "Hidden size has to match"
																				assert model.position_embeddings.weight.shape[0] == 512, "We want 512 position_embeddings."
																				_snake_case													=			nn.Parameter(old_model.embed_positions.weight[:512, :]  )
																				_snake_case													=			True
																				break

															if attribute.isdigit():
																				_snake_case													=			model[int(UpperCamelCase__  )]
																				_snake_case													=			old_model[int(UpperCamelCase__  )]
															else:
																				_snake_case													=			getattr(UpperCamelCase__    ,      UpperCamelCase__  )

																				if old_attribute == "":
																									_snake_case													=			old_model
																				else:
																									if not hasattr(UpperCamelCase__    ,      UpperCamelCase__  ):
																														raise ValueError(F'''{old_model} does not have {old_attribute}'''  )
																									_snake_case													=			getattr(UpperCamelCase__    ,      UpperCamelCase__  )

										if not is_key_init:
															raise ValueError(F'''{key} was not correctly initialized!'''  )

					print(F'''Saving model to {pytorch_dump_folder_path}'''  )
					prophet.save_pretrained(UpperCamelCase__  )


if __name__ == "__main__":
					UpperCAmelCase_ 							=       argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    """--prophetnet_checkpoint_path""", default=None, type=str, required=True, help="""Path the official PyTorch dump."""
					)
					parser.add_argument(
					    """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
					)
					UpperCAmelCase_ 							=       parser.parse_args()
					convert_prophetnet_checkpoint_to_pytorch(args.prophetnet_checkpoint_path, args.pytorch_dump_folder_path)



