






# 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"]









from __future__ import annotations

import unittest

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

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


if is_tf_available():
		import numpy
		import tensorflow as tf

		from transformers import (
		    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    BertConfig,
		    DPRConfig,
		    TFDPRContextEncoder,
		    TFDPRQuestionEncoder,
		    TFDPRReader,
		)



class 				__UpperCAmelCase		:







			def __init__(     self:  Optional[int]				,     UpperCAmelCase_:  List[str]				,     UpperCAmelCase_:  int=13				,     UpperCAmelCase_:  List[str]=7				,     UpperCAmelCase_:  Union[str, Any]=True				,     UpperCAmelCase_:  List[str]=True				,     UpperCAmelCase_:  Dict=True				,     UpperCAmelCase_:  str=True				,     UpperCAmelCase_:  List[str]=99				,     UpperCAmelCase_:  int=32				,     UpperCAmelCase_:  List[Any]=2				,     UpperCAmelCase_:  Union[str, Any]=4				,     UpperCAmelCase_:  Dict=37				,     UpperCAmelCase_:  Tuple="gelu"				,     UpperCAmelCase_:  Any=0.1				,     UpperCAmelCase_:  List[str]=0.1				,     UpperCAmelCase_:  List[Any]=512				,     UpperCAmelCase_:  Optional[int]=16				,     UpperCAmelCase_:  Optional[Any]=2				,     UpperCAmelCase_:  Optional[int]=0.02				,     UpperCAmelCase_:  Union[str, Any]=3				,     UpperCAmelCase_:  Optional[int]=4				,     UpperCAmelCase_:  Tuple=None				,     UpperCAmelCase_:  Any=0				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  parent
									_SCREAMING_SNAKE_CASE     				=  batch_size
									_SCREAMING_SNAKE_CASE     				=  seq_length
									_SCREAMING_SNAKE_CASE     				=  is_training
									_SCREAMING_SNAKE_CASE     				=  use_input_mask
									_SCREAMING_SNAKE_CASE     				=  use_token_type_ids
									_SCREAMING_SNAKE_CASE     				=  use_labels
									_SCREAMING_SNAKE_CASE     				=  vocab_size
									_SCREAMING_SNAKE_CASE     				=  hidden_size
									_SCREAMING_SNAKE_CASE     				=  num_hidden_layers
									_SCREAMING_SNAKE_CASE     				=  num_attention_heads
									_SCREAMING_SNAKE_CASE     				=  intermediate_size
									_SCREAMING_SNAKE_CASE     				=  hidden_act
									_SCREAMING_SNAKE_CASE     				=  hidden_dropout_prob
									_SCREAMING_SNAKE_CASE     				=  attention_probs_dropout_prob
									_SCREAMING_SNAKE_CASE     				=  max_position_embeddings
									_SCREAMING_SNAKE_CASE     				=  type_vocab_size
									_SCREAMING_SNAKE_CASE     				=  type_sequence_label_size
									_SCREAMING_SNAKE_CASE     				=  initializer_range
									_SCREAMING_SNAKE_CASE     				=  num_labels
									_SCREAMING_SNAKE_CASE     				=  num_choices
									_SCREAMING_SNAKE_CASE     				=  scope
									_SCREAMING_SNAKE_CASE     				=  projection_dim







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     self.vocab_size					)

									_SCREAMING_SNAKE_CASE     				=  None
									if self.use_input_mask:
															# follow test_modeling_tf_ctrl.py
															_SCREAMING_SNAKE_CASE     				=  random_attention_mask([self.batch_size, self.seq_length]					)

									_SCREAMING_SNAKE_CASE     				=  None
									if self.use_token_type_ids:
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     self.type_vocab_size					)

									_SCREAMING_SNAKE_CASE     				=  None
									_SCREAMING_SNAKE_CASE     				=  None
									_SCREAMING_SNAKE_CASE     				=  None
									if self.use_labels:
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size]				,     self.type_sequence_label_size					)
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     self.num_labels					)
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size]				,     self.num_choices					)

									_SCREAMING_SNAKE_CASE     				=  BertConfig(
									    vocab_size=self.vocab_size				,     hidden_size=self.hidden_size				,     num_hidden_layers=self.num_hidden_layers				,     num_attention_heads=self.num_attention_heads				,     intermediate_size=self.intermediate_size				,     hidden_act=self.hidden_act				,     hidden_dropout_prob=self.hidden_dropout_prob				,     attention_probs_dropout_prob=self.attention_probs_dropout_prob				,     max_position_embeddings=self.max_position_embeddings				,     type_vocab_size=self.type_vocab_size				,     is_decoder=UpperCAmelCase_				,     initializer_range=self.initializer_range				,     )
									_SCREAMING_SNAKE_CASE     				=  DPRConfig(projection_dim=self.projection_dim				,     **config.to_dict()					)

									return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







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







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFDPRContextEncoder(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     attention_mask=UpperCAmelCase_				,     token_type_ids=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     token_type_ids=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.pooler_output.shape				,     (self.batch_size, self.projection_dim or self.hidden_size)					)







			def 							UpperCamelCase				(     self:  List[Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Any				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFDPRQuestionEncoder(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     attention_mask=UpperCAmelCase_				,     token_type_ids=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     token_type_ids=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.pooler_output.shape				,     (self.batch_size, self.projection_dim or self.hidden_size)					)







			def 							UpperCamelCase				(     self:  Optional[Any]				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  Optional[Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFDPRReader(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     attention_mask=UpperCAmelCase_					)

									self.parent.assertEqual(result.start_logits.shape				,     (self.batch_size, self.seq_length)					)
									self.parent.assertEqual(result.end_logits.shape				,     (self.batch_size, self.seq_length)					)
									self.parent.assertEqual(result.relevance_logits.shape				,     (self.batch_size,)					)







			def 							UpperCamelCase				(     self:  Optional[int]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.prepare_config_and_inputs()
									(
									    (
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					
									)     				=  config_and_inputs
									_SCREAMING_SNAKE_CASE     				=  {"""input_ids""": input_ids}
									return config, inputs_dict



@require_tf
class 				__UpperCAmelCase		(_UpperCAmelCase	,_UpperCAmelCase	,unittest.TestCase   ):
			__snake_case    :				Dict									=							(
			    (
			        TFDPRContextEncoder,
			        TFDPRQuestionEncoder,
			        TFDPRReader,
			    )
			    if is_tf_available()
			    else ()
			)
			__snake_case    :				Optional[Any]									=							{"feature-extraction": TFDPRQuestionEncoder} if is_tf_available() else {}

			__snake_case    :				str									=							False
			__snake_case    :				List[Any]									=							False
			__snake_case    :				Any									=							False
			__snake_case    :				List[Any]									=							False
			__snake_case    :				Tuple									=							False







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFDPRModelTester(self					)
									_SCREAMING_SNAKE_CASE     				=  ConfigTester(self				,     config_class=UpperCAmelCase_				,     hidden_size=37					)







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									self.config_tester.run_common_tests()







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_dpr_context_encoder(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_dpr_question_encoder(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_dpr_reader(*UpperCAmelCase_					)







			@slow
			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  TFDPRContextEncoder.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)

									for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  TFDPRContextEncoder.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)

									for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  TFDPRQuestionEncoder.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)

									for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  TFDPRReader.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)



@require_tf
class 				__UpperCAmelCase		(unittest.TestCase   ):







			@slow
			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base"""					)

									_SCREAMING_SNAKE_CASE     				=  tf.constant(
									    [[101, 7_592, 1_010, 2_003, 2_026, 3_899, 10_140, 1_029, 102]]					)  # [CLS] hello, is my dog cute? [SEP]
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)[0]  # embedding shape = (1, 768)
									# compare the actual values for a slice.
									_SCREAMING_SNAKE_CASE     				=  tf.constant(
									    [
									        [
									            0.03_23_62_53,
									            0.12_75_33_35,
									            0.16_81_85_09,
									            0.00_27_97_86,
									            0.3_89_69_33,
									            0.24_26_49_45,
									            0.2_17_89_71,
									            -0.02_33_52_27,
									            -0.08_48_19_59,
									            -0.14_32_41_17,
									        ]
									    ]					)
									self.assertTrue(numpy.allclose(output[:, :10].numpy()				,     expected_slice.numpy()				,     atol=1E-4					)					)

