




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     snake_case__							(_UpperCamelCase   ):
				"""simple docstring"""



				def __init__(      self				:     Optional[Any]					,      __lowerCamelCase				:     List[Any]					,      __lowerCamelCase				:     Optional[int]=13					,      __lowerCamelCase				:     List[str]=7					,      __lowerCamelCase				:     Union[str, Any]=True					,      __lowerCamelCase				:     Dict=True					,      __lowerCamelCase				:     Optional[int]=False					,      __lowerCamelCase				:     str=True					,      __lowerCamelCase				:     List[str]=99					,      __lowerCamelCase				:     List[Any]=32					,      __lowerCamelCase				:     Any=5					,      __lowerCamelCase				:     str=4					,      __lowerCamelCase				:     List[Any]=64					,      __lowerCamelCase				:     Dict="gelu"					,      __lowerCamelCase				:     int=0.1					,      __lowerCamelCase				:     List[str]=0.1					,      __lowerCamelCase				:     Union[str, Any]=5_12					,      __lowerCamelCase				:     List[Any]=16					,      __lowerCamelCase				:     str=2					,      __lowerCamelCase				:     Optional[Any]=0.02					,      __lowerCamelCase				:     Optional[int]=3					,      __lowerCamelCase				:     List[Any]=4					,      __lowerCamelCase				:     Tuple=None					,      __lowerCamelCase				:     List[Any]=2					,      __lowerCamelCase				:     List[Any]=2					,      __lowerCamelCase				:     Optional[Any]=2					,      __lowerCamelCase				:     str=2					,      __lowerCamelCase				:     Dict=4					,      __lowerCamelCase				:     Optional[int]=1					,      )					->     Any:
							a    =			parent
							a    =			batch_size
							a    =			seq_length
							a    =			is_training
							a    =			use_input_mask
							a    =			use_token_type_ids
							a    =			use_labels
							a    =			vocab_size
							a    =			hidden_size
							a    =			num_hidden_layers
							a    =			num_attention_heads
							a    =			intermediate_size
							a    =			hidden_act
							a    =			hidden_dropout_prob
							a    =			attention_probs_dropout_prob
							a    =			max_position_embeddings
							a    =			type_vocab_size
							a    =			type_sequence_label_size
							a    =			initializer_range
							a    =			num_labels
							a    =			num_choices
							a    =			scope
							a    =			q_groups
							a    =			k_groups
							a    =			v_groups
							a    =			post_attention_groups
							a    =			intermediate_groups
							a    =			output_groups





				def 							__UpperCAmelCase							(      self				:     int							)					->     Optional[int]:
							a    =			ids_tensor([self.batch_size, self.seq_length]					,      self.vocab_size							)

							a    =			None
							if self.use_input_mask:
										a    =			random_attention_mask([self.batch_size, self.seq_length]							)

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

							a    =			self.get_config()

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





				def 							__UpperCAmelCase							(      self				:     Dict							)					->     Dict:
							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 							__UpperCAmelCase							(      self				:     str					,      __lowerCamelCase				:     Optional[int]					,      __lowerCamelCase				:     List[Any]					,      __lowerCamelCase				:     Optional[Any]					,      __lowerCamelCase				:     List[Any]					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     Optional[Any]							)					->     str:
							a    =			SqueezeBertModel(config=__lowerCamelCase							)
							model.to(__lowerCamelCase							)
							model.eval()
							a    =			model(__lowerCamelCase					,      __lowerCamelCase							)
							a    =			model(__lowerCamelCase							)
							self.parent.assertEqual(result.last_hidden_state.shape					,      (self.batch_size, self.seq_length, self.hidden_size)							)





				def 							__UpperCAmelCase							(      self				:     List[Any]					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     Any					,      __lowerCamelCase				:     int					,      __lowerCamelCase				:     List[str]					,      __lowerCamelCase				:     List[Any]					,      __lowerCamelCase				:     List[Any]							)					->     Dict:
							a    =			SqueezeBertForMaskedLM(config=__lowerCamelCase							)
							model.to(__lowerCamelCase							)
							model.eval()
							a    =			model(__lowerCamelCase					,      attention_mask=__lowerCamelCase					,      labels=__lowerCamelCase							)
							self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)							)





				def 							__UpperCAmelCase							(      self				:     Optional[Any]					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     Optional[int]					,      __lowerCamelCase				:     Optional[Any]					,      __lowerCamelCase				:     Dict					,      __lowerCamelCase				:     Optional[int]					,      __lowerCamelCase				:     str							)					->     Union[str, Any]:
							a    =			SqueezeBertForQuestionAnswering(config=__lowerCamelCase							)
							model.to(__lowerCamelCase							)
							model.eval()
							a    =			model(
							    __lowerCamelCase					,      attention_mask=__lowerCamelCase					,      start_positions=__lowerCamelCase					,      end_positions=__lowerCamelCase							)
							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 							__UpperCAmelCase							(      self				:     List[str]					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     Tuple					,      __lowerCamelCase				:     Any					,      __lowerCamelCase				:     Optional[int]					,      __lowerCamelCase				:     Any					,      __lowerCamelCase				:     str							)					->     Any:
							a    =			self.num_labels
							a    =			SqueezeBertForSequenceClassification(__lowerCamelCase							)
							model.to(__lowerCamelCase							)
							model.eval()
							a    =			model(__lowerCamelCase					,      attention_mask=__lowerCamelCase					,      labels=__lowerCamelCase							)
							self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.num_labels)							)





				def 							__UpperCAmelCase							(      self				:     List[Any]					,      __lowerCamelCase				:     Tuple					,      __lowerCamelCase				:     int					,      __lowerCamelCase				:     str					,      __lowerCamelCase				:     str					,      __lowerCamelCase				:     Optional[Any]					,      __lowerCamelCase				:     Tuple							)					->     Any:
							a    =			self.num_labels
							a    =			SqueezeBertForTokenClassification(config=__lowerCamelCase							)
							model.to(__lowerCamelCase							)
							model.eval()

							a    =			model(__lowerCamelCase					,      attention_mask=__lowerCamelCase					,      labels=__lowerCamelCase							)
							self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.num_labels)							)





				def 							__UpperCAmelCase							(      self				:     Optional[Any]					,      __lowerCamelCase				:     str					,      __lowerCamelCase				:     int					,      __lowerCamelCase				:     Tuple					,      __lowerCamelCase				:     List[str]					,      __lowerCamelCase				:     Dict					,      __lowerCamelCase				:     Tuple							)					->     Optional[int]:
							a    =			self.num_choices
							a    =			SqueezeBertForMultipleChoice(config=__lowerCamelCase							)
							model.to(__lowerCamelCase							)
							model.eval()
							a    =			input_ids.unsqueeze(1							).expand(-1					,      self.num_choices					,      -1							).contiguous()
							a    =			input_mask.unsqueeze(1							).expand(-1					,      self.num_choices					,      -1							).contiguous()
							a    =			model(
							    __lowerCamelCase					,      attention_mask=__lowerCamelCase					,      labels=__lowerCamelCase					,      )
							self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.num_choices)							)





				def 							__UpperCAmelCase							(      self				:     Any							)					->     Optional[int]:
							a    =			self.prepare_config_and_inputs()
							((a)    ,		(a)    ,		(a)    ,		(a)    ,		(a)    ,		(a))    =			config_and_inputs
							a    =			{"input_ids": input_ids, "attention_mask": input_mask}
							return config, inputs_dict




