
'''simple docstring'''







import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch
							from torch import nn

							from transformers import (
							    MODEL_MAPPING,
							    BeitForImageClassification,
							    BeitForMaskedImageModeling,
							    BeitForSemanticSegmentation,
							    BeitModel,
							)
							from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							import PIL
							from PIL import Image

							from transformers import BeitImageProcessor






class  lowerCAmelCase__   :






					def __init__(   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE=1_00				,							__SCREAMING_SNAKE_CASE=13				,							__SCREAMING_SNAKE_CASE=30				,							__SCREAMING_SNAKE_CASE=2				,							__SCREAMING_SNAKE_CASE=3				,							__SCREAMING_SNAKE_CASE=True				,							__SCREAMING_SNAKE_CASE=True				,							__SCREAMING_SNAKE_CASE=32				,							__SCREAMING_SNAKE_CASE=4				,							__SCREAMING_SNAKE_CASE=4				,							__SCREAMING_SNAKE_CASE=37				,							__SCREAMING_SNAKE_CASE="gelu"				,							__SCREAMING_SNAKE_CASE=0.1				,							__SCREAMING_SNAKE_CASE=0.1				,							__SCREAMING_SNAKE_CASE=10				,							__SCREAMING_SNAKE_CASE=0.02				,							__SCREAMING_SNAKE_CASE=3				,							__SCREAMING_SNAKE_CASE=None				,							__SCREAMING_SNAKE_CASE=[0, 1, 2, 3]				,							):


											"""simple docstring"""





											lowercase_	:     Optional[int]													=       parent
											lowercase_	:     Tuple													=       1_00
											lowercase_	:     Any													=       batch_size
											lowercase_	:     Union[str, Any]													=       image_size
											lowercase_	:     Dict													=       patch_size
											lowercase_	:     List[str]													=       num_channels
											lowercase_	:     str													=       is_training
											lowercase_	:     Any													=       use_labels
											lowercase_	:     Optional[int]													=       hidden_size
											lowercase_	:     List[Any]													=       num_hidden_layers
											lowercase_	:     Optional[int]													=       num_attention_heads
											lowercase_	:     Union[str, Any]													=       intermediate_size
											lowercase_	:     Any													=       hidden_act
											lowercase_	:     List[Any]													=       hidden_dropout_prob
											lowercase_	:     Tuple													=       attention_probs_dropout_prob
											lowercase_	:     Dict													=       type_sequence_label_size
											lowercase_	:     Optional[Any]													=       initializer_range
											lowercase_	:     str													=       scope
											lowercase_	:     str													=       out_indices
											lowercase_	:     List[str]													=       num_labels

											# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
											lowercase_	:     Tuple													=       (image_size // patch_size) ** 2
											lowercase_	:     int													=       num_patches + 1






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     List[str]													=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

											lowercase_	:     List[str]													=       None
											lowercase_	:     int													=       None
											if self.use_labels:
																	lowercase_	:     str													=       ids_tensor([self.batch_size]				,							self.type_sequence_label_size		)
																	lowercase_	:     Dict													=       ids_tensor([self.batch_size, self.image_size, self.image_size]				,							self.num_labels		)

											lowercase_	:     int													=       self.get_config()

											return config, pixel_values, labels, pixel_labels






					def 	_snake_case    (   self		):


											"""simple docstring"""





											return BeitConfig(
											    vocab_size=self.vocab_size				,							image_size=self.image_size				,							patch_size=self.patch_size				,							num_channels=self.num_channels				,							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				,							is_decoder=__SCREAMING_SNAKE_CASE				,							initializer_range=self.initializer_range				,							out_indices=self.out_indices				,							)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





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






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											lowercase_	:     str													=       BeitForMaskedImageModeling(config=__SCREAMING_SNAKE_CASE		)
											model.to(__SCREAMING_SNAKE_CASE		)
											model.eval()
											lowercase_	:     int													=       model(__SCREAMING_SNAKE_CASE		)
											self.parent.assertEqual(result.logits.shape				,							(self.batch_size, self.seq_length - 1, self.vocab_size)		)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											lowercase_	:     List[Any]													=       self.type_sequence_label_size
											lowercase_	:     List[Any]													=       BeitForImageClassification(__SCREAMING_SNAKE_CASE		)
											model.to(__SCREAMING_SNAKE_CASE		)
											model.eval()
											lowercase_	:     List[str]													=       model(__SCREAMING_SNAKE_CASE				,							labels=__SCREAMING_SNAKE_CASE		)
											self.parent.assertEqual(result.logits.shape				,							(self.batch_size, self.type_sequence_label_size)		)

											# test greyscale images
											lowercase_	:     Optional[int]													=       1
											lowercase_	:     Optional[Any]													=       BeitForImageClassification(__SCREAMING_SNAKE_CASE		)
											model.to(__SCREAMING_SNAKE_CASE		)
											model.eval()

											lowercase_	:     Dict													=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
											lowercase_	:     Tuple													=       model(__SCREAMING_SNAKE_CASE				,							labels=__SCREAMING_SNAKE_CASE		)
											self.parent.assertEqual(result.logits.shape				,							(self.batch_size, self.type_sequence_label_size)		)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											lowercase_	:     Tuple													=       self.num_labels
											lowercase_	:     List[str]													=       BeitForSemanticSegmentation(__SCREAMING_SNAKE_CASE		)
											model.to(__SCREAMING_SNAKE_CASE		)
											model.eval()
											lowercase_	:     List[str]													=       model(__SCREAMING_SNAKE_CASE		)
											self.parent.assertEqual(
											    result.logits.shape				,							(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)		)
											lowercase_	:     int													=       model(__SCREAMING_SNAKE_CASE				,							labels=__SCREAMING_SNAKE_CASE		)
											self.parent.assertEqual(
											    result.logits.shape				,							(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     List[str]													=       self.prepare_config_and_inputs()
											lowercase_						, lowercase_						, lowercase_						, lowercase_	:     Optional[Any]													=       config_and_inputs
											lowercase_	:     Dict													=       {'''pixel_values''': pixel_values}
											return config, inputs_dict







@require_torch
class  lowerCAmelCase__   (  lowerCamelCase_    ,						lowerCamelCase_    ,						unittest.TestCase     ):
					lowerCAmelCase_        =     (
					    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
					    if is_torch_available()
					    else ()
					)
					lowerCAmelCase_        =     (
					    {
					        '''feature-extraction''': BeitModel,
					        '''image-classification''': BeitForImageClassification,
					        '''image-segmentation''': BeitForSemanticSegmentation,
					    }
					    if is_torch_available()
					    else {}
					)

					lowerCAmelCase_        =     False
					lowerCAmelCase_        =     False
					lowerCAmelCase_        =     False






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Any													=       BeitModelTester(self		)
											lowercase_	:     Dict													=       ConfigTester(self				,							config_class=__SCREAMING_SNAKE_CASE				,							has_text_modality=__SCREAMING_SNAKE_CASE				,							hidden_size=37		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											self.config_tester.run_common_tests()






					@unittest.skip(reason='''BEiT does not use inputs_embeds'''		)
					def 	_snake_case    (   self		):


											"""simple docstring"""





											pass






					@require_torch_multi_gpu
					@unittest.skip(reason='''BEiT has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''		)
					def 	_snake_case    (   self		):


											"""simple docstring"""





											pass






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_						, lowercase_	:     List[str]													=       self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	lowercase_	:     Any													=       model_class(__SCREAMING_SNAKE_CASE		)
																	self.assertIsInstance(model.get_input_embeddings()				,							(nn.Module)		)
																	lowercase_	:     int													=       model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(__SCREAMING_SNAKE_CASE				,							nn.Linear		)		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_						, lowercase_	:     str													=       self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	lowercase_	:     List[str]													=       model_class(__SCREAMING_SNAKE_CASE		)
																	lowercase_	:     Tuple													=       inspect.signature(model.forward		)
																	# signature.parameters is an OrderedDict => so arg_names order is deterministic
																	lowercase_	:     int													=       [*signature.parameters.keys()]

																	lowercase_	:     List[Any]													=       ['''pixel_values''']
																	self.assertListEqual(arg_names[:1]				,							__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     List[Any]													=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     int													=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Optional[int]													=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Union[str, Any]													=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_semantic_segmentation(*__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											if not self.model_tester.is_training:
																	return

											lowercase_						, lowercase_	:     Any													=       self.model_tester.prepare_config_and_inputs_for_common()
											lowercase_	:     List[Any]													=       True

											for model_class in self.all_model_classes:
																	# we don't test BeitForMaskedImageModeling
																	if model_class in [*get_values(__SCREAMING_SNAKE_CASE		), BeitForMaskedImageModeling]:
																							continue

																	lowercase_	:     int													=       model_class(__SCREAMING_SNAKE_CASE		)
																	model.to(__SCREAMING_SNAKE_CASE		)
																	model.train()
																	lowercase_	:     str													=       self._prepare_for_class(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							return_labels=__SCREAMING_SNAKE_CASE		)
																	lowercase_	:     Dict													=       model(**__SCREAMING_SNAKE_CASE		).loss
																	loss.backward()






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_						, lowercase_	:     str													=       self.model_tester.prepare_config_and_inputs_for_common()
											if not self.model_tester.is_training:
																	return

											lowercase_	:     Any													=       False
											lowercase_	:     List[str]													=       True

											for model_class in self.all_model_classes:
																	# we don't test BeitForMaskedImageModeling
																	if (
																	    model_class in [*get_values(__SCREAMING_SNAKE_CASE		), BeitForMaskedImageModeling]
																	    or not model_class.supports_gradient_checkpointing
																	):
																							continue

																	lowercase_	:     int													=       model_class(__SCREAMING_SNAKE_CASE		)
																	model.gradient_checkpointing_enable()
																	model.to(__SCREAMING_SNAKE_CASE		)
																	model.train()
																	lowercase_	:     Dict													=       self._prepare_for_class(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							return_labels=__SCREAMING_SNAKE_CASE		)
																	lowercase_	:     Tuple													=       model(**__SCREAMING_SNAKE_CASE		).loss
																	loss.backward()






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_						, lowercase_	:     Dict													=       self.model_tester.prepare_config_and_inputs_for_common()

											lowercase_	:     str													=       _config_zero_init(__SCREAMING_SNAKE_CASE		)
											for model_class in self.all_model_classes:
																	lowercase_	:     Optional[Any]													=       model_class(config=__SCREAMING_SNAKE_CASE		)
																	for name, param in model.named_parameters():
																							# we skip lambda parameters as these require special initial values
																							# determined by config.layer_scale_init_value
																							if "lambda" in name:
																													continue
																							if param.requires_grad:
																													self.assertIn(
																													    ((param.data.mean() * 1E9).round() / 1E9).item()				,							[0.0, 1.0]				,							msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''				,							)






					@slow
					def 	_snake_case    (   self		):


											"""simple docstring"""





											for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowercase_	:     Dict													=       BeitModel.from_pretrained(__SCREAMING_SNAKE_CASE		)
																	self.assertIsNotNone(__SCREAMING_SNAKE_CASE		)







def 					snake_case_		(				):



						"""simple docstring"""





						lowercase_	:     Optional[Any]													=       Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)
						return image






@require_torch
@require_vision
class  lowerCAmelCase__   (  unittest.TestCase     ):






					@cached_property
					def 	_snake_case    (   self		):


											"""simple docstring"""





											return BeitImageProcessor.from_pretrained('''microsoft/beit-base-patch16-224'''		) if is_vision_available() else None






					@slow
					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     int													=       BeitForMaskedImageModeling.from_pretrained('''microsoft/beit-base-patch16-224-pt22k'''		).to(__SCREAMING_SNAKE_CASE		)

											lowercase_	:     int													=       self.default_image_processor
											lowercase_	:     Optional[int]													=       prepare_img()
											lowercase_	:     str													=       image_processor(images=__SCREAMING_SNAKE_CASE				,							return_tensors='''pt'''		).pixel_values.to(__SCREAMING_SNAKE_CASE		)

											# prepare bool_masked_pos
											lowercase_	:     Union[str, Any]													=       torch.ones((1, 1_96)				,							dtype=torch.bool		).to(__SCREAMING_SNAKE_CASE		)

											# forward pass
											with torch.no_grad():
																	lowercase_	:     str													=       model(pixel_values=__SCREAMING_SNAKE_CASE				,							bool_masked_pos=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Any													=       outputs.logits

											# verify the logits
											lowercase_	:     Optional[int]													=       torch.Size((1, 1_96, 81_92)		)
											self.assertEqual(logits.shape				,							__SCREAMING_SNAKE_CASE		)

											lowercase_	:     List[str]													=       torch.tensor(
											    [[-3.2_437, 0.5_072, -13.9_174], [-3.2_456, 0.4_948, -13.9_401], [-3.2_033, 0.5_121, -13.8_550]]		).to(__SCREAMING_SNAKE_CASE		)

											self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]				,							__SCREAMING_SNAKE_CASE				,							atol=1E-2		)		)






					@slow
					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Optional[Any]													=       BeitForImageClassification.from_pretrained('''microsoft/beit-base-patch16-224'''		).to(__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Dict													=       self.default_image_processor
											lowercase_	:     Union[str, Any]													=       prepare_img()
											lowercase_	:     Any													=       image_processor(images=__SCREAMING_SNAKE_CASE				,							return_tensors='''pt'''		).to(__SCREAMING_SNAKE_CASE		)

											# forward pass
											with torch.no_grad():
																	lowercase_	:     List[Any]													=       model(**__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Dict													=       outputs.logits

											# verify the logits
											lowercase_	:     Any													=       torch.Size((1, 10_00)		)
											self.assertEqual(logits.shape				,							__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Optional[Any]													=       torch.tensor([-1.2_385, -1.0_987, -1.0_108]		).to(__SCREAMING_SNAKE_CASE		)

											self.assertTrue(torch.allclose(logits[0, :3]				,							__SCREAMING_SNAKE_CASE				,							atol=1E-4		)		)

											lowercase_	:     Any													=       2_81
											self.assertEqual(logits.argmax(-1		).item()				,							__SCREAMING_SNAKE_CASE		)






					@slow
					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Dict													=       BeitForImageClassification.from_pretrained('''microsoft/beit-large-patch16-224-pt22k-ft22k'''		).to(
											    __SCREAMING_SNAKE_CASE		)

											lowercase_	:     List[Any]													=       self.default_image_processor
											lowercase_	:     Any													=       prepare_img()
											lowercase_	:     Any													=       image_processor(images=__SCREAMING_SNAKE_CASE				,							return_tensors='''pt'''		).to(__SCREAMING_SNAKE_CASE		)

											# forward pass
											with torch.no_grad():
																	lowercase_	:     List[str]													=       model(**__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Union[str, Any]													=       outputs.logits

											# verify the logits
											lowercase_	:     Union[str, Any]													=       torch.Size((1, 2_18_41)		)
											self.assertEqual(logits.shape				,							__SCREAMING_SNAKE_CASE		)

											lowercase_	:     List[Any]													=       torch.tensor([1.6_881, -0.2_787, 0.5_901]		).to(__SCREAMING_SNAKE_CASE		)

											self.assertTrue(torch.allclose(logits[0, :3]				,							__SCREAMING_SNAKE_CASE				,							atol=1E-4		)		)

											lowercase_	:     str													=       23_96
											self.assertEqual(logits.argmax(-1		).item()				,							__SCREAMING_SNAKE_CASE		)






					@slow
					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Optional[int]													=       BeitForSemanticSegmentation.from_pretrained('''microsoft/beit-base-finetuned-ade-640-640'''		)
											lowercase_	:     Any													=       model.to(__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Optional[Any]													=       BeitImageProcessor(do_resize=__SCREAMING_SNAKE_CASE				,							size=6_40				,							do_center_crop=__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Union[str, Any]													=       load_dataset('''hf-internal-testing/fixtures_ade20k'''				,							split='''test'''		)
											lowercase_	:     Any													=       Image.open(ds[0]['''file''']		)
											lowercase_	:     Tuple													=       image_processor(images=__SCREAMING_SNAKE_CASE				,							return_tensors='''pt'''		).to(__SCREAMING_SNAKE_CASE		)

											# forward pass
											with torch.no_grad():
																	lowercase_	:     Union[str, Any]													=       model(**__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Optional[Any]													=       outputs.logits

											# verify the logits
											lowercase_	:     Any													=       torch.Size((1, 1_50, 1_60, 1_60)		)
											self.assertEqual(logits.shape				,							__SCREAMING_SNAKE_CASE		)

											lowercase_	:     int													=       version.parse(PIL.__version__		) < version.parse('''9.0.0'''		)

											if is_pillow_less_than_a:
																	lowercase_	:     Dict													=       torch.tensor(
																	    [
																	        [[-4.9_225, -2.3_954, -3.0_522], [-2.8_822, -1.0_046, -1.7_561], [-2.9_549, -1.3_228, -2.1_347]],
																	        [[-5.8_168, -3.4_129, -4.0_778], [-3.8_651, -2.2_214, -3.0_277], [-3.8_356, -2.4_643, -3.3_535]],
																	        [[-0.0_078, 3.9_952, 4.0_754], [2.9_856, 4.6_944, 5.0_035], [3.2_413, 4.7_813, 4.9_969]],
																	    ]				,							device=__SCREAMING_SNAKE_CASE				,							)
											else:
																	lowercase_	:     int													=       torch.tensor(
																	    [
																	        [[-4.8_960, -2.3_688, -3.0_355], [-2.8_478, -0.9_836, -1.7_418], [-2.9_449, -1.3_332, -2.1_456]],
																	        [[-5.8_081, -3.4_124, -4.1_006], [-3.8_561, -2.2_081, -3.0_323], [-3.8_365, -2.4_601, -3.3_669]],
																	        [[-0.0_309, 3.9_868, 4.0_540], [2.9_640, 4.6_877, 4.9_976], [3.2_081, 4.7_690, 4.9_942]],
																	    ]				,							device=__SCREAMING_SNAKE_CASE				,							)

											self.assertTrue(torch.allclose(logits[0, :3, :3, :3]				,							__SCREAMING_SNAKE_CASE				,							atol=1E-4		)		)






					@slow
					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Tuple													=       BeitForSemanticSegmentation.from_pretrained('''microsoft/beit-base-finetuned-ade-640-640'''		)
											lowercase_	:     str													=       model.to(__SCREAMING_SNAKE_CASE		)

											lowercase_	:     List[str]													=       BeitImageProcessor(do_resize=__SCREAMING_SNAKE_CASE				,							size=6_40				,							do_center_crop=__SCREAMING_SNAKE_CASE		)

											lowercase_	:     int													=       load_dataset('''hf-internal-testing/fixtures_ade20k'''				,							split='''test'''		)
											lowercase_	:     Optional[int]													=       Image.open(ds[0]['''file''']		)
											lowercase_	:     Union[str, Any]													=       image_processor(images=__SCREAMING_SNAKE_CASE				,							return_tensors='''pt'''		).to(__SCREAMING_SNAKE_CASE		)

											# forward pass
											with torch.no_grad():
																	lowercase_	:     str													=       model(**__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Optional[int]													=       outputs.logits.detach().cpu()

											lowercase_	:     Optional[int]													=       image_processor.post_process_semantic_segmentation(outputs=__SCREAMING_SNAKE_CASE				,							target_sizes=[(5_00, 3_00)]		)
											lowercase_	:     List[Any]													=       torch.Size((5_00, 3_00)		)
											self.assertEqual(segmentation[0].shape				,							__SCREAMING_SNAKE_CASE		)

											lowercase_	:     List[str]													=       image_processor.post_process_semantic_segmentation(outputs=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Union[str, Any]													=       torch.Size((1_60, 1_60)		)
											self.assertEqual(segmentation[0].shape				,							__SCREAMING_SNAKE_CASE		)





from typing import Dict

import numpy as np

from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, GenericTensor, Pipeline, PipelineException


if is_tf_available():
				import tensorflow as tf

				from ..tf_utils import stable_softmax


if is_torch_available():
				import torch


__snake_case								=  logging.get_logger(__name__)




@add_end_docstrings(
    _UpperCAmelCase     ,       R"""
        top_k (`int`, defaults to 5):
            The number of predictions to return.
        targets (`str` or `List[str]`, *optional*):
            When passed, the model will limit the scores to the passed targets instead of looking up in the whole
            vocab. If the provided targets are not in the model vocab, they will be tokenized and the first resulting
            token will be used (with a warning, and that might be slower).

    """     ,       )
class 						lowercase__  (						_UpperCAmelCase  ):
	def      A_				(			self							:   Any       ,   UpperCAmelCase_							:   GenericTensor							):
						if self.framework == "tf":
											SCREAMING_SNAKE_CASE__													= tf.where(input_ids == self.tokenizer.mask_token_id							).numpy()
						elif self.framework == "pt":
											SCREAMING_SNAKE_CASE__													= torch.nonzero(input_ids == self.tokenizer.mask_token_id       ,   as_tuple=UpperCAmelCase_							)
						else:
											raise ValueError('Unsupported framework'							)
						return masked_index
	def      A_				(			self							:   Optional[Any]       ,   UpperCAmelCase_							:   GenericTensor							):
						SCREAMING_SNAKE_CASE__													= self.get_masked_index(UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= np.prod(masked_index.shape							)
						if numel < 1:
											raise PipelineException(
											    'fill-mask'       ,   self.model.base_model_prefix       ,   F'No mask_token ({self.tokenizer.mask_token}) found on the input'       ,   )
	def      A_				(			self							:   Union[str, Any]       ,   UpperCAmelCase_							:   GenericTensor							):
						if isinstance(UpperCAmelCase_       ,   UpperCAmelCase_							):
											for model_input in model_inputs:
																self._ensure_exactly_one_mask_token(model_input['input_ids'][0]							)
						else:
											for input_ids in model_inputs["input_ids"]:
																self._ensure_exactly_one_mask_token(UpperCAmelCase_							)
	def      A_				(			self							:   Any       ,   UpperCAmelCase_							:   Optional[Any]       ,   UpperCAmelCase_							:   str=None       ,   **UpperCAmelCase_							:   Union[str, Any]							):
						if return_tensors is None:
											SCREAMING_SNAKE_CASE__													= self.framework
						SCREAMING_SNAKE_CASE__													= self.tokenizer(UpperCAmelCase_       ,   return_tensors=UpperCAmelCase_							)
						self.ensure_exactly_one_mask_token(UpperCAmelCase_							)
						return model_inputs
	def      A_				(			self							:   Tuple       ,   UpperCAmelCase_							:   Dict							):
						SCREAMING_SNAKE_CASE__													= self.model(**UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= model_inputs['input_ids']
						return model_outputs
	def      A_				(			self							:   Optional[int]       ,   UpperCAmelCase_							:   List[Any]       ,   UpperCAmelCase_							:   Union[str, Any]=5       ,   UpperCAmelCase_							:   Dict=None							):
						# Cap top_k if there are targets
						if target_ids is not None and target_ids.shape[0] < top_k:
											SCREAMING_SNAKE_CASE__													= target_ids.shape[0]
						SCREAMING_SNAKE_CASE__													= model_outputs['input_ids'][0]
						SCREAMING_SNAKE_CASE__													= model_outputs['logits']

						if self.framework == "tf":
											SCREAMING_SNAKE_CASE__													= tf.where(input_ids == self.tokenizer.mask_token_id							).numpy()[:, 0]

											SCREAMING_SNAKE_CASE__													= outputs.numpy()

											SCREAMING_SNAKE_CASE__													= outputs[0, masked_index, :]
											SCREAMING_SNAKE_CASE__													= stable_softmax(UpperCAmelCase_       ,   axis=-1							)
											if target_ids is not None:
																SCREAMING_SNAKE_CASE__													= tf.gather_nd(tf.squeeze(UpperCAmelCase_       ,   0							)       ,   target_ids.reshape(-1       ,   1							)							)
																SCREAMING_SNAKE_CASE__													= tf.expand_dims(UpperCAmelCase_       ,   0							)

											SCREAMING_SNAKE_CASE__													= tf.math.top_k(UpperCAmelCase_       ,   k=UpperCAmelCase_							)
											SCREAMING_SNAKE_CASE__						, SCREAMING_SNAKE_CASE__													= topk.values.numpy(), topk.indices.numpy()
						else:
											SCREAMING_SNAKE_CASE__													= torch.nonzero(input_ids == self.tokenizer.mask_token_id       ,   as_tuple=UpperCAmelCase_							).squeeze(-1							)
											# Fill mask pipeline supports only one ${mask_token} per sample

											SCREAMING_SNAKE_CASE__													= outputs[0, masked_index, :]
											SCREAMING_SNAKE_CASE__													= logits.softmax(dim=-1							)
											if target_ids is not None:
																SCREAMING_SNAKE_CASE__													= probs[..., target_ids]

											SCREAMING_SNAKE_CASE__						, SCREAMING_SNAKE_CASE__													= probs.topk(UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= []
						SCREAMING_SNAKE_CASE__													= values.shape[0] == 1
						for i, (_values, _predictions) in enumerate(zip(values.tolist()       ,   predictions.tolist()							)							):
											SCREAMING_SNAKE_CASE__													= []
											for v, p in zip(_values       ,   _predictions							):
																# Copy is important since we're going to modify this array in place
																SCREAMING_SNAKE_CASE__													= input_ids.numpy().copy()
																if target_ids is not None:
																					SCREAMING_SNAKE_CASE__													= target_ids[p].tolist()

																SCREAMING_SNAKE_CASE__													= p
																# Filter padding out:
																SCREAMING_SNAKE_CASE__													= tokens[np.where(tokens != self.tokenizer.pad_token_id							)]
																# Originally we skip special tokens to give readable output.
																# For multi masks though, the other [MASK] would be removed otherwise
																# making the output look odd, so we add them back
																SCREAMING_SNAKE_CASE__													= self.tokenizer.decode(UpperCAmelCase_       ,   skip_special_tokens=UpperCAmelCase_							)
																SCREAMING_SNAKE_CASE__													= {'score': v, 'token': p, 'token_str': self.tokenizer.decode([p]							), 'sequence': sequence}
																row.append(UpperCAmelCase_							)
											result.append(UpperCAmelCase_							)
						if single_mask:
											return result[0]
						return result
	def      A_				(			self							:   Dict       ,   UpperCAmelCase_							:   Tuple       ,   UpperCAmelCase_							:   List[Any]=None							):
						if isinstance(UpperCAmelCase_       ,   UpperCAmelCase_							):
											SCREAMING_SNAKE_CASE__													= [targets]
						try:
											SCREAMING_SNAKE_CASE__													= self.tokenizer.get_vocab()
						except Exception:
											SCREAMING_SNAKE_CASE__													= {}
						SCREAMING_SNAKE_CASE__													= []
						for target in targets:
											SCREAMING_SNAKE_CASE__													= vocab.get(UpperCAmelCase_       ,   UpperCAmelCase_							)
											if id_ is None:
																SCREAMING_SNAKE_CASE__													= self.tokenizer(
																    UpperCAmelCase_       ,   add_special_tokens=UpperCAmelCase_       ,   return_attention_mask=UpperCAmelCase_       ,   return_token_type_ids=UpperCAmelCase_       ,   max_length=1       ,   truncation=UpperCAmelCase_       ,   )['input_ids']
																if len(UpperCAmelCase_							) == 0:
																					logger.warning(
																					    F'The specified target token `{target}` does not exist in the model vocabulary. '
																					    'We cannot replace it with anything meaningful, ignoring it'							)
																					continue
																SCREAMING_SNAKE_CASE__													= input_ids[0]
																# XXX: If users encounter this pass
																# it becomes pretty slow, so let's make sure
																# The warning enables them to fix the input to
																# get faster performance.
																logger.warning(
																    F'The specified target token `{target}` does not exist in the model vocabulary. '
																    F'Replacing with `{self.tokenizer.convert_ids_to_tokens(id_							)}`.'							)
											target_ids.append(id_							)
						SCREAMING_SNAKE_CASE__													= list(set(UpperCAmelCase_							)							)
						if len(UpperCAmelCase_							) == 0:
											raise ValueError('At least one target must be provided when passed.'							)
						SCREAMING_SNAKE_CASE__													= np.array(UpperCAmelCase_							)
						return target_ids
	def      A_				(			self							:   List[str]       ,   UpperCAmelCase_							:   Optional[Any]=None       ,   UpperCAmelCase_							:   List[Any]=None							):
						SCREAMING_SNAKE_CASE__													= {}

						if targets is not None:
											SCREAMING_SNAKE_CASE__													= self.get_target_ids(UpperCAmelCase_       ,   UpperCAmelCase_							)
											SCREAMING_SNAKE_CASE__													= target_ids

						if top_k is not None:
											SCREAMING_SNAKE_CASE__													= top_k

						if self.tokenizer.mask_token_id is None:
											raise PipelineException(
											    'fill-mask'       ,   self.model.base_model_prefix       ,   'The tokenizer does not define a `mask_token`.'							)
						return {}, {}, postprocess_params





	def __call__(			self							:   Tuple       ,   UpperCAmelCase_							:   Union[str, Any]       ,   *UpperCAmelCase_							:   int       ,   **UpperCAmelCase_							:   Tuple							):
						SCREAMING_SNAKE_CASE__													= super().__call__(UpperCAmelCase_       ,   **UpperCAmelCase_							)
						if isinstance(UpperCAmelCase_       ,   UpperCAmelCase_							) and len(UpperCAmelCase_							) == 1:
											return outputs[0]
						return outputs





