
"""simple docstring"""
import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          SqueezeBertForMaskedLM,
          SqueezeBertForMultipleChoice,
          SqueezeBertForQuestionAnswering,
          SqueezeBertForSequenceClassification,
          SqueezeBertForTokenClassification,
          SqueezeBertModel,
      )
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE  ):
   def __init__(     self  ,     snake_case__  ,     snake_case__=13  ,     snake_case__=7  ,     snake_case__=True  ,     snake_case__=True  ,     snake_case__=False  ,     snake_case__=True  ,     snake_case__=99  ,     snake_case__=32  ,     snake_case__=5  ,     snake_case__=4  ,     snake_case__=64  ,     snake_case__="gelu"  ,     snake_case__=0.1  ,     snake_case__=0.1  ,     snake_case__=512  ,     snake_case__=16  ,     snake_case__=2  ,     snake_case__=0.02  ,     snake_case__=3  ,     snake_case__=4  ,     snake_case__=None  ,     snake_case__=2  ,     snake_case__=2  ,     snake_case__=2  ,     snake_case__=2  ,     snake_case__=4  ,     snake_case__=1  ,     ):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				str							    = parent
     _SCREAMING_SNAKE_CASE   :				Dict							    = batch_size
     _SCREAMING_SNAKE_CASE   :				List[str]							    = seq_length
     _SCREAMING_SNAKE_CASE   :				List[str]							    = is_training
     _SCREAMING_SNAKE_CASE   :				Optional[int]							    = use_input_mask
     _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = use_token_type_ids
     _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = use_labels
     _SCREAMING_SNAKE_CASE   :				List[Any]							    = vocab_size
     _SCREAMING_SNAKE_CASE   :				Any							    = hidden_size
     _SCREAMING_SNAKE_CASE   :				Any							    = num_hidden_layers
     _SCREAMING_SNAKE_CASE   :				str							    = num_attention_heads
     _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = intermediate_size
     _SCREAMING_SNAKE_CASE   :				List[str]							    = hidden_act
     _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = hidden_dropout_prob
     _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = attention_probs_dropout_prob
     _SCREAMING_SNAKE_CASE   :				int							    = max_position_embeddings
     _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = type_vocab_size
     _SCREAMING_SNAKE_CASE   :				List[Any]							    = type_sequence_label_size
     _SCREAMING_SNAKE_CASE   :				Tuple							    = initializer_range
     _SCREAMING_SNAKE_CASE   :				Any							    = num_labels
     _SCREAMING_SNAKE_CASE   :				Tuple							    = num_choices
     _SCREAMING_SNAKE_CASE   :				Dict							    = scope
     _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = q_groups
     _SCREAMING_SNAKE_CASE   :				Tuple							    = k_groups
     _SCREAMING_SNAKE_CASE   :				str							    = v_groups
     _SCREAMING_SNAKE_CASE   :				Any							    = post_attention_groups
     _SCREAMING_SNAKE_CASE   :				Any							    = intermediate_groups
     _SCREAMING_SNAKE_CASE   :				int							    = output_groups
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = ids_tensor([self.batch_size, self.seq_length]  ,     self.vocab_size				)

     _SCREAMING_SNAKE_CASE   :				str							    = None
     if self.use_input_mask:
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = random_attention_mask([self.batch_size, self.seq_length]				)

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

     _SCREAMING_SNAKE_CASE   :				Tuple							    = self.get_config()

     return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     return SqueezeBertConfig(
         embedding_size=self.hidden_size  ,     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  ,     attention_probs_dropout_prob=self.hidden_dropout_prob  ,     attention_dropout=self.attention_probs_dropout_prob  ,     max_position_embeddings=self.max_position_embeddings  ,     initializer_range=self.initializer_range  ,     q_groups=self.q_groups  ,     k_groups=self.k_groups  ,     v_groups=self.v_groups  ,     post_attention_groups=self.post_attention_groups  ,     intermediate_groups=self.intermediate_groups  ,     output_groups=self.output_groups  ,     )
   def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				List[Any]							    = SqueezeBertModel(config=snake_case__				)
     model.to(snake_case__				)
     model.eval()
     _SCREAMING_SNAKE_CASE   :				List[Any]							    = model(snake_case__  ,     snake_case__				)
     _SCREAMING_SNAKE_CASE   :				str							    = model(snake_case__				)
     self.parent.assertEqual(result.last_hidden_state.shape  ,     (self.batch_size, self.seq_length, self.hidden_size)				)
   def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				List[str]							    = SqueezeBertForMaskedLM(config=snake_case__				)
     model.to(snake_case__				)
     model.eval()
     _SCREAMING_SNAKE_CASE   :				Any							    = model(snake_case__  ,     attention_mask=snake_case__  ,     labels=snake_case__				)
     self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.vocab_size)				)
   def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				str							    = SqueezeBertForQuestionAnswering(config=snake_case__				)
     model.to(snake_case__				)
     model.eval()
     _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = model(
         snake_case__  ,     attention_mask=snake_case__  ,     start_positions=snake_case__  ,     end_positions=snake_case__				)
     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 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = self.num_labels
     _SCREAMING_SNAKE_CASE   :				List[str]							    = SqueezeBertForSequenceClassification(snake_case__				)
     model.to(snake_case__				)
     model.eval()
     _SCREAMING_SNAKE_CASE   :				int							    = model(snake_case__  ,     attention_mask=snake_case__  ,     labels=snake_case__				)
     self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.num_labels)				)
   def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.num_labels
     _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = SqueezeBertForTokenClassification(config=snake_case__				)
     model.to(snake_case__				)
     model.eval()

     _SCREAMING_SNAKE_CASE   :				List[str]							    = model(snake_case__  ,     attention_mask=snake_case__  ,     labels=snake_case__				)
     self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.num_labels)				)
   def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				str							    = self.num_choices
     _SCREAMING_SNAKE_CASE   :				List[str]							    = SqueezeBertForMultipleChoice(config=snake_case__				)
     model.to(snake_case__				)
     model.eval()
     _SCREAMING_SNAKE_CASE   :				str							    = input_ids.unsqueeze(1				).expand(-1  ,     self.num_choices  ,     -1				).contiguous()
     _SCREAMING_SNAKE_CASE   :				Dict							    = input_mask.unsqueeze(1				).expand(-1  ,     self.num_choices  ,     -1				).contiguous()
     _SCREAMING_SNAKE_CASE   :				List[str]							    = model(
         snake_case__  ,     attention_mask=snake_case__  ,     labels=snake_case__  ,     )
     self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.num_choices)				)




   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				List[Any]							    = self.prepare_config_and_inputs()
     (_SCREAMING_SNAKE_CASE)   :				Any							    = config_and_inputs
     _SCREAMING_SNAKE_CASE   :				List[str]							    = {"input_ids": input_ids, "attention_mask": input_mask}
     return config, inputs_dict
