
"""simple docstring"""


from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import tensorflow as tf

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







class 				lowerCAmelCase_				:



						'''simple docstring'''


						def __init__( self     :      str  ,A_     :      Optional[Any]  ,)     ->   Union[str, Any]:
							A						   =   parent
							A						   =   13
							A						   =   7
							A						   =   True
							A						   =   True
							A						   =   False
							A						   =   True
							A						   =   99
							A						   =   32
							A						   =   2
							A						   =   4
							A						   =   37
							A						   =   'gelu'
							A						   =   0.1
							A						   =   0.1
							A						   =   512
							A						   =   16
							A						   =   2
							A						   =   0.02
							A						   =   3
							A						   =   4
							A						   =   None


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]      )     ->   Optional[Any]:
							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						   =   DistilBertConfig(
							    vocab_size=self.vocab_size  ,dim=self.hidden_size  ,n_layers=self.num_hidden_layers  ,n_heads=self.num_attention_heads  ,hidden_dim=self.intermediate_size  ,hidden_act=self.hidden_act  ,dropout=self.hidden_dropout_prob  ,attention_dropout=self.attention_probs_dropout_prob  ,max_position_embeddings=self.max_position_embeddings  ,initializer_range=self.initializer_range  ,)

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


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]  ,A_     :      int  ,A_     :      Tuple  ,A_     :      Optional[Any]  ,A_     :      str  ,A_     :      Union[str, Any]  ,A_     :      str      )     ->   List[str]:
							A						   =   TFDistilBertModel(config=A_      )
							A						   =   {'input_ids': input_ids, 'attention_mask': input_mask}

							A						   =   model(A_      )

							A						   =   [input_ids, input_mask]

							A						   =   model(A_      )

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


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]  ,A_     :      Tuple  ,A_     :      List[Any]  ,A_     :      Union[str, Any]  ,A_     :      Union[str, Any]  ,A_     :      Union[str, Any]  ,A_     :      str      )     ->   Any:
							A						   =   TFDistilBertForMaskedLM(config=A_      )
							A						   =   {'input_ids': input_ids, 'attention_mask': input_mask}
							A						   =   model(A_      )
							self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.vocab_size)      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      int  ,A_     :      Any  ,A_     :      Any  ,A_     :      Dict  ,A_     :      Optional[Any]  ,A_     :      Union[str, Any]  ,A_     :      str      )     ->   Tuple:
							A						   =   TFDistilBertForQuestionAnswering(config=A_      )
							A						   =   {
							    'input_ids': input_ids,
							    'attention_mask': input_mask,
							}
							A						   =   model(A_      )
							self.parent.assertEqual(result.start_logits.shape  ,(self.batch_size, self.seq_length)      )
							self.parent.assertEqual(result.end_logits.shape  ,(self.batch_size, self.seq_length)      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple  ,A_     :      Optional[Any]  ,A_     :      Optional[int]  ,A_     :      Optional[int]  ,A_     :      Optional[Any]  ,A_     :      List[str]  ,A_     :      Dict      )     ->   Tuple:
							A						   =   self.num_labels
							A						   =   TFDistilBertForSequenceClassification(A_      )
							A						   =   {'input_ids': input_ids, 'attention_mask': input_mask}
							A						   =   model(A_      )
							self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.num_labels)      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      str  ,A_     :      Tuple  ,A_     :      List[str]  ,A_     :      str  ,A_     :      Tuple  ,A_     :      Tuple  ,A_     :      Tuple      )     ->   int:
							A						   =   self.num_choices
							A						   =   TFDistilBertForMultipleChoice(A_      )
							A						   =   tf.tile(tf.expand_dims(A_  ,1      )  ,(1, self.num_choices, 1)      )
							A						   =   tf.tile(tf.expand_dims(A_  ,1      )  ,(1, self.num_choices, 1)      )
							A						   =   {
							    'input_ids': multiple_choice_inputs_ids,
							    'attention_mask': multiple_choice_input_mask,
							}
							A						   =   model(A_      )
							self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.num_choices)      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]  ,A_     :      List[str]  ,A_     :      List[str]  ,A_     :      Optional[int]  ,A_     :      int  ,A_     :      int  ,A_     :      List[Any]      )     ->   str:
							A						   =   self.num_labels
							A						   =   TFDistilBertForTokenClassification(A_      )
							A						   =   {'input_ids': input_ids, 'attention_mask': input_mask}
							A						   =   model(A_      )
							self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.num_labels)      )




						def 						_SCREAMING_SNAKE_CASE      ( self     :      Union[str, Any]      )     ->   Union[str, Any]:
							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_tf
