'''simple docstring'''

import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
							import tensorflow as tf
							from tensorflow.python.framework.errors_impl import ResourceExhaustedError

							from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
							import pyanvml.pyanvml as nvml

lowercase_     		= logging.get_logger(__name__)



def 							UpperCamelCase__			(			a__			,      a__      ):


						'''simple docstring'''






						def run_func(a__      ):
												@wraps(A__      )
												def run_in_eager_mode(*a__			,      **a__      ):
																		return func(*A__			,      **A__      )

												@wraps(A__      )
												@tf.function(experimental_compile=A__      )
												def run_in_graph_mode(*a__			,      **a__      ):
																		return func(*A__			,      **A__      )

												if do_eager_mode is True:
																		if use_xla is not False:
																								raise ValueError(
																								    'Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.'      )
																		return run_in_eager_mode
												else:
																		return run_in_graph_mode

						return run_func



def 							UpperCamelCase__			(			a__			,      a__			,      a__      ):


						'''simple docstring'''






						_lowerCAmelCase										=random.Random()
						_lowerCAmelCase										=[rng.randint(0			,      vocab_size - 1      ) for i in range(batch_size * sequence_length      )]
						return tf.constant(A__			,      shape=(batch_size, sequence_length)			,      dtype=tf.intaa      )


class   SCREAMING_SNAKE_CASE			(						__lowerCamelCase):







		"""simple docstring"""
		lowercase   :      TensorFlowBenchmarkArguments
		lowercase   :      PretrainedConfig
		lowercase   :      str       =							"TensorFlow"






		@property
		def 		UpperCamelCase__       (   self				) ->     Optional[int]:
								return tf.__version__






		def 		UpperCamelCase__       (   self		,						__A		,						__A		,						__A				) ->     Union[str, Any]:
								# initialize GPU on separate process
								_lowerCAmelCase										=self.args.strategy
								if strategy is None:
														raise ValueError('A device strategy has to be initialized before using TensorFlow.'				)
								_lowerCAmelCase										=self._prepare_inference_func(UpperCamelCase_		,						UpperCamelCase_		,						UpperCamelCase_				)
								return self._measure_speed(_inference				)






		def 		UpperCamelCase__       (   self		,						__A		,						__A		,						__A				) ->     List[str]:
								_lowerCAmelCase										=self.args.strategy
								if strategy is None:
														raise ValueError('A device strategy has to be initialized before using TensorFlow.'				)
								_lowerCAmelCase										=self._prepare_train_func(UpperCamelCase_		,						UpperCamelCase_		,						UpperCamelCase_				)
								return self._measure_speed(_train				)






		def 		UpperCamelCase__       (   self		,						__A		,						__A		,						__A				) ->     Tuple:
								# initialize GPU on separate process
								if self.args.is_gpu:
														tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]		,						UpperCamelCase_				)
								_lowerCAmelCase										=self.args.strategy
								if strategy is None:
														raise ValueError('A device strategy has to be initialized before using TensorFlow.'				)
								_lowerCAmelCase										=self._prepare_inference_func(UpperCamelCase_		,						UpperCamelCase_		,						UpperCamelCase_				)
								return self._measure_memory(_inference				)






		def 		UpperCamelCase__       (   self		,						__A		,						__A		,						__A				) ->     List[str]:
								if self.args.is_gpu:
														tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]		,						UpperCamelCase_				)
								_lowerCAmelCase										=self.args.strategy
								if strategy is None:
														raise ValueError('A device strategy has to be initialized before using TensorFlow.'				)

								_lowerCAmelCase										=self._prepare_train_func(UpperCamelCase_		,						UpperCamelCase_		,						UpperCamelCase_				)
								return self._measure_memory(_train				)






		def 		UpperCamelCase__       (   self		,						__A		,						__A		,						__A				) ->     List[Any]:
								_lowerCAmelCase										=self.config_dict[model_name]

								if self.args.fpaa:
														raise NotImplementedError('Mixed precision is currently not supported.'				)

								_lowerCAmelCase										=(
								    hasattr(UpperCamelCase_		,						'architectures'				)
								    and isinstance(config.architectures		,						UpperCamelCase_				)
								    and len(config.architectures				) > 0
								)
								if not self.args.only_pretrain_model and has_model_class_in_config:
														try:
																				_lowerCAmelCase										='TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
																				_lowerCAmelCase										=__import__('transformers'		,						fromlist=[model_class]				)
																				_lowerCAmelCase										=getattr(UpperCamelCase_		,						UpperCamelCase_				)
																				_lowerCAmelCase										=model_cls(UpperCamelCase_				)
														except ImportError:
																				raise ImportError(
																				    F'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
																				    ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'				)
								else:
														_lowerCAmelCase										=TF_MODEL_MAPPING[config.__class__](UpperCamelCase_				)

								# encoder-decoder has vocab size saved differently
								_lowerCAmelCase										=config.vocab_size if hasattr(UpperCamelCase_		,						'vocab_size'				) else config.encoder.vocab_size
								_lowerCAmelCase										=random_input_ids(UpperCamelCase_		,						UpperCamelCase_		,						UpperCamelCase_				)

								@run_with_tf_optimizations(self.args.eager_mode		,						self.args.use_xla				)
								def encoder_decoder_forward():
														return model(UpperCamelCase_		,						decoder_input_ids=UpperCamelCase_		,						training=UpperCamelCase_				)

								@run_with_tf_optimizations(self.args.eager_mode		,						self.args.use_xla				)
								def encoder_forward():
														return model(UpperCamelCase_		,						training=UpperCamelCase_				)

								_lowerCAmelCase										=encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

								return _inference






		def 		UpperCamelCase__       (   self		,						__A		,						__A		,						__A				) ->     Any:
								_lowerCAmelCase										=self.config_dict[model_name]

								if self.args.eager_mode is not False:
														raise ValueError('Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.'				)

								if self.args.fpaa:
														raise NotImplementedError('Mixed precision is currently not supported.'				)

								_lowerCAmelCase										=(
								    hasattr(UpperCamelCase_		,						'architectures'				)
								    and isinstance(config.architectures		,						UpperCamelCase_				)
								    and len(config.architectures				) > 0
								)
								if not self.args.only_pretrain_model and has_model_class_in_config:
														try:
																				_lowerCAmelCase										='TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
																				_lowerCAmelCase										=__import__('transformers'		,						fromlist=[model_class]				)
																				_lowerCAmelCase										=getattr(UpperCamelCase_		,						UpperCamelCase_				)
																				_lowerCAmelCase										=model_cls(UpperCamelCase_				)
														except ImportError:
																				raise ImportError(
																				    F'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
																				    ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'				)
								else:
														_lowerCAmelCase										=TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](UpperCamelCase_				)

								# encoder-decoder has vocab size saved differently
								_lowerCAmelCase										=config.vocab_size if hasattr(UpperCamelCase_		,						'vocab_size'				) else config.encoder.vocab_size
								_lowerCAmelCase										=random_input_ids(UpperCamelCase_		,						UpperCamelCase_		,						UpperCamelCase_				)

								@run_with_tf_optimizations(self.args.eager_mode		,						self.args.use_xla				)
								def encoder_decoder_train():
														_lowerCAmelCase										=model(UpperCamelCase_		,						decoder_input_ids=UpperCamelCase_		,						labels=UpperCamelCase_		,						training=UpperCamelCase_				)[0]
														_lowerCAmelCase										=tf.gradients(UpperCamelCase_		,						model.trainable_variables				)
														return gradients

								@run_with_tf_optimizations(self.args.eager_mode		,						self.args.use_xla				)
								def encoder_train():
														_lowerCAmelCase										=model(UpperCamelCase_		,						labels=UpperCamelCase_		,						training=UpperCamelCase_				)[0]
														_lowerCAmelCase										=tf.gradients(UpperCamelCase_		,						model.trainable_variables				)
														return gradients

								_lowerCAmelCase										=encoder_decoder_train if config.is_encoder_decoder else encoder_train

								return _train






		def 		UpperCamelCase__       (   self		,						__A				) ->     Union[str, Any]:
								with self.args.strategy.scope():
														try:
																				if self.args.is_tpu or self.args.use_xla:
																										# run additional 10 times to stabilize compilation for tpu
																										logger.info('Do inference on TPU. Running model 5 times to stabilize compilation'				)
																										timeit.repeat(UpperCamelCase_		,						repeat=1		,						number=5				)

																				# as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
																				_lowerCAmelCase										=timeit.repeat(
																				    UpperCamelCase_		,						repeat=self.args.repeat		,						number=10		,						)

																				return min(UpperCamelCase_				) / 10.0
														except ResourceExhaustedError as e:
																				self.print_fn(F'''Doesn\'t fit on GPU. {e}'''				)






		def 		UpperCamelCase__       (   self		,						__A				) ->     List[Any]:
								logger.info(
								    'Note that TensorFlow allocates more memory than '
								    'it might need to speed up computation. '
								    'The memory reported here corresponds to the memory '
								    'reported by `nvidia-smi`, which can vary depending '
								    'on total available memory on the GPU that is used.'				)
								with self.args.strategy.scope():
														try:
																				if self.args.trace_memory_line_by_line:
																										if not self.args.eager_mode:
																																raise ValueError(
																																    '`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory'
																																    ' consumption line by line.'				)
																										_lowerCAmelCase										=start_memory_tracing('transformers'				)

																				if self.args.is_tpu:
																										# tpu
																										raise NotImplementedError(
																										    'Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking'
																										    ' with `args.memory=False`'				)
																				elif self.args.is_gpu:
																										# gpu
																										if not is_pyanvml_available():
																																logger.warning(
																																    'py3nvml not installed, we won\'t log GPU memory usage. '
																																    'Install py3nvml (pip install py3nvml) to log information about GPU.'				)
																																_lowerCAmelCase										='N/A'
																										else:
																																logger.info(
																																    'Measuring total GPU usage on GPU device. Make sure to not have additional processes'
																																    ' running on the same GPU.'				)
																																# init nvml
																																nvml.nvmlInit()
																																func()
																																_lowerCAmelCase										=nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx				)
																																_lowerCAmelCase										=nvml.nvmlDeviceGetMemoryInfo(UpperCamelCase_				)
																																_lowerCAmelCase										=meminfo.used
																																_lowerCAmelCase										=Memory(UpperCamelCase_				)
																																# shutdown nvml
																																nvml.nvmlShutdown()
																				else:
																										# cpu
																										if self.args.trace_memory_line_by_line:
																																logger.info(
																																    'When enabling line by line tracing, the max peak memory for CPU is inaccurate in'
																																    ' TensorFlow.'				)
																																_lowerCAmelCase										=None
																										else:
																																_lowerCAmelCase										=measure_peak_memory_cpu(UpperCamelCase_				)
																																_lowerCAmelCase										=Memory(UpperCamelCase_				) if isinstance(UpperCamelCase_		,						UpperCamelCase_				) else memory_bytes
																				if self.args.trace_memory_line_by_line:
																										_lowerCAmelCase										=stop_memory_tracing(UpperCamelCase_				)
																										if memory is None:
																																_lowerCAmelCase										=summary.total
																				else:
																										_lowerCAmelCase										=None

																				return memory, summary
														except ResourceExhaustedError as e:
																				self.print_fn(F'''Doesn\'t fit on GPU. {e}'''				)
																				return "N/A", None






