



from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


UpperCAmelCase   :							List[Any]							 =logging.get_logger(__name__)

UpperCAmelCase   :							Dict							 ={
    """facebook/blenderbot_small-90M""": """https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/config.json""",
    # See all BlenderbotSmall models at https://huggingface.co/models?filter=blenderbot_small
}


class       _lowercase   (a_      ):
				'''simple docstring'''







				lowercase__   						=						"""blenderbot-small"""
				lowercase__   						=						["""past_key_values"""]
				lowercase__   						=						{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}

				def __init__( self     , snake_case__=5_0265     , snake_case__=512     , snake_case__=8     , snake_case__=2048     , snake_case__=16     , snake_case__=8     , snake_case__=2048     , snake_case__=16     , snake_case__=0.0     , snake_case__=0.0     , snake_case__=True     , snake_case__=True     , snake_case__="gelu"     , snake_case__=512     , snake_case__=0.1     , snake_case__=0.0     , snake_case__=0.0     , snake_case__=0.02     , snake_case__=1     , snake_case__=False     , snake_case__=0     , snake_case__=1     , snake_case__=2     , snake_case__=2     , **snake_case__     , ):

									'''simple docstring'''


									UpperCamelCase_		  =   vocab_size
									UpperCamelCase_		  =   max_position_embeddings
									UpperCamelCase_		  =   d_model
									UpperCamelCase_		  =   encoder_ffn_dim
									UpperCamelCase_		  =   encoder_layers
									UpperCamelCase_		  =   encoder_attention_heads
									UpperCamelCase_		  =   decoder_ffn_dim
									UpperCamelCase_		  =   decoder_layers
									UpperCamelCase_		  =   decoder_attention_heads
									UpperCamelCase_		  =   dropout
									UpperCamelCase_		  =   attention_dropout
									UpperCamelCase_		  =   activation_dropout
									UpperCamelCase_		  =   activation_function
									UpperCamelCase_		  =   init_std
									UpperCamelCase_		  =   encoder_layerdrop
									UpperCamelCase_		  =   decoder_layerdrop
									UpperCamelCase_		  =   use_cache
									UpperCamelCase_		  =   encoder_layers
									UpperCamelCase_		  =   scale_embedding  # scale factor will be sqrt(d_model) if True

									super().__init__(
									    pad_token_id=snake_case__     , bos_token_id=snake_case__     , eos_token_id=snake_case__     , is_encoder_decoder=snake_case__     , decoder_start_token_id=snake_case__     , forced_eos_token_id=snake_case__     , **snake_case__     , )