@require_torch
class     snake_case__							(_UpperCamelCase    ,       _UpperCamelCase    ,       unittest.TestCase   ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_       :							int							 =					(
				    (
				        SqueezeBertModel,
				        SqueezeBertForMaskedLM,
				        SqueezeBertForMultipleChoice,
				        SqueezeBertForQuestionAnswering,
				        SqueezeBertForSequenceClassification,
				        SqueezeBertForTokenClassification,
				    )
				    if is_torch_available()
				    else None
				)
				SCREAMING_SNAKE_CASE_       :							Tuple							 =					(
				    {
				        """feature-extraction""": SqueezeBertModel,
				        """fill-mask""": SqueezeBertForMaskedLM,
				        """question-answering""": SqueezeBertForQuestionAnswering,
				        """text-classification""": SqueezeBertForSequenceClassification,
				        """token-classification""": SqueezeBertForTokenClassification,
				        """zero-shot""": SqueezeBertForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				SCREAMING_SNAKE_CASE_       :							Optional[int]							 =					False
				SCREAMING_SNAKE_CASE_       :							str							 =					True
				SCREAMING_SNAKE_CASE_       :							Optional[int]							 =					False





				def 							__UpperCAmelCase							(      self				:     Tuple							)					->     Any:
							a    =			SqueezeBertModelTester(self							)
							a    =			ConfigTester(self					,      config_class=__lowerCamelCase					,      dim=37							)





				def 							__UpperCAmelCase							(      self				:     int							)					->     int:
							self.config_tester.run_common_tests()





				def 							__UpperCAmelCase							(      self				:     Tuple							)					->     Dict:
							a    =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_model(*__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Optional[int]							)					->     Dict:
							a    =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_for_masked_lm(*__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Optional[Any]							)					->     Optional[Any]:
							a    =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_for_question_answering(*__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Dict							)					->     Any:
							a    =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_for_sequence_classification(*__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Union[str, Any]							)					->     str:
							a    =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_for_token_classification(*__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Optional[Any]							)					->     List[Any]:
							a    =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_for_multiple_choice(*__lowerCamelCase							)





				@slow
				def 							__UpperCAmelCase							(      self				:     Any							)					->     Any:
							for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										a    =			SqueezeBertModel.from_pretrained(__lowerCamelCase							)
										self.assertIsNotNone(__lowerCamelCase							)




@require_sentencepiece
@require_tokenizers
@require_torch
class     snake_case__							(unittest.TestCase   ):
				"""simple docstring"""



				@slow
				def 							__UpperCAmelCase							(      self				:     str							)					->     Union[str, Any]:
							a    =			SqueezeBertForSequenceClassification.from_pretrained("squeezebert/squeezebert-mnli"							)

							a    =			torch.tensor([[1, 2_94_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 13, 15_88, 2]]							)
							a    =			model(__lowerCamelCase							)[0]
							a    =			torch.Size((1, 3)							)
							self.assertEqual(output.shape					,      __lowerCamelCase							)
							a    =			torch.tensor([[0.6_401, -0.0_349, -0.6_041]]							)
							self.assertTrue(torch.allclose(__lowerCamelCase					,      __lowerCamelCase					,      atol=1e-4							)							)














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

					'''simple docstring'''

					_A  :	List[str]     =					'''naver-clova-ix/donut-base-finetuned-docvqa'''
					_A  :	Any     =					(
					    '''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  :	Tuple     =					'''document_qa'''
					_A  :	Dict     =					AutoProcessor
					_A  :	Tuple     =					VisionEncoderDecoderModel

					_A  :	Optional[int]     =					['''image''', '''text''']
					_A  :	Optional[int]     =					['''text''']



					def __init__(     self :		Any						,   *__a :		List[str]						,   **__a :		Any							)	->     Optional[Any]:


												"""simple docstring"""

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

												super().__init__(*__a						,   **__a							)



					def 					lowerCAmelCase (     self :		List[Any]						,   __a :		"Image"						,   __a :		str							)	->     List[str]:


												"""simple docstring"""

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

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



					def 					lowerCAmelCase (     self :		Optional[int]						,   __a :		int							)	->     int:


												"""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 					lowerCAmelCase (     self :		Union[str, Any]						,   __a :		Union[str, Any]							)	->     Union[str, Any]:


												"""simple docstring"""

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

												return sequence["answer"]