
"""simple docstring"""
from collections import OrderedDict
from typing import Any, Mapping, Optional

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


lowercase_       :						List[str]     					=	logging.get_logger(__name__)

lowercase_       :						str     					=	{
    '''Helsinki-NLP/opus-mt-en-de''': '''https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/config.json''',
    # See all Marian models at https://huggingface.co/models?filter=marian
}
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE  ):
   A__      =			"""marian"""
   A__      =			["""past_key_values"""]
   A__      =			{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}
   def __init__(     self  ,     snake_case__=58101  ,     snake_case__=None  ,     snake_case__=1024  ,     snake_case__=12  ,     snake_case__=4096  ,     snake_case__=16  ,     snake_case__=12  ,     snake_case__=4096  ,     snake_case__=16  ,     snake_case__=0.0  ,     snake_case__=0.0  ,     snake_case__=True  ,     snake_case__=True  ,     snake_case__="gelu"  ,     snake_case__=1024  ,     snake_case__=0.1  ,     snake_case__=0.0  ,     snake_case__=0.0  ,     snake_case__=0.02  ,     snake_case__=58100  ,     snake_case__=False  ,     snake_case__=58100  ,     snake_case__=0  ,     snake_case__=0  ,     snake_case__=True  ,     **snake_case__  ,     ):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				str							    = vocab_size
     _SCREAMING_SNAKE_CASE   :				Any							    = decoder_vocab_size or vocab_size
     _SCREAMING_SNAKE_CASE   :				Optional[int]							    = max_position_embeddings
     _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = d_model
     _SCREAMING_SNAKE_CASE   :				Optional[int]							    = encoder_ffn_dim
     _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = encoder_layers
     _SCREAMING_SNAKE_CASE   :				Any							    = encoder_attention_heads
     _SCREAMING_SNAKE_CASE   :				Dict							    = decoder_ffn_dim
     _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = decoder_layers
     _SCREAMING_SNAKE_CASE   :				Tuple							    = decoder_attention_heads
     _SCREAMING_SNAKE_CASE   :				str							    = dropout
     _SCREAMING_SNAKE_CASE   :				int							    = attention_dropout
     _SCREAMING_SNAKE_CASE   :				int							    = activation_dropout
     _SCREAMING_SNAKE_CASE   :				int							    = activation_function
     _SCREAMING_SNAKE_CASE   :				int							    = init_std
     _SCREAMING_SNAKE_CASE   :				Tuple							    = encoder_layerdrop
     _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = decoder_layerdrop
     _SCREAMING_SNAKE_CASE   :				Dict							    = use_cache
     _SCREAMING_SNAKE_CASE   :				Optional[int]							    = encoder_layers
     _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = scale_embedding  # scale factor will be sqrt(d_model) if True
     _SCREAMING_SNAKE_CASE   :				int							    = share_encoder_decoder_embeddings
     super().__init__(
         pad_token_id=__A  ,     eos_token_id=__A  ,     is_encoder_decoder=__A  ,     decoder_start_token_id=__A  ,     forced_eos_token_id=__A  ,     **__A  ,     )
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE  ):
   @property
   # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


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

       if self.use_past:
         _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = {0: "batch"}
         _SCREAMING_SNAKE_CASE   :				str							    = {0: "batch", 1: "past_decoder_sequence + sequence"}
       else:
         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = {0: "batch", 1: "decoder_sequence"}
         _SCREAMING_SNAKE_CASE   :				Dict							    = {0: "batch", 1: "decoder_sequence"}

       if self.use_past:
         self.fill_with_past_key_values_(__A  ,     direction="inputs"				)
     elif self.task == "causal-lm":
       # TODO: figure this case out.
       _SCREAMING_SNAKE_CASE   :				Dict							    = OrderedDict(
           [
               ("input_ids", {0: "batch", 1: "encoder_sequence"}),
               ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
           ]				)
       if self.use_past:
         _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.num_layers
         for i in range(__A				):
           _SCREAMING_SNAKE_CASE   :				Any							    = {0: "batch", 2: "past_sequence + sequence"}
           _SCREAMING_SNAKE_CASE   :				Optional[int]							    = {0: "batch", 2: "past_sequence + sequence"}
     else:
       _SCREAMING_SNAKE_CASE   :				Tuple							    = 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
   # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     if self.task in ["default", "seq2seq-lm"]:
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = super().outputs
     else:
       _SCREAMING_SNAKE_CASE   :				str							    = super(__A  ,     self				).outputs
       if self.use_past:
         _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				Dict							    = self.num_layers
         for i in range(__A				):
           _SCREAMING_SNAKE_CASE   :				str							    = {0: "batch", 2: "past_sequence + sequence"}
           _SCREAMING_SNAKE_CASE   :				Optional[int]							    = {0: "batch", 2: "past_sequence + sequence"}
     return common_outputs
   def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = -1  ,     snake_case__ = -1  ,     snake_case__ = False  ,     snake_case__ = None  ,     ):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				Dict							    = self._generate_dummy_inputs_for_encoder_and_decoder(
         __A  ,     __A  ,     __A  ,     __A  ,     __A				)

     # Generate decoder inputs
     _SCREAMING_SNAKE_CASE   :				List[Any]							    = seq_length if not self.use_past else 1
     _SCREAMING_SNAKE_CASE   :				List[Any]							    = self._generate_dummy_inputs_for_encoder_and_decoder(
         __A  ,     __A  ,     __A  ,     __A  ,     __A				)
     _SCREAMING_SNAKE_CASE   :				str							    = {F'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
     _SCREAMING_SNAKE_CASE   :				Any							    = dict(**__A  ,     **__A				)

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

       _SCREAMING_SNAKE_CASE   :				Any							    = torch.cat(
           [common_inputs["decoder_attention_mask"], torch.ones(__A  ,     __A				)]  ,     dim=1				)

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

       for _ in range(__A				):
         common_inputs["past_key_values"].append(
             (
                 torch.zeros(__A				),
                 torch.zeros(__A				),
                 torch.zeros(__A				),
                 torch.zeros(__A				),
             )				)
       # TODO: test this.
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = encoder_shape if remaining_side_name == "encoder" else decoder_shape
       for _ in range(__A  ,     __A				):
         common_inputs["past_key_values"].append((torch.zeros(__A				), torch.zeros(__A				))				)
     return common_inputs
   def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = -1  ,     snake_case__ = -1  ,     snake_case__ = False  ,     snake_case__ = None  ,     ):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = self._generate_dummy_inputs_for_encoder_and_decoder(
         __A  ,     __A  ,     __A  ,     __A  ,     __A				)

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

       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = common_inputs["attention_mask"].dtype
       _SCREAMING_SNAKE_CASE   :				int							    = torch.cat(
           [common_inputs["attention_mask"], torch.ones(__A  ,     __A  ,     dtype=__A				)]  ,     dim=1				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = [
           (torch.zeros(__A				), torch.zeros(__A				)) for _ in range(__A				)
       ]
     return common_inputs
   def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = -1  ,     snake_case__ = -1  ,     snake_case__ = False  ,     snake_case__ = None  ,     ):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				Any							    = compute_effective_axis_dimension(
         __A  ,     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
     _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = tokenizer.num_special_tokens_to_add(__A				)
     _SCREAMING_SNAKE_CASE   :				List[str]							    = compute_effective_axis_dimension(
         __A  ,     fixed_dimension=OnnxConfig.default_fixed_sequence  ,     num_token_to_add=__A				)

     # Generate dummy inputs according to compute batch and sequence
     _SCREAMING_SNAKE_CASE   :				List[Any]							    = [" ".join([tokenizer.unk_token]				) * seq_length] * batch_size
     _SCREAMING_SNAKE_CASE   :				Any							    = dict(tokenizer(__A  ,     return_tensors=__A				)				)
     return common_inputs
   def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = -1  ,     snake_case__ = -1  ,     snake_case__ = False  ,     snake_case__ = None  ,     ):




     """simple docstring"""


     if self.task in ["default", "seq2seq-lm"]:
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self._generate_dummy_inputs_for_default_and_seqaseq_lm(
           __A  ,     batch_size=__A  ,     seq_length=__A  ,     is_pair=__A  ,     framework=__A				)

     else:
       _SCREAMING_SNAKE_CASE   :				Tuple							    = self._generate_dummy_inputs_for_causal_lm(
           __A  ,     batch_size=__A  ,     seq_length=__A  ,     is_pair=__A  ,     framework=__A				)

     return common_inputs
   def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				):




     """simple docstring"""


     if self.task in ["default", "seq2seq-lm"]:
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = super()._flatten_past_key_values_(__A  ,     __A  ,     __A  ,     __A				)
     else:
       _SCREAMING_SNAKE_CASE   :				Dict							    = super(__A  ,     self				)._flatten_past_key_values_(
           __A  ,     __A  ,     __A  ,     __A				)




   @property
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     return 1E-4