class       _lowercase   (a_      ):
				'''simple docstring'''







				@property
				def       _lowerCamelCase					( self    ):

									'''simple docstring'''


									if self.task in ["default", "seq2seq-lm"]:
														UpperCamelCase_		  =   OrderedDict(
														    [
														        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
														        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
														    ]    )

														if self.use_past:
																			UpperCamelCase_		  =   {0: "batch"}
																			UpperCamelCase_		  =   {0: "batch", 1: "past_decoder_sequence + sequence"}
														else:
																			UpperCamelCase_		  =   {0: "batch", 1: "decoder_sequence"}
																			UpperCamelCase_		  =   {0: "batch", 1: "decoder_sequence"}

														if self.use_past:
																			self.fill_with_past_key_values_(snake_case__     , direction="inputs"    )
									elif self.task == "causal-lm":
														# TODO: figure this case out.
														UpperCamelCase_		  =   OrderedDict(
														    [
														        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
														        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
														    ]    )
														if self.use_past:
																			UpperCamelCase_  ,    UpperCamelCase_		  =   self.num_layers
																			for i in range(snake_case__    ):
																								UpperCamelCase_		  =   {0: "batch", 2: "past_sequence + sequence"}
																								UpperCamelCase_		  =   {0: "batch", 2: "past_sequence + sequence"}
									else:
														UpperCamelCase_		  =   OrderedDict(
														    [
														        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
														        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
														        ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
														        ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
														    ]    )

									return common_inputs

				@property
				def       _lowerCamelCase					( self    ):

									'''simple docstring'''


									if self.task in ["default", "seq2seq-lm"]:
														UpperCamelCase_		  =   super().outputs
									else:
														UpperCamelCase_		  =   super(snake_case__     , self    ).outputs
														if self.use_past:
																			UpperCamelCase_  ,    UpperCamelCase_		  =   self.num_layers
																			for i in range(snake_case__    ):
																								UpperCamelCase_		  =   {0: "batch", 2: "past_sequence + sequence"}
																								UpperCamelCase_		  =   {0: "batch", 2: "past_sequence + sequence"}
									return common_outputs

				def       _lowerCamelCase					( self     , snake_case__     , snake_case__ = -1     , snake_case__ = -1     , snake_case__ = False     , snake_case__ = None     , ):

									'''simple docstring'''


									UpperCamelCase_		  =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
									    snake_case__     , snake_case__     , snake_case__     , snake_case__     , snake_case__    )

									# Generate decoder inputs
									UpperCamelCase_		  =   seq_length if not self.use_past else 1
									UpperCamelCase_		  =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
									    snake_case__     , snake_case__     , snake_case__     , snake_case__     , snake_case__    )
									UpperCamelCase_		  =   {F"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
									UpperCamelCase_		  =   dict(**snake_case__     , **snake_case__    )

									if self.use_past:
														if not is_torch_available():
																			raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."    )
														else:
																			import torch
														UpperCamelCase_  ,    UpperCamelCase_		  =   common_inputs["input_ids"].shape
														UpperCamelCase_		  =   common_inputs["decoder_input_ids"].shape[1]
														UpperCamelCase_  ,    UpperCamelCase_		  =   self.num_attention_heads
														UpperCamelCase_		  =   (
														    batch,
														    num_encoder_attention_heads,
														    encoder_seq_length,
														    self._config.hidden_size // num_encoder_attention_heads,
														)
														UpperCamelCase_		  =   decoder_seq_length + 3
														UpperCamelCase_		  =   (
														    batch,
														    num_decoder_attention_heads,
														    decoder_past_length,
														    self._config.hidden_size // num_decoder_attention_heads,
														)

														UpperCamelCase_		  =   torch.cat(
														    [common_inputs["decoder_attention_mask"], torch.ones(snake_case__     , snake_case__    )]     , dim=1    )

														UpperCamelCase_		  =   []
														# If the number of encoder and decoder layers are present in the model configuration, both are considered
														UpperCamelCase_  ,    UpperCamelCase_		  =   self.num_layers
														UpperCamelCase_		  =   min(snake_case__     , snake_case__    )
														UpperCamelCase_		  =   max(snake_case__     , snake_case__    ) - min_num_layers
														UpperCamelCase_		  =   "encoder" if num_encoder_layers > num_decoder_layers else "decoder"

														for _ in range(snake_case__    ):
																			common_inputs["past_key_values"].append(
																			    (
																			        torch.zeros(snake_case__    ),
																			        torch.zeros(snake_case__    ),
																			        torch.zeros(snake_case__    ),
																			        torch.zeros(snake_case__    ),
																			    )    )
														# TODO: test this.
														UpperCamelCase_		  =   encoder_shape if remaining_side_name == "encoder" else decoder_shape
														for _ in range(snake_case__     , snake_case__    ):
																			common_inputs["past_key_values"].append((torch.zeros(snake_case__    ), torch.zeros(snake_case__    ))    )
									return common_inputs

				def       _lowerCamelCase					( self     , snake_case__     , snake_case__ = -1     , snake_case__ = -1     , snake_case__ = False     , snake_case__ = None     , ):

									'''simple docstring'''


									UpperCamelCase_		  =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
									    snake_case__     , snake_case__     , snake_case__     , snake_case__     , snake_case__    )

									if self.use_past:
														if not is_torch_available():
																			raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."    )
														else:
																			import torch
														UpperCamelCase_  ,    UpperCamelCase_		  =   common_inputs["input_ids"].shape
														# Not using the same length for past_key_values
														UpperCamelCase_		  =   seqlen + 2
														UpperCamelCase_  ,    UpperCamelCase_		  =   self.num_layers
														UpperCamelCase_  ,    UpperCamelCase_		  =   self.num_attention_heads
														UpperCamelCase_		  =   (
														    batch,
														    num_encoder_attention_heads,
														    past_key_values_length,
														    self._config.hidden_size // num_encoder_attention_heads,
														)

														UpperCamelCase_		  =   common_inputs["attention_mask"].dtype
														UpperCamelCase_		  =   torch.cat(
														    [common_inputs["attention_mask"], torch.ones(snake_case__     , snake_case__     , dtype=snake_case__    )]     , dim=1    )
														UpperCamelCase_		  =   [
														    (torch.zeros(snake_case__    ), torch.zeros(snake_case__    )) for _ in range(snake_case__    )
														]
									return common_inputs

				def       _lowerCamelCase					( self     , snake_case__     , snake_case__ = -1     , snake_case__ = -1     , snake_case__ = False     , snake_case__ = None     , ):

									'''simple docstring'''


									UpperCamelCase_		  =   compute_effective_axis_dimension(
									    snake_case__     , fixed_dimension=OnnxConfig.default_fixed_batch     , num_token_to_add=0    )

									# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
									UpperCamelCase_		  =   tokenizer.num_special_tokens_to_add(snake_case__    )
									UpperCamelCase_		  =   compute_effective_axis_dimension(
									    snake_case__     , fixed_dimension=OnnxConfig.default_fixed_sequence     , num_token_to_add=snake_case__    )

									# Generate dummy inputs according to compute batch and sequence
									UpperCamelCase_		  =   [" ".join([tokenizer.unk_token]    ) * seq_length] * batch_size
									UpperCamelCase_		  =   dict(tokenizer(snake_case__     , return_tensors=snake_case__    )    )
									return common_inputs

				def       _lowerCamelCase					( self     , snake_case__     , snake_case__ = -1     , snake_case__ = -1     , snake_case__ = False     , snake_case__ = None     , ):

									'''simple docstring'''


									if self.task in ["default", "seq2seq-lm"]:
														UpperCamelCase_		  =   self._generate_dummy_inputs_for_default_and_seqaseq_lm(
														    snake_case__     , batch_size=snake_case__     , seq_length=snake_case__     , is_pair=snake_case__     , framework=snake_case__    )

									elif self.task == "causal-lm":
														UpperCamelCase_		  =   self._generate_dummy_inputs_for_causal_lm(
														    snake_case__     , batch_size=snake_case__     , seq_length=snake_case__     , is_pair=snake_case__     , framework=snake_case__    )
									else:
														UpperCamelCase_		  =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
														    snake_case__     , batch_size=snake_case__     , seq_length=snake_case__     , is_pair=snake_case__     , framework=snake_case__    )

									return common_inputs

				def       _lowerCamelCase					( self     , snake_case__     , snake_case__     , snake_case__     , snake_case__    ):

									'''simple docstring'''


									if self.task in ["default", "seq2seq-lm"]:
														UpperCamelCase_		  =   super()._flatten_past_key_values_(snake_case__     , snake_case__     , snake_case__     , snake_case__    )
									else:
														UpperCamelCase_		  =   super(snake_case__     , self    )._flatten_past_key_values_(
														    snake_case__     , snake_case__     , snake_case__     , snake_case__    )





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   :							str							 =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       _lowercase   (a_      ):
				'''simple docstring'''







				def __init__( self     , *snake_case__     , snake_case__=None     , snake_case__=None     , snake_case__=None     , **snake_case__    ):

									'''simple docstring'''


									super().__init__(*snake_case__     , **snake_case__    )
									UpperCamelCase_		  =   eval_examples
									UpperCamelCase_		  =   post_process_function
									UpperCamelCase_		  =   quant_trainer_args
									UpperCamelCase_		  =   128  # default number of calibration samples

				def       _lowerCamelCase					( self     , snake_case__=None    ):

									'''simple docstring'''


									if calib_dataset is None and self.calib_dataset is None:
														raise ValueError("Trainer: calibration requires an calib_dataset."    )
									UpperCamelCase_		  =   calib_dataset if calib_dataset is not None else self.calib_dataset

									UpperCamelCase_		  =   self._remove_unused_columns(snake_case__     , description="Calibration"    )

									return DataLoader(
									    snake_case__     , 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=snake_case__     , )

				def       _lowerCamelCase					( self     , snake_case__=None    ):

									'''simple docstring'''


									UpperCamelCase_		  =   self.train_dataset if calib_dataset is None else calib_dataset
									UpperCamelCase_		  =   self.get_calib_dataloader(snake_case__    )

									UpperCamelCase_		  =   self.model
									quant_trainer.configure_model(snake_case__     , self.quant_trainer_args     , calib=snake_case__    )
									model.eval()
									quant_trainer.enable_calibration(snake_case__    )

									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(snake_case__    ):
														# Prediction step
														UpperCamelCase_  ,    UpperCamelCase_  ,    UpperCamelCase_		  =   self.prediction_step(snake_case__     , snake_case__     , prediction_loss_only=snake_case__    )
														if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
																			break

									quant_trainer.finish_calibration(snake_case__     , self.quant_trainer_args    )
									UpperCamelCase_		  =   model

				def       _lowerCamelCase					( self     , snake_case__=None     , snake_case__=None     , snake_case__=None     , snake_case__ = "eval"    ):

									'''simple docstring'''


									UpperCamelCase_		  =   self.eval_dataset if eval_dataset is None else eval_dataset
									UpperCamelCase_		  =   self.get_eval_dataloader(snake_case__    )
									UpperCamelCase_		  =   self.eval_examples if eval_examples is None else eval_examples

									# Temporarily disable metric computation, we will do it in the loop here.
									UpperCamelCase_		  =   self.compute_metrics
									UpperCamelCase_		  =   None
									UpperCamelCase_		  =   self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
									try:
														UpperCamelCase_		  =   eval_loop(
														    snake_case__     , description="Evaluation"     , prediction_loss_only=True if compute_metrics is None else None     , ignore_keys=snake_case__     , )
									finally:
														UpperCamelCase_		  =   compute_metrics

									if self.post_process_function is not None and self.compute_metrics is not None:
														UpperCamelCase_		  =   self.post_process_function(snake_case__     , snake_case__     , output.predictions    )
														UpperCamelCase_		  =   self.compute_metrics(snake_case__    )

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

														self.log(snake_case__    )
									else:
														UpperCamelCase_		  =   {}

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

									UpperCamelCase_		  =   self.callback_handler.on_evaluate(self.args     , self.state     , self.control     , snake_case__    )
									return metrics

				def       _lowerCamelCase					( self     , snake_case__     , snake_case__     , snake_case__=None     , snake_case__ = "test"    ):

									'''simple docstring'''


									UpperCamelCase_		  =   self.get_test_dataloader(snake_case__    )

									# Temporarily disable metric computation, we will do it in the loop here.
									UpperCamelCase_		  =   self.compute_metrics
									UpperCamelCase_		  =   None
									UpperCamelCase_		  =   self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
									try:
														UpperCamelCase_		  =   eval_loop(
														    snake_case__     , description="Prediction"     , prediction_loss_only=True if compute_metrics is None else None     , ignore_keys=snake_case__     , )
									finally:
														UpperCamelCase_		  =   compute_metrics

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

									UpperCamelCase_		  =   self.post_process_function(snake_case__     , snake_case__     , output.predictions     , "predict"    )
									UpperCamelCase_		  =   self.compute_metrics(snake_case__    )

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

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

				def       _lowerCamelCase					( self     , snake_case__="./"    ):

									'''simple docstring'''


									UpperCamelCase_		  =   self.eval_dataset
									UpperCamelCase_		  =   self.get_eval_dataloader(snake_case__    )

									UpperCamelCase_		  =   next(iter(snake_case__    )    )

									# saving device - to make it consistent
									UpperCamelCase_		  =   torch.device("cuda" if torch.cuda.is_available() else "cpu"    )

									# convert to tuple
									UpperCamelCase_		  =   tuple(v.to(snake_case__    ) for k, v in batch.items()    )

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

									UpperCamelCase_		  =   True

									UpperCamelCase_		  =   self.model.to(snake_case__    )

									model.eval()
									model.float()

									UpperCamelCase_		  =   model.module if hasattr(snake_case__     , "module"    ) else model
									quant_trainer.configure_model(snake_case__     , self.quant_trainer_args    )

									UpperCamelCase_		  =   os.path.join(snake_case__     , "model.onnx"    )
									logger.info(F"""exporting model to {output_model_file}"""    )

									UpperCamelCase_		  =   {0: "batch_size", 1: "seq_len"}

									torch.onnx.export(
									    snake_case__     , snake_case__     , snake_case__     , export_params=snake_case__     , opset_version=13     , do_constant_folding=snake_case__     , 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=snake_case__     , )
									logger.info("onnx export finished"    )
