


import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    MODEL_FOR_PRETRAINING_MAPPING,
							    NezhaForMaskedLM,
							    NezhaForMultipleChoice,
							    NezhaForNextSentencePrediction,
							    NezhaForPreTraining,
							    NezhaForQuestionAnswering,
							    NezhaForSequenceClassification,
							    NezhaForTokenClassification,
							    NezhaModel,
							)
							from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST







class        UpperCamelCase__   :

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






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Union[str, Any]       ):






										'''simple docstring'''






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

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

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

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

										lowercase_  							=			self.get_config()

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

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										return NezhaConfig(
										    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							,					)

					def   UpperCAmelCase__       (     self       :   str       ):






										'''simple docstring'''






										(
										    (
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,										
										)  							=			self.prepare_config_and_inputs()

										lowercase_  							=			True
										lowercase_  							=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
										lowercase_  							=			ids_tensor([self.batch_size, self.seq_length]							,					vocab_size=2       )

										return (
										    config,
										    input_ids,
										    token_type_ids,
										    input_mask,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										    encoder_hidden_states,
										    encoder_attention_mask,
										)

					def   UpperCAmelCase__       (     self       :   str							,					UpperCamelCase__       :   List[Any]							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   Optional[int]       ):






										'''simple docstring'''






										lowercase_  							=			NezhaModel(config=UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()
										lowercase_  							=			model(UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__       )
										lowercase_  							=			model(UpperCamelCase__							,					token_type_ids=UpperCamelCase__       )
										lowercase_  							=			model(UpperCamelCase__       )
										self.parent.assertEqual(result.last_hidden_state.shape							,					(self.batch_size, self.seq_length, self.hidden_size)       )
										self.parent.assertEqual(result.pooler_output.shape							,					(self.batch_size, self.hidden_size)       )

					def   UpperCAmelCase__       (     self       :   int							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   List[Any]							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Dict							,					):






										'''simple docstring'''






										lowercase_  							=			True
										lowercase_  							=			NezhaModel(UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()
										lowercase_  							=			model(
										    UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__							,					encoder_hidden_states=UpperCamelCase__							,					encoder_attention_mask=UpperCamelCase__							,					)
										lowercase_  							=			model(
										    UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__							,					encoder_hidden_states=UpperCamelCase__							,					)
										lowercase_  							=			model(UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__       )
										self.parent.assertEqual(result.last_hidden_state.shape							,					(self.batch_size, self.seq_length, self.hidden_size)       )
										self.parent.assertEqual(result.pooler_output.shape							,					(self.batch_size, self.hidden_size)       )

					def   UpperCAmelCase__       (     self       :   Optional[int]							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   List[Any]							,					UpperCamelCase__       :   Any							,					UpperCamelCase__       :   Union[str, Any]       ):






										'''simple docstring'''






										lowercase_  							=			NezhaForMaskedLM(config=UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()
										lowercase_  							=			model(UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__							,					labels=UpperCamelCase__       )
										self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.vocab_size)       )

					def   UpperCAmelCase__       (     self       :   Any							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   Any							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   Optional[int]       ):






										'''simple docstring'''






										lowercase_  							=			NezhaForNextSentencePrediction(config=UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()
										lowercase_  							=			model(
										    UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__							,					labels=UpperCamelCase__							,					)
										self.parent.assertEqual(result.logits.shape							,					(self.batch_size, 2)       )

					def   UpperCAmelCase__       (     self       :   Any							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   List[Any]							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			NezhaForPreTraining(config=UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()
										lowercase_  							=			model(
										    UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__							,					labels=UpperCamelCase__							,					next_sentence_label=UpperCamelCase__							,					)
										self.parent.assertEqual(result.prediction_logits.shape							,					(self.batch_size, self.seq_length, self.vocab_size)       )
										self.parent.assertEqual(result.seq_relationship_logits.shape							,					(self.batch_size, 2)       )

					def   UpperCAmelCase__       (     self       :   Optional[int]							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Dict       ):






										'''simple docstring'''






										lowercase_  							=			NezhaForQuestionAnswering(config=UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()
										lowercase_  							=			model(
										    UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__							,					start_positions=UpperCamelCase__							,					end_positions=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   UpperCAmelCase__       (     self       :   Any							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   List[Any]							,					UpperCamelCase__       :   int       ):






										'''simple docstring'''






										lowercase_  							=			self.num_labels
										lowercase_  							=			NezhaForSequenceClassification(UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()
										lowercase_  							=			model(UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__							,					labels=UpperCamelCase__       )
										self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.num_labels)       )

					def   UpperCAmelCase__       (     self       :   Union[str, Any]							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   List[Any]							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			self.num_labels
										lowercase_  							=			NezhaForTokenClassification(config=UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()
										lowercase_  							=			model(UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__							,					labels=UpperCamelCase__       )
										self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.num_labels)       )

					def   UpperCAmelCase__       (     self       :   Any							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   List[str]       ):






										'''simple docstring'''






										lowercase_  							=			self.num_choices
										lowercase_  							=			NezhaForMultipleChoice(config=UpperCamelCase__       )
										model.to(UpperCamelCase__       )
										model.eval()
										lowercase_  							=			input_ids.unsqueeze(1       ).expand(-1							,					self.num_choices							,					-1       ).contiguous()
										lowercase_  							=			token_type_ids.unsqueeze(1       ).expand(-1							,					self.num_choices							,					-1       ).contiguous()
										lowercase_  							=			input_mask.unsqueeze(1       ).expand(-1							,					self.num_choices							,					-1       ).contiguous()
										lowercase_  							=			model(
										    UpperCamelCase__							,					attention_mask=UpperCamelCase__							,					token_type_ids=UpperCamelCase__							,					labels=UpperCamelCase__							,					)
										self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.num_choices)       )







					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			self.prepare_config_and_inputs()
										(
										    (
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,										
										)  							=			config_and_inputs
										lowercase_  							=			{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
										return config, inputs_dict







@require_torch
class        UpperCamelCase__   ( __magic_name__							,				__magic_name__							,				__magic_name__							,				unittest.TestCase   ):
					__SCREAMING_SNAKE_CASE       :  Optional[int]						=   (
					    (
					        NezhaModel,
					        NezhaForMaskedLM,
					        NezhaForMultipleChoice,
					        NezhaForNextSentencePrediction,
					        NezhaForPreTraining,
					        NezhaForQuestionAnswering,
					        NezhaForSequenceClassification,
					        NezhaForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					__SCREAMING_SNAKE_CASE       :  List[Any]						=   (
					    {
					        'feature-extraction': NezhaModel,
					        'fill-mask': NezhaForMaskedLM,
					        'question-answering': NezhaForQuestionAnswering,
					        'text-classification': NezhaForSequenceClassification,
					        'token-classification': NezhaForTokenClassification,
					        'zero-shot': NezhaForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					__SCREAMING_SNAKE_CASE       :  Dict						=   True

					def   UpperCAmelCase__       (     self       :   Optional[int]							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   List[Any]=False       ):






										'''simple docstring'''






										lowercase_  							=			super()._prepare_for_class(UpperCamelCase__							,					UpperCamelCase__							,					return_labels=UpperCamelCase__       )

										if return_labels:
															if model_class in get_values(UpperCamelCase__       ):
																				lowercase_  							=			torch.zeros(
																				    (self.model_tester.batch_size, self.model_tester.seq_length)							,					dtype=torch.long							,					device=UpperCamelCase__       )
																				lowercase_  							=			torch.zeros(
																				    self.model_tester.batch_size							,					dtype=torch.long							,					device=UpperCamelCase__       )
										return inputs_dict

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			NezhaModelTester(self       )
										lowercase_  							=			ConfigTester(self							,					config_class=UpperCamelCase__							,					hidden_size=37       )

					def   UpperCAmelCase__       (     self       :   str       ):






										'''simple docstring'''






										self.config_tester.run_common_tests()

					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






										lowercase_  							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   str       ):






										'''simple docstring'''






										lowercase_  							=			self.model_tester.prepare_config_and_inputs_for_decoder()
										self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										(
										    (
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,					(
										    lowercase_					
										)							,										
										)  							=			self.model_tester.prepare_config_and_inputs_for_decoder()

										lowercase_  							=			None

										self.model_tester.create_and_check_model_as_decoder(
										    UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					)

					def   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''simple docstring'''






										lowercase_  							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






										lowercase_  							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_next_sequence_prediction(*UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






										lowercase_  							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_pretraining(*UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										lowercase_  							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										lowercase_  							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__       )

					@slow
					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															lowercase_  							=			NezhaModel.from_pretrained(UpperCamelCase__       )
															self.assertIsNotNone(UpperCamelCase__       )







					@slow
					@require_torch_gpu
					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										lowercase_							,					lowercase_  							=			self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes:
															# NezhaForMultipleChoice behaves incorrectly in JIT environments.
															if model_class == NezhaForMultipleChoice:
																				return

															lowercase_  							=			True
															lowercase_  							=			model_class(config=UpperCamelCase__       )

															lowercase_  							=			self._prepare_for_class(UpperCamelCase__							,					UpperCamelCase__       )
															lowercase_  							=			torch.jit.trace(
															    UpperCamelCase__							,					(inputs_dict["""input_ids"""].to("""cpu"""       ), inputs_dict["""attention_mask"""].to("""cpu"""       ))       )

															with tempfile.TemporaryDirectory() as tmp:
																				torch.jit.save(UpperCamelCase__							,					os.path.join(UpperCamelCase__							,					"""bert.pt"""       )       )
																				lowercase_  							=			torch.jit.load(os.path.join(UpperCamelCase__							,					"""bert.pt"""       )							,					map_location=UpperCamelCase__       )
																				loaded(inputs_dict["""input_ids"""].to(UpperCamelCase__       )							,					inputs_dict["""attention_mask"""].to(UpperCamelCase__       )       )







@require_torch
class        UpperCamelCase__   ( unittest.TestCase   ):

					@slow
					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			NezhaModel.from_pretrained("""sijunhe/nezha-cn-base"""       )
										lowercase_  							=			torch.tensor([[0, 1, 2, 3, 4, 5]]       )
										lowercase_  							=			torch.tensor([[0, 1, 1, 1, 1, 1]]       )
										with torch.no_grad():
															lowercase_  							=			model(UpperCamelCase__							,					attention_mask=UpperCamelCase__       )[0]
										lowercase_  							=			torch.Size((1, 6, 768)       )
										self.assertEqual(output.shape							,					UpperCamelCase__       )
										lowercase_  							=			torch.tensor([[[0.0_685, 0.2_441, 0.1_102], [0.0_600, 0.1_906, 0.1_349], [0.0_221, 0.0_819, 0.0_586]]]       )

										self.assertTrue(torch.allclose(output[:, 1:4, 1:4]							,					UpperCamelCase__							,					atol=1e-4       )       )







					@slow
					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






										lowercase_  							=			NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""       )
										lowercase_  							=			torch.tensor([[0, 1, 2, 3, 4, 5]]       )
										lowercase_  							=			torch.tensor([[1, 1, 1, 1, 1, 1]]       )
										with torch.no_grad():
															lowercase_  							=			model(UpperCamelCase__							,					attention_mask=UpperCamelCase__       )[0]
										lowercase_  							=			torch.Size((1, 6, 21_128)       )
										self.assertEqual(output.shape							,					UpperCamelCase__       )
										lowercase_  							=			torch.tensor(
										    [[-2.7_939, -1.7_902, -2.2_189], [-2.8_585, -1.8_908, -2.3_723], [-2.6_499, -1.7_750, -2.2_558]]       )

										self.assertTrue(torch.allclose(output[:, 1:4, 1:4]							,					UpperCamelCase__							,					atol=1e-4       )       )






import io
import math
from typing import Dict, Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import convert_to_rgb, normalize, to_channel_dimension_format, to_pil_image
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    get_image_size,
    infer_channel_dimension_format,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_vision_available, logging
from ...utils.import_utils import requires_backends


if is_vision_available():
							import textwrap

							from PIL import Image, ImageDraw, ImageFont

if is_torch_available():
							import torch

							from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_11
else:
							a					  =     False


a					  =     logging.get_logger(__name__)
a					  =     'ybelkada/fonts'


def 			UpperCAmelCase_				(      ):
					if is_torch_available() and not is_torch_greater_or_equal_than_1_11:
										raise ImportError(
										    F'''You are using torch=={torch.__version__}, but torch>=1.11.0 is required to use '''
										    """Pix2StructImageProcessor. Please upgrade torch."""  )


def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__  ):
					requires_backends(UpperCAmelCase__						,							["""torch"""]  )
					_check_torch_version()

					lowercase_  							=			image_tensor.unsqueeze(0  )
					lowercase_  							=			torch.nn.functional.unfold(UpperCAmelCase__						,							(patch_height, patch_width)						,							stride=(patch_height, patch_width)  )
					lowercase_  							=			patches.reshape(image_tensor.size(0  )						,							image_tensor.size(1  )						,							UpperCAmelCase__						,							UpperCAmelCase__						,							-1  )
					lowercase_  							=			patches.permute(0						,							4						,							2						,							3						,							1  ).reshape(
					    image_tensor.size(2  ) // patch_height						,							image_tensor.size(3  ) // patch_width						,							image_tensor.size(1  ) * patch_height * patch_width						,							)
					return patches.unsqueeze(0  )


def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__ = 3_6						,							UpperCAmelCase__ = "black"						,							UpperCAmelCase__ = "white"						,							UpperCAmelCase__ = 5						,							UpperCAmelCase__ = 5						,							UpperCAmelCase__ = 5						,							UpperCAmelCase__ = 5						,							UpperCAmelCase__ = None						,							UpperCAmelCase__ = None						,							):
					requires_backends(UpperCAmelCase__						,							"""vision"""  )
					# Add new lines so that each line is no more than 80 characters.

					lowercase_  							=			textwrap.TextWrapper(width=8_0  )
					lowercase_  							=			wrapper.wrap(text=UpperCAmelCase__  )
					lowercase_  							=			"""\n""".join(UpperCAmelCase__  )

					if font_bytes is not None and font_path is None:
										lowercase_  							=			io.BytesIO(UpperCAmelCase__  )
					elif font_path is not None:
										lowercase_  							=			font_path
					else:
										lowercase_  							=			hf_hub_download(UpperCAmelCase__						,							"""Arial.TTF"""  )
					lowercase_  							=			ImageFont.truetype(UpperCAmelCase__						,							encoding="""UTF-8"""						,							size=UpperCAmelCase__  )

					# Use a temporary canvas to determine the width and height in pixels when
					# rendering the text.
					lowercase_  							=			ImageDraw.Draw(Image.new("""RGB"""						,							(1, 1)						,							UpperCAmelCase__  )  )
					lowercase_							,					lowercase_							,					lowercase_							,					lowercase_  							=			temp_draw.textbbox((0, 0)						,							UpperCAmelCase__						,							UpperCAmelCase__  )

					# Create the actual image with a bit of padding around the text.
					lowercase_  							=			text_width + left_padding + right_padding
					lowercase_  							=			text_height + top_padding + bottom_padding
					lowercase_  							=			Image.new("""RGB"""						,							(image_width, image_height)						,							UpperCAmelCase__  )
					lowercase_  							=			ImageDraw.Draw(UpperCAmelCase__  )
					draw.text(xy=(left_padding, top_padding)						,							text=UpperCAmelCase__						,							fill=UpperCAmelCase__						,							font=UpperCAmelCase__  )
					return image


def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							**UpperCAmelCase__  ):
					requires_backends(UpperCAmelCase__						,							"""vision"""  )

					# Convert to PIL image if necessary
					lowercase_  							=			to_pil_image(UpperCAmelCase__  )

					lowercase_  							=			render_text(UpperCAmelCase__						,							**UpperCAmelCase__  )
					lowercase_  							=			max(header_image.width						,							image.width  )

					lowercase_  							=			int(image.height * (new_width / image.width)  )
					lowercase_  							=			int(header_image.height * (new_width / header_image.width)  )

					lowercase_  							=			Image.new("""RGB"""						,							(new_width, new_height + new_header_height)						,							"""white"""  )
					new_image.paste(header_image.resize((new_width, new_header_height)  )						,							(0, 0)  )
					new_image.paste(image.resize((new_width, new_height)  )						,							(0, new_header_height)  )

					# Convert back to the original framework if necessary
					lowercase_  							=			to_numpy_array(UpperCAmelCase__  )

					if infer_channel_dimension_format(UpperCAmelCase__  ) == ChannelDimension.LAST:
										lowercase_  							=			to_channel_dimension_format(UpperCAmelCase__						,							ChannelDimension.LAST  )

					return new_image







class        UpperCamelCase__   (	__magic_name__   ):
					__SCREAMING_SNAKE_CASE       :  Tuple						=   ['flattened_patches']

					def __init__(     self       :   str							,					UpperCamelCase__       :   bool = True							,					UpperCamelCase__       :   bool = True							,					UpperCamelCase__       :   Dict[str, int] = None							,					UpperCamelCase__       :   int = 2_048							,					UpperCamelCase__       :   bool = False							,					**UpperCamelCase__       :   Optional[int]							,					):






										'''simple docstring'''






										super().__init__(**UpperCamelCase__       )
										lowercase_  							=			patch_size if patch_size is not None else {"""height""": 16, """width""": 16}
										lowercase_  							=			do_normalize
										lowercase_  							=			do_convert_rgb
										lowercase_  							=			max_patches
										lowercase_  							=			is_vqa

					def   UpperCAmelCase__       (     self       :   Optional[Any]							,					UpperCamelCase__       :   np.ndarray							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   dict							,					**UpperCamelCase__       :   Optional[int]       ):






										'''simple docstring'''






										requires_backends(self.extract_flattened_patches							,					"""torch"""       )
										_check_torch_version()

										# convert to torch
										lowercase_  							=			to_channel_dimension_format(UpperCamelCase__							,					ChannelDimension.FIRST       )
										lowercase_  							=			torch.from_numpy(UpperCamelCase__       )

										lowercase_							,					lowercase_  							=			patch_size["""height"""], patch_size["""width"""]
										lowercase_							,					lowercase_  							=			get_image_size(UpperCamelCase__       )

										# maximize scale s.t.
										lowercase_  							=			math.sqrt(max_patches * (patch_height / image_height) * (patch_width / image_width)       )
										lowercase_  							=			max(min(math.floor(scale * image_height / patch_height       )							,					UpperCamelCase__       )							,					1       )
										lowercase_  							=			max(min(math.floor(scale * image_width / patch_width       )							,					UpperCamelCase__       )							,					1       )
										lowercase_  							=			max(num_feasible_rows * patch_height							,					1       )
										lowercase_  							=			max(num_feasible_cols * patch_width							,					1       )

										lowercase_  							=			torch.nn.functional.interpolate(
										    image.unsqueeze(0       )							,					size=(resized_height, resized_width)							,					mode="""bilinear"""							,					align_corners=UpperCamelCase__							,					antialias=UpperCamelCase__							,					).squeeze(0       )

										# [1, rows, columns, patch_height * patch_width * image_channels]
										lowercase_  							=			torch_extract_patches(UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__       )

										lowercase_  							=			patches.shape
										lowercase_  							=			patches_shape[1]
										lowercase_  							=			patches_shape[2]
										lowercase_  							=			patches_shape[3]

										# [rows * columns, patch_height * patch_width * image_channels]
										lowercase_  							=			patches.reshape([rows * columns, depth]       )

										# [rows * columns, 1]
										lowercase_  							=			torch.arange(UpperCamelCase__       ).reshape([rows, 1]       ).repeat(1							,					UpperCamelCase__       ).reshape([rows * columns, 1]       )
										lowercase_  							=			torch.arange(UpperCamelCase__       ).reshape([1, columns]       ).repeat(UpperCamelCase__							,					1       ).reshape([rows * columns, 1]       )

										# Offset by 1 so the ids do not contain zeros, which represent padding.
										row_ids += 1
										col_ids += 1

										# Prepare additional patch features.
										# [rows * columns, 1]
										lowercase_  							=			row_ids.to(torch.floataa       )
										lowercase_  							=			col_ids.to(torch.floataa       )

										# [rows * columns, 2 + patch_height * patch_width * image_channels]
										lowercase_  							=			torch.cat([row_ids, col_ids, patches]							,					-1       )

										# [max_patches, 2 + patch_height * patch_width * image_channels]
										lowercase_  							=			torch.nn.functional.pad(UpperCamelCase__							,					[0, 0, 0, max_patches - (rows * columns)]       ).float()

										lowercase_  							=			to_numpy_array(UpperCamelCase__       )

										return result

					def   UpperCAmelCase__       (     self       :   List[Any]							,					UpperCamelCase__       :   np.ndarray							,					UpperCamelCase__       :   Optional[Union[str, ChannelDimension]] = None							,					**UpperCamelCase__       :   Dict       ):






										'''simple docstring'''






										if image.dtype == np.uinta:
															lowercase_  							=			image.astype(np.floataa       )

										# take mean across the whole `image`
										lowercase_  							=			np.mean(UpperCamelCase__       )
										lowercase_  							=			np.std(UpperCamelCase__       )
										lowercase_  							=			max(UpperCamelCase__							,					1.0 / math.sqrt(np.prod(image.shape       )       )       )

										return normalize(UpperCamelCase__							,					mean=UpperCamelCase__							,					std=UpperCamelCase__							,					**UpperCamelCase__       )







					def   UpperCAmelCase__       (     self       :   str							,					UpperCamelCase__       :   ImageInput							,					UpperCamelCase__       :   Optional[str] = None							,					UpperCamelCase__       :   bool = None							,					UpperCamelCase__       :   Optional[bool] = None							,					UpperCamelCase__       :   Optional[int] = None							,					UpperCamelCase__       :   Optional[Dict[str, int]] = None							,					UpperCamelCase__       :   Optional[Union[str, TensorType]] = None							,					UpperCamelCase__       :   ChannelDimension = ChannelDimension.FIRST							,					**UpperCamelCase__       :   Union[str, Any]							,					):






										'''simple docstring'''






										lowercase_  							=			do_normalize if do_normalize is not None else self.do_normalize
										lowercase_  							=			do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
										lowercase_  							=			patch_size if patch_size is not None else self.patch_size
										lowercase_  							=			max_patches if max_patches is not None else self.max_patches
										lowercase_  							=			self.is_vqa

										if kwargs.get("""data_format"""							,					UpperCamelCase__       ) is not None:
															raise ValueError("""data_format is not an accepted input as the outputs are """       )

										lowercase_  							=			make_list_of_images(UpperCamelCase__       )

										if not valid_images(UpperCamelCase__       ):
															raise ValueError(
															    """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
															    """torch.Tensor, tf.Tensor or jax.ndarray."""       )

										# PIL RGBA images are converted to RGB
										if do_convert_rgb:
															lowercase_  							=			[convert_to_rgb(UpperCamelCase__       ) for image in images]

										# All transformations expect numpy arrays.
										lowercase_  							=			[to_numpy_array(UpperCamelCase__       ) for image in images]

										if is_vqa:
															if header_text is None:
																				raise ValueError("""A header text must be provided for VQA models."""       )
															lowercase_  							=			kwargs.pop("""font_bytes"""							,					UpperCamelCase__       )
															lowercase_  							=			kwargs.pop("""font_path"""							,					UpperCamelCase__       )

															if isinstance(UpperCamelCase__							,					UpperCamelCase__       ):
																				lowercase_  							=			[header_text] * len(UpperCamelCase__       )

															lowercase_  							=			[
															    render_header(UpperCamelCase__							,					header_text[i]							,					font_bytes=UpperCamelCase__							,					font_path=UpperCamelCase__       )
															    for i, image in enumerate(UpperCamelCase__       )
															]

										if do_normalize:
															lowercase_  							=			[self.normalize(image=UpperCamelCase__       ) for image in images]

										# convert to torch tensor and permute
										lowercase_  							=			[
										    self.extract_flattened_patches(image=UpperCamelCase__							,					max_patches=UpperCamelCase__							,					patch_size=UpperCamelCase__       )
										    for image in images
										]

										# create attention mask in numpy
										lowercase_  							=			[(image.sum(axis=-1       ) != 0).astype(np.floataa       ) for image in images]

										lowercase_  							=			BatchFeature(
										    data={"""flattened_patches""": images, """attention_mask""": attention_masks}							,					tensor_type=UpperCamelCase__       )

										return encoded_outputs