"""simple docstring"""
# 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      (							__SCREAMING_SNAKE_CASE  ):
     A__      =			"""naver-clova-ix/donut-base-finetuned-docvqa"""
     A__      =			(
         """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."""
     )
     A__      =			"""document_qa"""
     A__      =			AutoProcessor
     A__      =			VisionEncoderDecoderModel

     A__      =			["""image""", """text"""]
     A__      =			["""text"""]
     def __init__(     self  ,     *snake_case__  ,     **snake_case__				):




       """simple docstring"""


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

       super().__init__(*snake_case__  ,     **snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[str]							    = "<s_docvqa><s_question>{user_input}</s_question><s_answer>"
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = task_prompt.replace("{user_input}"  ,     snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = self.pre_processor.tokenizer(
           snake_case__  ,     add_special_tokens=snake_case__  ,     return_tensors="pt"				).input_ids
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.pre_processor(snake_case__  ,     return_tensors="pt"				).pixel_values

       return {"decoder_input_ids": decoder_input_ids, "pixel_values": pixel_values}
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__				):




       """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=snake_case__  ,     pad_token_id=self.pre_processor.tokenizer.pad_token_id  ,     eos_token_id=self.pre_processor.tokenizer.eos_token_id  ,     use_cache=snake_case__  ,     num_beams=1  ,     bad_words_ids=[[self.pre_processor.tokenizer.unk_token_id]]  ,     return_dict_in_generate=snake_case__  ,     ).sequences




     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = self.pre_processor.batch_decode(snake_case__				)[0]
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = sequence.replace(self.pre_processor.tokenizer.eos_token  ,     ""				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = sequence.replace(self.pre_processor.tokenizer.pad_token  ,     ""				)
       _SCREAMING_SNAKE_CASE   :				Dict							    = re.sub(r"<.*?>"  ,     ""  ,     snake_case__  ,     count=1				).strip()  # remove first task start token
       _SCREAMING_SNAKE_CASE   :				Dict							    = self.pre_processor.tokenajson(snake_case__				)

       return sequence["answer"]



