

from __future__ import annotations

import unittest

from transformers import DistilBertConfig, 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 tensorflow as tf

     from transformers.models.distilbert.modeling_tf_distilbert import (
         TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFDistilBertForMaskedLM,
         TFDistilBertForMultipleChoice,
         TFDistilBertForQuestionAnswering,
         TFDistilBertForSequenceClassification,
         TFDistilBertForTokenClassification,
         TFDistilBertModel,
     )





class 					__UpperCAmelCase     :



   def __init__(							self							:     int,		__A							:     Union[str, Any],		):
     UpperCAmelCase       :							Optional[Any]  = parent
     UpperCAmelCase       :							str  = 1_3
     UpperCAmelCase       :							Optional[int]  = 7
     UpperCAmelCase       :							Tuple  = True
     UpperCAmelCase       :							str  = True
     UpperCAmelCase       :							Tuple  = False
     UpperCAmelCase       :							List[str]  = True
     UpperCAmelCase       :							Dict  = 9_9
     UpperCAmelCase       :							int  = 3_2
     UpperCAmelCase       :							Any  = 2
     UpperCAmelCase       :							int  = 4
     UpperCAmelCase       :							Optional[Any]  = 3_7
     UpperCAmelCase       :							int  = "gelu"
     UpperCAmelCase       :							int  = 0.1
     UpperCAmelCase       :							int  = 0.1
     UpperCAmelCase       :							Optional[int]  = 5_1_2
     UpperCAmelCase       :							int  = 1_6
     UpperCAmelCase       :							Tuple  = 2
     UpperCAmelCase       :							str  = 0.0_2
     UpperCAmelCase       :							List[str]  = 3
     UpperCAmelCase       :							Optional[Any]  = 4
     UpperCAmelCase       :							List[Any]  = None



   def        __magic_name__    (							self							:     Optional[int]						):
     UpperCAmelCase       :							Any  = ids_tensor([self.batch_size, self.seq_length],		self.vocab_size						)

     UpperCAmelCase       :							Any  = None
     if self.use_input_mask:
       UpperCAmelCase       :							List[str]  = random_attention_mask([self.batch_size, self.seq_length]						)

     UpperCAmelCase       :							List[Any]  = None
     UpperCAmelCase       :							int  = None
     UpperCAmelCase       :							Any  = None
     if self.use_labels:
       UpperCAmelCase       :							List[Any]  = ids_tensor([self.batch_size],		self.type_sequence_label_size						)
       UpperCAmelCase       :							Any  = ids_tensor([self.batch_size, self.seq_length],		self.num_labels						)
       UpperCAmelCase       :							Optional[Any]  = ids_tensor([self.batch_size],		self.num_choices						)

     UpperCAmelCase       :							Optional[int]  = DistilBertConfig(
         vocab_size=self.vocab_size,		dim=self.hidden_size,		n_layers=self.num_hidden_layers,		n_heads=self.num_attention_heads,		hidden_dim=self.intermediate_size,		hidden_act=self.hidden_act,		dropout=self.hidden_dropout_prob,		attention_dropout=self.attention_probs_dropout_prob,		max_position_embeddings=self.max_position_embeddings,		initializer_range=self.initializer_range,		)

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



   def        __magic_name__    (							self							:     int,		__A							:     List[str],		__A							:     List[Any],		__A							:     int,		__A							:     Optional[Any],		__A							:     Dict,		__A							:     List[str]						):
     UpperCAmelCase       :							Optional[Any]  = TFDistilBertModel(config=__UpperCAmelCase						)
     UpperCAmelCase       :							Dict  = {"input_ids": input_ids, "attention_mask": input_mask}

     UpperCAmelCase       :							List[Any]  = model(__UpperCAmelCase						)

     UpperCAmelCase       :							List[Any]  = [input_ids, input_mask]

     UpperCAmelCase       :							Any  = model(__UpperCAmelCase						)

     self.parent.assertEqual(result.last_hidden_state.shape,		(self.batch_size, self.seq_length, self.hidden_size)						)



   def        __magic_name__    (							self							:     str,		__A							:     str,		__A							:     Optional[Any],		__A							:     Union[str, Any],		__A							:     int,		__A							:     Tuple,		__A							:     int						):
     UpperCAmelCase       :							Optional[int]  = TFDistilBertForMaskedLM(config=__UpperCAmelCase						)
     UpperCAmelCase       :							int  = {"input_ids": input_ids, "attention_mask": input_mask}
     UpperCAmelCase       :							int  = model(__UpperCAmelCase						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.vocab_size)						)



   def        __magic_name__    (							self							:     Union[str, Any],		__A							:     Any,		__A							:     Union[str, Any],		__A							:     List[Any],		__A							:     Optional[Any],		__A							:     Optional[Any],		__A							:     str						):
     UpperCAmelCase       :							Optional[int]  = TFDistilBertForQuestionAnswering(config=__UpperCAmelCase						)
     UpperCAmelCase       :							Union[str, Any]  = {
         "input_ids": input_ids,
         "attention_mask": input_mask,
     }
     UpperCAmelCase       :							List[str]  = model(__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)						)



   def        __magic_name__    (							self							:     Tuple,		__A							:     str,		__A							:     Optional[int],		__A							:     Union[str, Any],		__A							:     Dict,		__A							:     Optional[int],		__A							:     Optional[Any]						):
     UpperCAmelCase       :							int  = self.num_labels
     UpperCAmelCase       :							Optional[int]  = TFDistilBertForSequenceClassification(__UpperCAmelCase						)
     UpperCAmelCase       :							Union[str, Any]  = {"input_ids": input_ids, "attention_mask": input_mask}
     UpperCAmelCase       :							Dict  = model(__UpperCAmelCase						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_labels)						)



   def        __magic_name__    (							self							:     Dict,		__A							:     Optional[int],		__A							:     str,		__A							:     str,		__A							:     Any,		__A							:     List[str],		__A							:     List[str]						):
     UpperCAmelCase       :							str  = self.num_choices
     UpperCAmelCase       :							Optional[Any]  = TFDistilBertForMultipleChoice(__UpperCAmelCase						)
     UpperCAmelCase       :							Any  = tf.tile(tf.expand_dims(__UpperCAmelCase,		1						),		(1, self.num_choices, 1)						)
     UpperCAmelCase       :							Union[str, Any]  = tf.tile(tf.expand_dims(__UpperCAmelCase,		1						),		(1, self.num_choices, 1)						)
     UpperCAmelCase       :							str  = {
         "input_ids": multiple_choice_inputs_ids,
         "attention_mask": multiple_choice_input_mask,
     }
     UpperCAmelCase       :							Tuple  = model(__UpperCAmelCase						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_choices)						)



   def        __magic_name__    (							self							:     List[Any],		__A							:     Optional[Any],		__A							:     Tuple,		__A							:     Union[str, Any],		__A							:     Optional[int],		__A							:     Any,		__A							:     List[Any]						):
     UpperCAmelCase       :							Tuple  = self.num_labels
     UpperCAmelCase       :							Dict  = TFDistilBertForTokenClassification(__UpperCAmelCase						)
     UpperCAmelCase       :							int  = {"input_ids": input_ids, "attention_mask": input_mask}
     UpperCAmelCase       :							Tuple  = model(__UpperCAmelCase						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.num_labels)						)





   def        __magic_name__    (							self							:     Any						):
     UpperCAmelCase       :							List[str]  = self.prepare_config_and_inputs()
     (UpperCAmelCase)       :							Optional[Any]  = config_and_inputs
     UpperCAmelCase       :							List[Any]  = {"input_ids": input_ids, "attention_mask": input_mask}
     return config, inputs_dict