'''simple docstring'''

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


lowercase_     		= logging.get_logger(__name__)

lowercase_     		= {
    '''salesforce/blip2-opt-2.7b''': '''https://huggingface.co/salesforce/blip2-opt-2.7b/resolve/main/config.json''',
}
class   SCREAMING_SNAKE_CASE			(						__lowercase):







		"""simple docstring"""
		lowercase   :      Tuple       =							'blip_2_vision_model'






		def __init__(   self		,						__A=1408		,						__A=6144		,						__A=39		,						__A=16		,						__A=224		,						__A=14		,						__A="gelu"		,						__A=0.00_001		,						__A=0.0		,						__A=1E-10		,						__A=True		,						**__A		,						) ->     int:
								super().__init__(**__A				)

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






		@classmethod
		def 		UpperCamelCase__       (   cls		,						__A		,						**__A				) ->     "PretrainedConfig":
								cls._set_token_in_kwargs(__A				)

								_lowerCAmelCase						,				_lowerCAmelCase										=cls.get_config_dict(__A		,						**__A				)

								# get the vision config dict if we are loading from Blip2Config
								if config_dict.get('model_type'				) == "blip-2":
														_lowerCAmelCase										=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(__A		,						**__A				)

class   SCREAMING_SNAKE_CASE			(						__lowercase):







		"""simple docstring"""
		lowercase   :      int       =							'blip_2_qformer'






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

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






		@classmethod
		def 		UpperCamelCase__       (   cls		,						__A		,						**__A				) ->     "PretrainedConfig":
								cls._set_token_in_kwargs(__A				)

								_lowerCAmelCase						,				_lowerCAmelCase										=cls.get_config_dict(__A		,						**__A				)

								# get the qformer config dict if we are loading from Blip2Config
								if config_dict.get('model_type'				) == "blip-2":
														_lowerCAmelCase										=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(__A		,						**__A				)