class 				lowerCAmelCase_				(      _lowercase		,					_lowercase		,					unittest.TestCase    ):



						'''simple docstring'''


						_lowerCamelCase:    Optional[int]      =    (
						    (
						        TFDistilBertModel,
						        TFDistilBertForMaskedLM,
						        TFDistilBertForQuestionAnswering,
						        TFDistilBertForSequenceClassification,
						        TFDistilBertForTokenClassification,
						        TFDistilBertForMultipleChoice,
						    )
						    if is_tf_available()
						    else None
						)
						_lowerCamelCase:    str      =    (
						    {
						        '''feature-extraction''': TFDistilBertModel,
						        '''fill-mask''': TFDistilBertForMaskedLM,
						        '''question-answering''': TFDistilBertForQuestionAnswering,
						        '''text-classification''': TFDistilBertForSequenceClassification,
						        '''token-classification''': TFDistilBertForTokenClassification,
						        '''zero-shot''': TFDistilBertForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						_lowerCamelCase:    Optional[int]      =    False
						_lowerCamelCase:    str      =    False


						def 						_SCREAMING_SNAKE_CASE      ( self     :      str      )     ->   Union[str, Any]:
							A						   =   TFDistilBertModelTester(self      )
							A						   =   ConfigTester(self  ,config_class=A_  ,dim=37      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   Optional[int]:
							self.config_tester.run_common_tests()


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Any      )     ->   int:
							A						   =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_model(*A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      int      )     ->   str:
							A						   =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_masked_lm(*A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      str      )     ->   Optional[int]:
							A						   =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_question_answering(*A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      str      )     ->   List[str]:
							A						   =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_sequence_classification(*A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   List[Any]:
							A						   =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_multiple_choice(*A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   Any:
							A						   =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_token_classification(*A_      )




						@slow
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   Any:
							for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]      ):
								A						   =   TFDistilBertModel.from_pretrained(A_      )
								self.assertIsNotNone(A_      )







@require_tf
class 				lowerCAmelCase_				(      unittest.TestCase    ):



						'''simple docstring'''


						@slow
						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]      )     ->   Union[str, Any]:
							A						   =   TFDistilBertModel.from_pretrained('distilbert-base-uncased'      )
							A						   =   tf.constant([[0, 1, 2, 3, 4, 5]]      )
							A						   =   model(A_      )[0]

							A						   =   [1, 6, 768]
							self.assertEqual(output.shape  ,A_      )

							A						   =   tf.constant(
							    [
							        [
							            [0.19_26_18_85, -0.13_73_29_55, 0.4_11_97_99],
							            [0.22_15_01_56, -0.07_42_26_61, 0.39_03_72_04],
							            [0.22_75_60_18, -0.0_89_64_14, 0.3_70_14_67],
							        ]
							    ]      )
							tf.debugging.assert_near(output[:, :3, :3]  ,A_  ,atol=1e-4      )

"""simple docstring"""


from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, logging


if is_torch_available():
		import torch


_lowercase							   =				logging.get_logger(__name__)