@require_tf
class 					__UpperCAmelCase     (       _a			,  _a			,  unittest.TestCase						):
   UpperCamelCase        =     (
       (
           TFDistilBertModel,
           TFDistilBertForMaskedLM,
           TFDistilBertForQuestionAnswering,
           TFDistilBertForSequenceClassification,
           TFDistilBertForTokenClassification,
           TFDistilBertForMultipleChoice,
       )
       if is_tf_available()
       else None
   )
   UpperCamelCase        =     (
       {
           """feature-extraction""": TFDistilBertModel,
           """fill-mask""": TFDistilBertForMaskedLM,
           """question-answering""": TFDistilBertForQuestionAnswering,
           """text-classification""": TFDistilBertForSequenceClassification,
           """token-classification""": TFDistilBertForTokenClassification,
           """zero-shot""": TFDistilBertForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   UpperCamelCase        =     False
   UpperCamelCase        =     False



   def        __magic_name__    (							self							:     str						):
     UpperCAmelCase       :							int  = TFDistilBertModelTester(self						)
     UpperCAmelCase       :							Any  = ConfigTester(self,		config_class=__UpperCAmelCase,		dim=3_7						)



   def        __magic_name__    (							self							:     Optional[Any]						):
     self.config_tester.run_common_tests()



   def        __magic_name__    (							self							:     Tuple						):
     UpperCAmelCase       :							Any  = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_model(*__UpperCAmelCase						)



   def        __magic_name__    (							self							:     Optional[int]						):
     UpperCAmelCase       :							Dict  = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_masked_lm(*__UpperCAmelCase						)



   def        __magic_name__    (							self							:     Tuple						):
     UpperCAmelCase       :							Any  = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_question_answering(*__UpperCAmelCase						)



   def        __magic_name__    (							self							:     Union[str, Any]						):
     UpperCAmelCase       :							Dict  = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_sequence_classification(*__UpperCAmelCase						)



   def        __magic_name__    (							self							:     Tuple						):
     UpperCAmelCase       :							List[Any]  = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_multiple_choice(*__UpperCAmelCase						)



   def        __magic_name__    (							self							:     Optional[int]						):
     UpperCAmelCase       :							List[str]  = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_token_classification(*__UpperCAmelCase						)





   @slow
   def        __magic_name__    (							self							:     Optional[Any]						):
     for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]						):
       UpperCAmelCase       :							Tuple  = TFDistilBertModel.from_pretrained(__UpperCAmelCase						)
       self.assertIsNotNone(__UpperCAmelCase						)






