
# 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 (							lowerCAmelCase__       ):






     """simple docstring"""







     lowerCAmelCase_							   =			'''naver-clova-ix/donut-base-finetuned-docvqa'''
     lowerCAmelCase_							   =			(
         '''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.'''
     )
     lowerCAmelCase_							   =			'''document_qa'''
     lowerCAmelCase_							   =			AutoProcessor
     lowerCAmelCase_							   =			VisionEncoderDecoderModel

     lowerCAmelCase_							   =			['''image''', '''text''']
     lowerCAmelCase_							   =			['''text''']



     def __init__(				self							:							Tuple     ,						*_A							:							Tuple     ,						**_A							:							Tuple ):







       """simple docstring"""




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

       super().__init__(*_A     ,						**_A )



     def  UpperCAmelCase__							(				self							:							List[Any]     ,						_A							:							"Image"     ,						_A							:							str ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		'''<s_docvqa><s_question>{user_input}</s_question><s_answer>'''
       __SCREAMING_SNAKE_CASE      :	str					=		task_prompt.replace('''{user_input}'''     ,						_A )
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		self.pre_processor.tokenizer(
           _A     ,						add_special_tokens=_A     ,						return_tensors='''pt''' ).input_ids
       __SCREAMING_SNAKE_CASE      :	int					=		self.pre_processor(_A     ,						return_tensors='''pt''' ).pixel_values

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



     def  UpperCAmelCase__							(				self							:							Dict     ,						_A							:							Optional[Any] ):







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




     def  UpperCAmelCase__							(				self							:							Dict     ,						_A							:							Dict ):







       """simple docstring"""




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

       return sequence["answer"]




import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import is_speech_available, is_vision_available
from transformers.testing_utils import require_torch


if is_vision_available():
   from transformers import TvltImageProcessor

if is_speech_available():
   from transformers import TvltFeatureExtractor

from transformers import TvltProcessor
@require_torch
class 		__magic_name__						(unittest.TestCase				):





     '''simple docstring'''



     def 	SCREAMING_SNAKE_CASE__						(					self:str   ):
       snake_case__						= '''ZinengTang/tvlt-base'''
       snake_case__						= tempfile.mkdtemp()

     def 	SCREAMING_SNAKE_CASE__						(					self:Dict     ,		**_a:List[Any]   ):
       return TvltImageProcessor.from_pretrained(self.checkpoint     ,		**_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[Any]     ,		**_a:Tuple   ):
       return TvltFeatureExtractor.from_pretrained(self.checkpoint     ,		**_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Dict   ):
       shutil.rmtree(self.tmpdirname   )

     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]   ):
       snake_case__						= self.get_image_processor()
       snake_case__						= self.get_feature_extractor()

       snake_case__						= TvltProcessor(image_processor=_a     ,		feature_extractor=_a   )
       processor.save_pretrained(self.tmpdirname   )
       snake_case__						= TvltProcessor.from_pretrained(self.tmpdirname   )

       self.assertIsInstance(processor.feature_extractor     ,		_a   )
       self.assertIsInstance(processor.image_processor     ,		_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Dict   ):
       snake_case__						= self.get_image_processor()
       snake_case__						= self.get_feature_extractor()

       snake_case__						= TvltProcessor(image_processor=_a     ,		feature_extractor=_a   )

       snake_case__						= np.ones([1_20_00]   )

       snake_case__						= feature_extractor(_a     ,		return_tensors='''np'''   )
       snake_case__						= processor(audio=_a     ,		return_tensors='''np'''   )

       for key in audio_dict.keys():
         self.assertAlmostEqual(audio_dict[key].sum()     ,		input_processor[key].sum()     ,		delta=1e-2   )

     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       snake_case__						= self.get_image_processor()
       snake_case__						= self.get_feature_extractor()

       snake_case__						= TvltProcessor(image_processor=_a     ,		feature_extractor=_a   )

       snake_case__						= np.ones([3, 2_24, 2_24]   )

       snake_case__						= image_processor(_a     ,		return_tensors='''np'''   )
       snake_case__						= processor(images=_a     ,		return_tensors='''np'''   )

       for key in image_dict.keys():
         self.assertAlmostEqual(image_dict[key].sum()     ,		input_processor[key].sum()     ,		delta=1e-2   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Any   ):
       snake_case__						= self.get_image_processor()
       snake_case__						= self.get_feature_extractor()

       snake_case__						= TvltProcessor(image_processor=_a     ,		feature_extractor=_a   )

       snake_case__						= np.ones([1_20_00]   )
       snake_case__						= np.ones([3, 2_24, 2_24]   )

       snake_case__						= processor(audio=_a     ,		images=_a   )

       self.assertListEqual(list(inputs.keys()   )     ,		['''audio_values''', '''audio_mask''', '''pixel_values''', '''pixel_mask''']   )

       # test if it raises when no input is passed
       with pytest.raises(_a   ):
         processor()







     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       snake_case__						= self.get_image_processor()
       snake_case__						= self.get_feature_extractor()

       snake_case__						= TvltProcessor(image_processor=_a     ,		feature_extractor=_a   )

       self.assertListEqual(
           processor.model_input_names     ,		image_processor.model_input_names + feature_extractor.model_input_names     ,		msg='''`processor` and `image_processor`+`feature_extractor` model input names do not match'''     ,		)

