

'''simple docstring'''





import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


UpperCamelCase_    						=       logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
							import torch_xla.core.xla_model as xm
							import torch_xla.debug.metrics as met





class 		__SCREAMING_SNAKE_CASE							(			lowercase__      ):







				def __init__(       self     :		Optional[int]		,					*UpperCAmelCase__     :		Any		,					UpperCAmelCase__     :		int=None		,					UpperCAmelCase__     :		Tuple=None		,					UpperCAmelCase__     :		Optional[int]=None		,					**UpperCAmelCase__     :		Tuple   ):
										'''simple docstring'''
										super().__init__(*UpperCAmelCase__		,					**UpperCAmelCase__   )
										lowercase :							Any    		=eval_examples
										lowercase :							Optional[int]    		=post_process_function
										lowercase :							List[Any]    		=quant_trainer_args
										lowercase :							Optional[int]    		=128  # default number of calibration samples







				def  lowerCamelCase_	(       self     :		str		,					UpperCAmelCase__     :		Dict=None   ):
										'''simple docstring'''
										if calib_dataset is None and self.calib_dataset is None:
																raise ValueError('''Trainer: calibration requires an calib_dataset.'''   )
										lowercase :							Tuple    		=calib_dataset if calib_dataset is not None else self.calib_dataset

										lowercase :							Union[str, Any]    		=self._remove_unused_columns(UpperCAmelCase__		,					description='''Calibration'''   )

										return DataLoader(
										    UpperCAmelCase__		,					batch_size=self.args.eval_batch_size		,					collate_fn=self.data_collator		,					drop_last=self.args.dataloader_drop_last		,					num_workers=self.args.dataloader_num_workers		,					pin_memory=self.args.dataloader_pin_memory		,					shuffle=UpperCAmelCase__		,					)







				def  lowerCamelCase_	(       self     :		Optional[int]		,					UpperCAmelCase__     :		List[Any]=None   ):
										'''simple docstring'''
										lowercase :							Optional[int]    		=self.train_dataset if calib_dataset is None else calib_dataset
										lowercase :							List[Any]    		=self.get_calib_dataloader(UpperCAmelCase__   )

										lowercase :							Optional[int]    		=self.model
										quant_trainer.configure_model(UpperCAmelCase__		,					self.quant_trainer_args		,					calib=UpperCAmelCase__   )
										model.eval()
										quant_trainer.enable_calibration(UpperCAmelCase__   )

										logger.info('''***** Running calibration *****'''   )
										logger.info(F'''  Num examples = {self.calib_num}'''   )
										logger.info(F'''  Batch size = {calib_dataloader.batch_size}'''   )

										for step, inputs in enumerate(UpperCAmelCase__   ):
																# Prediction step
																lowercase		,	lowercase		,	lowercase :							Optional[int]    		=self.prediction_step(UpperCAmelCase__		,					UpperCAmelCase__		,					prediction_loss_only=UpperCAmelCase__   )
																if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
																						break

										quant_trainer.finish_calibration(UpperCAmelCase__		,					self.quant_trainer_args   )
										lowercase :							int    		=model







				def  lowerCamelCase_	(       self     :		int		,					UpperCAmelCase__     :		str=None		,					UpperCAmelCase__     :		List[str]=None		,					UpperCAmelCase__     :		Tuple=None		,					UpperCAmelCase__     :		str = "eval"   ):
										'''simple docstring'''
										lowercase :							Tuple    		=self.eval_dataset if eval_dataset is None else eval_dataset
										lowercase :							str    		=self.get_eval_dataloader(UpperCAmelCase__   )
										lowercase :							str    		=self.eval_examples if eval_examples is None else eval_examples

										# Temporarily disable metric computation, we will do it in the loop here.
										lowercase :							int    		=self.compute_metrics
										lowercase :							Optional[Any]    		=None
										lowercase :							Tuple    		=self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
										try:
																lowercase :							Any    		=eval_loop(
																    UpperCAmelCase__		,					description='''Evaluation'''		,					prediction_loss_only=True if compute_metrics is None else None		,					ignore_keys=UpperCAmelCase__		,					)
										finally:
																lowercase :							Dict    		=compute_metrics

										if self.post_process_function is not None and self.compute_metrics is not None:
																lowercase :							int    		=self.post_process_function(UpperCAmelCase__		,					UpperCAmelCase__		,					output.predictions   )
																lowercase :							Union[str, Any]    		=self.compute_metrics(UpperCAmelCase__   )

																# Prefix all keys with metric_key_prefix + '_'
																for key in list(metrics.keys()   ):
																						if not key.startswith(F'''{metric_key_prefix}_'''   ):
																												lowercase :							Tuple    		=metrics.pop(UpperCAmelCase__   )

																self.log(UpperCAmelCase__   )
										else:
																lowercase :							int    		={}

										if self.args.tpu_metrics_debug or self.args.debug:
																# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
																xm.master_print(met.metrics_report()   )

										lowercase :							List[Any]    		=self.callback_handler.on_evaluate(self.args		,					self.state		,					self.control		,					UpperCAmelCase__   )
										return metrics







				def  lowerCamelCase_	(       self     :		List[str]		,					UpperCAmelCase__     :		Union[str, Any]		,					UpperCAmelCase__     :		Dict		,					UpperCAmelCase__     :		Union[str, Any]=None		,					UpperCAmelCase__     :		str = "test"   ):
										'''simple docstring'''
										lowercase :							str    		=self.get_test_dataloader(UpperCAmelCase__   )

										# Temporarily disable metric computation, we will do it in the loop here.
										lowercase :							List[Any]    		=self.compute_metrics
										lowercase :							List[Any]    		=None
										lowercase :							Optional[Any]    		=self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
										try:
																lowercase :							Any    		=eval_loop(
																    UpperCAmelCase__		,					description='''Prediction'''		,					prediction_loss_only=True if compute_metrics is None else None		,					ignore_keys=UpperCAmelCase__		,					)
										finally:
																lowercase :							Union[str, Any]    		=compute_metrics

										if self.post_process_function is None or self.compute_metrics is None:
																return output

										lowercase :							str    		=self.post_process_function(UpperCAmelCase__		,					UpperCAmelCase__		,					output.predictions		,					'''predict'''   )
										lowercase :							Dict    		=self.compute_metrics(UpperCAmelCase__   )

										# Prefix all keys with metric_key_prefix + '_'
										for key in list(metrics.keys()   ):
																if not key.startswith(F'''{metric_key_prefix}_'''   ):
																						lowercase :							Any    		=metrics.pop(UpperCAmelCase__   )

										return PredictionOutput(predictions=predictions.predictions		,					label_ids=predictions.label_ids		,					metrics=UpperCAmelCase__   )







				def  lowerCamelCase_	(       self     :		Optional[Any]		,					UpperCAmelCase__     :		int="./"   ):
										'''simple docstring'''
										lowercase :							List[Any]    		=self.eval_dataset
										lowercase :							Any    		=self.get_eval_dataloader(UpperCAmelCase__   )

										lowercase :							Tuple    		=next(iter(UpperCAmelCase__   )   )

										# saving device - to make it consistent
										lowercase :							int    		=torch.device('''cuda''' if torch.cuda.is_available() else '''cpu'''   )

										# convert to tuple
										lowercase :							Optional[Any]    		=tuple(v.to(UpperCAmelCase__   ) for k, v in batch.items()   )

										logger.info('''Converting model to be onnx compatible'''   )
										from pytorch_quantization.nn import TensorQuantizer

										lowercase :							List[Any]    		=True

										lowercase :							Optional[int]    		=self.model.to(UpperCAmelCase__   )

										model.eval()
										model.float()

										lowercase :							List[Any]    		=model.module if hasattr(UpperCAmelCase__		,					'''module'''   ) else model
										quant_trainer.configure_model(UpperCAmelCase__		,					self.quant_trainer_args   )

										lowercase :							List[str]    		=os.path.join(UpperCAmelCase__		,					'''model.onnx'''   )
										logger.info(F'''exporting model to {output_model_file}'''   )

										lowercase :							Dict    		={0: '''batch_size''', 1: '''seq_len'''}

										torch.onnx.export(
										    UpperCAmelCase__		,					UpperCAmelCase__		,					UpperCAmelCase__		,					export_params=UpperCAmelCase__		,					opset_version=13		,					do_constant_folding=UpperCAmelCase__		,					input_names=['''input_ids''', '''attention_mask''', '''token_type_ids''']		,					output_names=['''output_start_logits''', '''output_end_logits''']		,					dynamic_axes={
										        '''input_ids''': axes,
										        '''attention_mask''': axes,
										        '''token_type_ids''': axes,
										        '''output_start_logits''': axes,
										        '''output_end_logits''': axes,
										    }		,					verbose=UpperCAmelCase__		,					)
										logger.info('''onnx export finished'''   )