@require_tf
class 					__UpperCAmelCase     (       unittest.TestCase						):



   @slow
   def        __magic_name__    (							self							:     Tuple						):
     UpperCAmelCase       :							List[str]  = TFDistilBertModel.from_pretrained('''distilbert-base-uncased'''						)
     UpperCAmelCase       :							Tuple  = tf.constant([[0, 1, 2, 3, 4, 5]]						)
     UpperCAmelCase       :							List[Any]  = model(__UpperCAmelCase						)[0]

     UpperCAmelCase       :							Union[str, Any]  = [1, 6, 7_6_8]
     self.assertEqual(output.shape,		__UpperCAmelCase						)

     UpperCAmelCase       :							Optional[int]  = tf.constant(
         [
             [
                 [0.1_9_2_6_1_8_8_5, -0.1_3_7_3_2_9_5_5, 0.4_1_1_9_7_9_9],
                 [0.2_2_1_5_0_1_5_6, -0.0_7_4_2_2_6_6_1, 0.3_9_0_3_7_2_0_4],
                 [0.2_2_7_5_6_0_1_8, -0.0_8_9_6_4_1_4, 0.3_7_0_1_4_6_7],
             ]
         ]						)
     tf.debugging.assert_near(output[:, :3, :3],		__UpperCAmelCase,		atol=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 						_A  (  _a  ):







    """simple docstring"""




    UpperCAmelCase       :							str	=       """naver-clova-ix/donut-base-finetuned-docvqa"""
    UpperCAmelCase       :							Tuple	=       (
        """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."""
    )
    UpperCAmelCase       :							List[str]	=       """document_qa"""
    UpperCAmelCase       :							str	=       AutoProcessor
    UpperCAmelCase       :							Optional[int]	=       VisionEncoderDecoderModel

    UpperCAmelCase       :							int	=       ["""image""", """text"""]
    UpperCAmelCase       :							int	=       ["""text"""]






    def __init__(					self				:    Tuple				, *__UpperCAmelCase				:    Union[str, Any]				, **__UpperCAmelCase				:    Any):
      if not is_vision_available():
        raise ValueError("Pillow must be installed to use the DocumentQuestionAnsweringTool.")

      super().__init__(*__UpperCAmelCase				, **__UpperCAmelCase)






    def 	__snake_case							(					self				:    Tuple				, __UpperCAmelCase				:    "Image"				, __UpperCAmelCase				:    str):
      a						:		Any      							=			"<s_docvqa><s_question>{user_input}</s_question><s_answer>"
      a						:		Union[str, Any]      							=			task_prompt.replace("{user_input}"				, __UpperCAmelCase)
      a						:		Optional[Any]      							=			self.pre_processor.tokenizer(
          __UpperCAmelCase				, add_special_tokens=__UpperCAmelCase				, return_tensors="pt").input_ids
      a						:		Any      							=			self.pre_processor(__UpperCAmelCase				, return_tensors="pt").pixel_values

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






    def 	__snake_case							(					self				:    int				, __UpperCAmelCase				:    int):
      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 	__snake_case							(					self				:    str				, __UpperCAmelCase				:    List[Any]):
      a						:		Union[str, Any]      							=			self.pre_processor.batch_decode(__UpperCAmelCase)[0]
      a						:		Optional[Any]      							=			sequence.replace(self.pre_processor.tokenizer.eos_token				, "")
      a						:		Any      							=			sequence.replace(self.pre_processor.tokenizer.pad_token				, "")
      a						:		Optional[Any]      							=			re.sub(r"<.*?>"				, ""				, __UpperCAmelCase				, count=1).strip()  # remove first task start token
      a						:		List[str]      							=			self.pre_processor.tokenajson(__UpperCAmelCase)

      return sequence["answer"]