class   SCREAMING_SNAKE_CASE			(						__lowercase):







		"""simple docstring"""
		lowercase   :      Optional[int]       =							'blip-2'
		lowercase   :      Any       =							True






		def __init__(   self		,						__A=None		,						__A=None		,						__A=None		,						__A=32		,						**__A				) ->     int:
								super().__init__(**__A				)

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

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

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

								_lowerCAmelCase										=BlipaVisionConfig(**__A				)
								_lowerCAmelCase										=BlipaQFormerConfig(**__A				)
								_lowerCAmelCase										=text_config['model_type'] if 'model_type' in text_config else 'opt'
								_lowerCAmelCase										=CONFIG_MAPPING[text_model_type](**__A				)

								_lowerCAmelCase										=self.text_config.tie_word_embeddings
								_lowerCAmelCase										=self.text_config.is_encoder_decoder

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






		@classmethod
		def 		UpperCamelCase__       (   cls		,						__A		,						__A		,						__A		,						**__A		,						) ->     Any:
								return cls(
								    vision_config=vision_config.to_dict()		,						qformer_config=qformer_config.to_dict()		,						text_config=text_config.to_dict()		,						**__A		,						)






		def 		UpperCamelCase__       (   self				) ->     Tuple:
								_lowerCAmelCase										=copy.deepcopy(self.__dict__				)
								_lowerCAmelCase										=self.vision_config.to_dict()
								_lowerCAmelCase										=self.qformer_config.to_dict()
								_lowerCAmelCase										=self.text_config.to_dict()
								_lowerCAmelCase										=self.__class__.model_type
								return output