from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
				import tensorflow as tf

				from transformers import TFBlipTextModel
				from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST



class        lowerCamelCase__				:





			def __init__(       self	:	Optional[Any]  ,  lowercase__	:	int  ,  lowercase__	:	List[str]=12  ,  lowercase__	:	Optional[int]=7  ,  lowercase__	:	List[Any]=True  ,  lowercase__	:	str=True  ,  lowercase__	:	Optional[Any]=True  ,  lowercase__	:	List[str]=99  ,  lowercase__	:	Optional[int]=32  ,  lowercase__	:	Tuple=32  ,  lowercase__	:	Dict=2  ,  lowercase__	:	Optional[int]=4  ,  lowercase__	:	str=37  ,  lowercase__	:	Optional[Any]=0.1  ,  lowercase__	:	List[str]=0.1  ,  lowercase__	:	Union[str, Any]=5_12  ,  lowercase__	:	List[str]=0.0_2  ,  lowercase__	:	Tuple=0  ,  lowercase__	:	List[Any]=None  ,  ):
										_lowerCAmelCase							=							parent
										_lowerCAmelCase							=							batch_size
										_lowerCAmelCase							=							seq_length
										_lowerCAmelCase							=							is_training
										_lowerCAmelCase							=							use_input_mask
										_lowerCAmelCase							=							use_labels
										_lowerCAmelCase							=							vocab_size
										_lowerCAmelCase							=							hidden_size
										_lowerCAmelCase							=							projection_dim
										_lowerCAmelCase							=							num_hidden_layers
										_lowerCAmelCase							=							num_attention_heads
										_lowerCAmelCase							=							intermediate_size
										_lowerCAmelCase							=							dropout
										_lowerCAmelCase							=							attention_dropout
										_lowerCAmelCase							=							max_position_embeddings
										_lowerCAmelCase							=							initializer_range
										_lowerCAmelCase							=							scope
										_lowerCAmelCase							=							bos_token_id





			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										_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]						)

										if input_mask is not None:
																	_lowerCAmelCase							=							input_mask.numpy()
																	_lowerCAmelCase			,		_lowerCAmelCase							=							input_mask.shape
																	_lowerCAmelCase							=							np.random.randint(1  ,  seq_length - 1  ,  size=(batch_size,)						)
																	for batch_idx, start_index in enumerate(lowercase__						):
																								_lowerCAmelCase							=							1
																								_lowerCAmelCase							=							0

										_lowerCAmelCase							=							self.get_config()

										return config, input_ids, tf.convert_to_tensor(lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										return BlipTextConfig(
										    vocab_size=self.vocab_size  ,  hidden_size=self.hidden_size  ,  projection_dim=self.projection_dim  ,  num_hidden_layers=self.num_hidden_layers  ,  num_attention_heads=self.num_attention_heads  ,  intermediate_size=self.intermediate_size  ,  dropout=self.dropout  ,  attention_dropout=self.attention_dropout  ,  max_position_embeddings=self.max_position_embeddings  ,  initializer_range=self.initializer_range  ,  bos_token_id=self.bos_token_id  ,  )





			def   SCREAMING_SNAKE_CASE__     (       self	:	Tuple  ,  lowercase__	:	int  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Union[str, Any]						):
										_lowerCAmelCase							=							TFBlipTextModel(config=lowercase__						)
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  training=lowercase__						)
										_lowerCAmelCase							=							model(lowercase__  ,  training=lowercase__						)
										self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)						)
										self.parent.assertEqual(result.pooler_output.shape  ,  (self.batch_size, self.hidden_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										_lowerCAmelCase							=							self.prepare_config_and_inputs()
										_lowerCAmelCase			,		_lowerCAmelCase			,		_lowerCAmelCase							=							config_and_inputs
										_lowerCAmelCase							=							{'input_ids': input_ids, 'attention_mask': input_mask}
										return config, inputs_dict







@require_tf
class        lowerCamelCase__				(   UpperCAmelCase       ,unittest.TestCase				):
			UpperCamelCase__    							=(TFBlipTextModel,) if is_tf_available() else ()
			UpperCamelCase__    							=False
			UpperCamelCase__    							=False
			UpperCamelCase__    							=False





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										_lowerCAmelCase							=							BlipTextModelTester(self						)
										_lowerCAmelCase							=							ConfigTester(self  ,  config_class=lowercase__  ,  hidden_size=37						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Union[str, Any]						):
										self.config_tester.run_common_tests()





			def   SCREAMING_SNAKE_CASE__     (       self	:	Tuple						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict						):
										pass





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict						):
										pass





			@unittest.skip(reason='Blip does not use inputs_embeds'						)
			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										pass





			@unittest.skip(reason='BlipTextModel has no base class and is not available in MODEL_MAPPING'						)
			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										pass





			@unittest.skip(reason='BlipTextModel has no base class and is not available in MODEL_MAPPING'						)
			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										pass





			@slow
			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_lowerCAmelCase							=							TFBlipTextModel.from_pretrained(lowercase__						)
																	self.assertIsNotNone(lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any  ,  lowercase__	:	Any=True						):
										super().test_pt_tf_model_equivalence(allow_missing_keys=lowercase__						)