class 				lowerCAmelCase_				(      _lowercase    ):



						'''simple docstring'''


						_lowerCamelCase:    Dict      =    ['''pixel_values''']


						def __init__( self     :      Optional[Any]  ,A_     :      bool = True  ,A_     :      Optional[Dict[str, int]] = None  ,A_     :      PILImageResampling = PILImageResampling.BILINEAR  ,A_     :      bool = True  ,A_     :      Dict[str, int] = None  ,A_     :      bool = True  ,A_     :      Union[int, float] = 1 / 255  ,A_     :      bool = True  ,A_     :      Optional[Union[float, List[float]]] = None  ,A_     :      Optional[Union[float, List[float]]] = None  ,**A_     :      Optional[Any]  ,)     ->   None:
							super().__init__(**A_      )
							A						   =   size if size is not None else {'shortest_edge': 256}
							A						   =   get_size_dict(A_  ,default_to_square=A_      )
							A						   =   crop_size if crop_size is not None else {'height': 224, 'width': 224}
							A						   =   get_size_dict(A_  ,param_name='crop_size'      )
							A						   =   do_resize
							A						   =   size
							A						   =   resample
							A						   =   do_center_crop
							A						   =   crop_size
							A						   =   do_rescale
							A						   =   rescale_factor
							A						   =   do_normalize
							A						   =   image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
							A						   =   image_std if image_std is not None else IMAGENET_STANDARD_STD


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Union[str, Any]  ,A_     :      np.ndarray  ,A_     :      Dict[str, int]  ,A_     :      PILImageResampling = PILImageResampling.BICUBIC  ,A_     :      Optional[Union[str, ChannelDimension]] = None  ,**A_     :      int  ,)     ->   np.ndarray:
							A						   =   get_size_dict(A_  ,default_to_square=A_      )
							if "shortest_edge" not in size:
								raise ValueError(F'The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'      )
							A						   =   get_resize_output_image_size(A_  ,size=size['shortest_edge']  ,default_to_square=A_      )
							return resize(A_  ,size=A_  ,resample=A_  ,data_format=A_  ,**A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]  ,A_     :      np.ndarray  ,A_     :      Dict[str, int]  ,A_     :      Optional[Union[str, ChannelDimension]] = None  ,**A_     :      int  ,)     ->   np.ndarray:
							A						   =   get_size_dict(A_      )
							if "height" not in size or "width" not in size:
								raise ValueError(F'The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}'      )
							return center_crop(A_  ,size=(size['height'], size['width'])  ,data_format=A_  ,**A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]  ,A_     :      np.ndarray  ,A_     :      float  ,A_     :      Optional[Union[str, ChannelDimension]] = None  ,**A_     :      List[str]      )     ->   np.ndarray:
							return rescale(A_  ,scale=A_  ,data_format=A_  ,**A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]  ,A_     :      np.ndarray  ,A_     :      Union[float, List[float]]  ,A_     :      Union[float, List[float]]  ,A_     :      Optional[Union[str, ChannelDimension]] = None  ,**A_     :      Any  ,)     ->   np.ndarray:
							return normalize(A_  ,mean=A_  ,std=A_  ,data_format=A_  ,**A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Any  ,A_     :      ImageInput  ,A_     :      Optional[bool] = None  ,A_     :      Dict[str, int] = None  ,A_     :      PILImageResampling = None  ,A_     :      bool = None  ,A_     :      Dict[str, int] = None  ,A_     :      Optional[bool] = None  ,A_     :      Optional[float] = None  ,A_     :      Optional[bool] = None  ,A_     :      Optional[Union[float, List[float]]] = None  ,A_     :      Optional[Union[float, List[float]]] = None  ,A_     :      Optional[Union[str, TensorType]] = None  ,A_     :      Union[str, ChannelDimension] = ChannelDimension.FIRST  ,**A_     :      Tuple  ,)     ->   List[Any]:
							A						   =   do_resize if do_resize is not None else self.do_resize
							A						   =   size if size is not None else self.size
							A						   =   get_size_dict(A_  ,default_to_square=A_      )
							A						   =   resample if resample is not None else self.resample
							A						   =   do_center_crop if do_center_crop is not None else self.do_center_crop
							A						   =   crop_size if crop_size is not None else self.crop_size
							A						   =   get_size_dict(A_  ,param_name='crop_size'      )
							A						   =   do_rescale if do_rescale is not None else self.do_rescale
							A						   =   rescale_factor if rescale_factor is not None else self.rescale_factor
							A						   =   do_normalize if do_normalize is not None else self.do_normalize
							A						   =   image_mean if image_mean is not None else self.image_mean
							A						   =   image_std if image_std is not None else self.image_std

							A						   =   make_list_of_images(A_      )

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

							if do_resize and size is None:
								raise ValueError('Size must be specified if do_resize is True.'      )

							if do_center_crop and crop_size is None:
								raise ValueError('Crop size must be specified if do_center_crop is True.'      )

							if do_rescale and rescale_factor is None:
								raise ValueError('Rescale factor must be specified if do_rescale is True.'      )

							if do_normalize and (image_mean is None or image_std is None):
								raise ValueError('Image mean and std must be specified if do_normalize is True.'      )

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

							if do_resize:
								A						   =   [self.resize(image=A_  ,size=A_  ,resample=A_      ) for image in images]

							if do_center_crop:
								A						   =   [self.center_crop(image=A_  ,size=A_      ) for image in images]

							if do_rescale:
								A						   =   [self.rescale(image=A_  ,scale=A_      ) for image in images]

							if do_normalize:
								A						   =   [self.normalize(image=A_  ,mean=A_  ,std=A_      ) for image in images]

							A						   =   [to_channel_dimension_format(A_  ,A_      ) for image in images]

							A						   =   {'pixel_values': images}
							return BatchFeature(data=A_  ,tensor_type=A_      )




						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[int]  ,A_     :      Union[str, Any]  ,A_     :      List[Tuple] = None      )     ->   str:
							A						   =   outputs.logits

							# Resize logits and compute semantic segmentation maps
							if target_sizes is not None:
								if len(A_      ) != len(A_      ):
									raise ValueError(
									    'Make sure that you pass in as many target sizes as the batch dimension of the logits'      )

								if is_torch_tensor(A_      ):
									A						   =   target_sizes.numpy()

								A						   =   []

								for idx in range(len(A_      )      ):
									A						   =   torch.nn.functional.interpolate(
									    logits[idx].unsqueeze(dim=0      )  ,size=target_sizes[idx]  ,mode='bilinear'  ,align_corners=A_      )
									A						   =   resized_logits[0].argmax(dim=0      )
									semantic_segmentation.append(A_      )
							else:
								A						   =   logits.argmax(dim=1      )
								A						   =   [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]      )]

							return semantic_segmentation