@require_torch
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE      ,       __SCREAMING_SNAKE_CASE      ,       unittest.TestCase  ):
   A__      =			(
       (
           SqueezeBertModel,
           SqueezeBertForMaskedLM,
           SqueezeBertForMultipleChoice,
           SqueezeBertForQuestionAnswering,
           SqueezeBertForSequenceClassification,
           SqueezeBertForTokenClassification,
       )
       if is_torch_available()
       else None
   )
   A__      =			(
       {
           """feature-extraction""": SqueezeBertModel,
           """fill-mask""": SqueezeBertForMaskedLM,
           """question-answering""": SqueezeBertForQuestionAnswering,
           """text-classification""": SqueezeBertForSequenceClassification,
           """token-classification""": SqueezeBertForTokenClassification,
           """zero-shot""": SqueezeBertForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   A__      =			False
   A__      =			True
   A__      =			False
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				Dict							    = SqueezeBertModelTester(self				)
     _SCREAMING_SNAKE_CASE   :				int							    = ConfigTester(self  ,     config_class=snake_case__  ,     dim=37				)
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     self.config_tester.run_common_tests()
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				List[Any]							    = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_model(*snake_case__				)
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				int							    = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_for_masked_lm(*snake_case__				)
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				List[Any]							    = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_for_question_answering(*snake_case__				)
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				int							    = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_for_sequence_classification(*snake_case__				)
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				Optional[int]							    = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_for_token_classification(*snake_case__				)
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_for_multiple_choice(*snake_case__				)




   @slow
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _SCREAMING_SNAKE_CASE   :				List[str]							    = SqueezeBertModel.from_pretrained(snake_case__				)
       self.assertIsNotNone(snake_case__				)
@require_sentencepiece
@require_tokenizers
@require_torch
class 			UpperCamelCase      (							unittest.TestCase  ):
   @slow
   def 				__SCREAMING_SNAKE_CASE							(     self				):




     """simple docstring"""


     _SCREAMING_SNAKE_CASE   :				List[str]							    = SqueezeBertForSequenceClassification.from_pretrained("squeezebert/squeezebert-mnli"				)

     _SCREAMING_SNAKE_CASE   :				int							    = torch.tensor([[1, 29414, 232, 328, 740, 1140, 12695, 69, 13, 1588, 2]]				)
     _SCREAMING_SNAKE_CASE   :				Optional[int]							    = model(snake_case__				)[0]
     _SCREAMING_SNAKE_CASE   :				List[Any]							    = torch.Size((1, 3)				)
     self.assertEqual(output.shape  ,     snake_case__				)
     _SCREAMING_SNAKE_CASE   :				Optional[int]							    = torch.tensor([[0.6_401, -0.0_349, -0.6_041]]				)
     self.assertTrue(torch.allclose(snake_case__  ,     snake_case__  ,     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 			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"]



