






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 				__UpperCAmelCase		(_UpperCAmelCase   ):







			def __init__(     self:  Tuple				,     *UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  int=None				,     UpperCAmelCase_:  Optional[int]=None				,     UpperCAmelCase_:  Any=None				,     **UpperCAmelCase_:  Optional[Any]					):







									'''simple docstring'''




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







			def 							UpperCamelCase				(     self:  List[str]				,     UpperCAmelCase_:  Optional[int]=None					):







									'''simple docstring'''




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

									_SCREAMING_SNAKE_CASE     				=  self._remove_unused_columns(_SCREAMING_SNAKE_CASE				,     description="""Calibration"""					)

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







			def 							UpperCamelCase				(     self:  Tuple				,     UpperCAmelCase_:  Union[str, Any]=None					):







									'''simple docstring'''




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

									_SCREAMING_SNAKE_CASE     				=  self.model
									quant_trainer.configure_model(_SCREAMING_SNAKE_CASE				,     self.quant_trainer_args				,     calib=_SCREAMING_SNAKE_CASE					)
									model.eval()
									quant_trainer.enable_calibration(_SCREAMING_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(_SCREAMING_SNAKE_CASE					):
															# Prediction step
															_SCREAMING_SNAKE_CASE     				=  self.prediction_step(_SCREAMING_SNAKE_CASE				,     _SCREAMING_SNAKE_CASE				,     prediction_loss_only=_SCREAMING_SNAKE_CASE					)
															if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
																					break

									quant_trainer.finish_calibration(_SCREAMING_SNAKE_CASE				,     self.quant_trainer_args					)
									_SCREAMING_SNAKE_CASE     				=  model







			def 							UpperCamelCase				(     self:  Optional[int]				,     UpperCAmelCase_:  Any=None				,     UpperCAmelCase_:  Dict=None				,     UpperCAmelCase_:  List[str]=None				,     UpperCAmelCase_:  Any = "eval"					):







									'''simple docstring'''




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

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

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

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

															self.log(_SCREAMING_SNAKE_CASE					)
									else:
															_SCREAMING_SNAKE_CASE     				=  {}

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

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







			def 							UpperCamelCase				(     self:  str				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  List[Any]=None				,     UpperCAmelCase_:  Tuple = "test"					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.get_test_dataloader(_SCREAMING_SNAKE_CASE					)

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

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

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

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

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







			def 							UpperCamelCase				(     self:  List[str]				,     UpperCAmelCase_:  Optional[Any]="./"					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.eval_dataset
									_SCREAMING_SNAKE_CASE     				=  self.get_eval_dataloader(_SCREAMING_SNAKE_CASE					)

									_SCREAMING_SNAKE_CASE     				=  next(iter(_SCREAMING_SNAKE_CASE					)					)

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

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

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

									_SCREAMING_SNAKE_CASE     				=  True

									_SCREAMING_SNAKE_CASE     				=  self.model.to(_SCREAMING_SNAKE_CASE					)

									model.eval()
									model.float()

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

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

									_SCREAMING_SNAKE_CASE     				=  {0: '''batch_size''', 1: '''seq_len'''}

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









# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..models.auto import AutoProcessor
from ..models.vision_encoder_decoder import VisionEncoderDecoderModel
from ..utils import is_vision_available
from .base import PipelineTool


if is_vision_available():
		from PIL import Image



class 				__UpperCAmelCase		(_UpperCAmelCase   ):
			__snake_case    :				Dict									=							"naver-clova-ix/donut-base-finetuned-docvqa"
			__snake_case    :				Dict									=							(
			    "This is a tool that answers a question about an document (pdf). It takes an input named `document` which "
			    "should be the document containing the information, as well as a `question` that is the question about the "
			    "document. It returns a text that contains the answer to the question."
			)
			__snake_case    :				Dict									=							"document_qa"
			__snake_case    :				Any									=							AutoProcessor
			__snake_case    :				int									=							VisionEncoderDecoderModel

			__snake_case    :				Union[str, Any]									=							["image", "text"]
			__snake_case    :				Optional[int]									=							["text"]







			def __init__(     self:  List[Any]				,     *UpperCAmelCase_:  Optional[int]				,     **UpperCAmelCase_:  List[str]					):







									'''simple docstring'''




									if not is_vision_available():
															raise ValueError("""Pillow must be installed to use the DocumentQuestionAnsweringTool."""					)

									super().__init__(*UpperCAmelCase_				,     **UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[int]				,     UpperCAmelCase_:  "Image"				,     UpperCAmelCase_:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  """<s_docvqa><s_question>{user_input}</s_question><s_answer>"""
									_SCREAMING_SNAKE_CASE     				=  task_prompt.replace("""{user_input}"""				,     UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  self.pre_processor.tokenizer(
									    UpperCAmelCase_				,     add_special_tokens=UpperCAmelCase_				,     return_tensors="""pt"""					).input_ids
									_SCREAMING_SNAKE_CASE     				=  self.pre_processor(UpperCAmelCase_				,     return_tensors="""pt"""					).pixel_values

									return {"decoder_input_ids": decoder_input_ids, "pixel_values": pixel_values}







			def 							UpperCamelCase				(     self:  Dict				,     UpperCAmelCase_:  str					):







									'''simple docstring'''




									return self.model.generate(
									    inputs["""pixel_values"""].to(self.device					)				,     decoder_input_ids=inputs["""decoder_input_ids"""].to(self.device					)				,     max_length=self.model.decoder.config.max_position_embeddings				,     early_stopping=UpperCAmelCase_				,     pad_token_id=self.pre_processor.tokenizer.pad_token_id				,     eos_token_id=self.pre_processor.tokenizer.eos_token_id				,     use_cache=UpperCAmelCase_				,     num_beams=1				,     bad_words_ids=[[self.pre_processor.tokenizer.unk_token_id]]				,     return_dict_in_generate=UpperCAmelCase_				,     ).sequences







			def 							UpperCamelCase				(     self:  Dict				,     UpperCAmelCase_:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.pre_processor.batch_decode(UpperCAmelCase_					)[0]
									_SCREAMING_SNAKE_CASE     				=  sequence.replace(self.pre_processor.tokenizer.eos_token				,     """"""					)
									_SCREAMING_SNAKE_CASE     				=  sequence.replace(self.pre_processor.tokenizer.pad_token				,     """"""					)
									_SCREAMING_SNAKE_CASE     				=  re.sub(R"""<.*?>"""				,     """"""				,     UpperCAmelCase_				,     count=1					).strip()  # remove first task start token
									_SCREAMING_SNAKE_CASE     				=  self.pre_processor.tokenajson(UpperCAmelCase_					)

									return sequence["answer"]

