from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
				from transformers import (
				    TFAutoModel,
				    TFAutoModelForCausalLM,
				    TFAutoModelForMaskedLM,
				    TFAutoModelForPreTraining,
				    TFAutoModelForQuestionAnswering,
				    TFAutoModelForSeqaSeqLM,
				    TFAutoModelForSequenceClassification,
				    TFAutoModelForTableQuestionAnswering,
				    TFAutoModelForTokenClassification,
				    TFAutoModelWithLMHead,
				    TFBertForMaskedLM,
				    TFBertForPreTraining,
				    TFBertForQuestionAnswering,
				    TFBertForSequenceClassification,
				    TFBertModel,
				    TFFunnelBaseModel,
				    TFFunnelModel,
				    TFGPTaLMHeadModel,
				    TFRobertaForMaskedLM,
				    TFTaForConditionalGeneration,
				    TFTapasForQuestionAnswering,
				)
				from transformers.models.auto.modeling_tf_auto import (
				    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
				    TF_MODEL_FOR_MASKED_LM_MAPPING,
				    TF_MODEL_FOR_PRETRAINING_MAPPING,
				    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
				    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
				    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
				    TF_MODEL_MAPPING,
				)
				from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
				from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
				from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
				from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST



class        lowercase_   (					__snake_case		):
							_lowerCamelCase						=		'new-model'


if is_tf_available():



				class        lowercase_   (					__snake_case		):
											_lowerCamelCase						=		NewModelConfig



@require_tf
class        lowercase_   (					unittest.TestCase		):



							@slow
							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Dict    =   "bert-base-cased"
													_snake_case						:					str    =   AutoConfig.from_pretrained(lowercase_    )
													self.assertIsNotNone(lowercase_    )
													self.assertIsInstance(lowercase_				,    lowercase_    )

													_snake_case						:					Optional[int]    =   TFAutoModel.from_pretrained(lowercase_    )
													self.assertIsNotNone(lowercase_    )
													self.assertIsInstance(lowercase_				,    lowercase_    )



							@slow
							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Optional[int]    =   "bert-base-cased"
													_snake_case						:					List[Any]    =   AutoConfig.from_pretrained(lowercase_    )
													self.assertIsNotNone(lowercase_    )
													self.assertIsInstance(lowercase_				,    lowercase_    )

													_snake_case						:					List[str]    =   TFAutoModelForPreTraining.from_pretrained(lowercase_    )
													self.assertIsNotNone(lowercase_    )
													self.assertIsInstance(lowercase_				,    lowercase_    )



							@slow
							def 	UpperCamelCase	(						self    ):
													for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			_snake_case						:					List[Any]    =   AutoConfig.from_pretrained(lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )

																			_snake_case						:					int    =   TFAutoModelForCausalLM.from_pretrained(lowercase_    )
																			_snake_case    ,_snake_case						:					int    =   TFAutoModelForCausalLM.from_pretrained(lowercase_				,    output_loading_info=lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )



							@slow
							def 	UpperCamelCase	(						self    ):
													for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			_snake_case						:					Optional[Any]    =   AutoConfig.from_pretrained(lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )

																			_snake_case						:					Dict    =   TFAutoModelWithLMHead.from_pretrained(lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )



							@slow
							def 	UpperCamelCase	(						self    ):
													for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			_snake_case						:					Optional[Any]    =   AutoConfig.from_pretrained(lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )

																			_snake_case						:					Dict    =   TFAutoModelForMaskedLM.from_pretrained(lowercase_    )
																			_snake_case    ,_snake_case						:					str    =   TFAutoModelForMaskedLM.from_pretrained(lowercase_				,    output_loading_info=lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )



							@slow
							def 	UpperCamelCase	(						self    ):
													for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			_snake_case						:					int    =   AutoConfig.from_pretrained(lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )

																			_snake_case						:					Tuple    =   TFAutoModelForSeqaSeqLM.from_pretrained(lowercase_    )
																			_snake_case    ,_snake_case						:					Dict    =   TFAutoModelForSeqaSeqLM.from_pretrained(lowercase_				,    output_loading_info=lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )



							@slow
							def 	UpperCamelCase	(						self    ):
													# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													for model_name in ["bert-base-uncased"]:
																			_snake_case						:					List[Any]    =   AutoConfig.from_pretrained(lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )

																			_snake_case						:					Any    =   TFAutoModelForSequenceClassification.from_pretrained(lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )



							@slow
							def 	UpperCamelCase	(						self    ):
													# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													for model_name in ["bert-base-uncased"]:
																			_snake_case						:					int    =   AutoConfig.from_pretrained(lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )

																			_snake_case						:					List[Any]    =   TFAutoModelForQuestionAnswering.from_pretrained(lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )



							@slow
							@require_tensorflow_probability
							def 	UpperCamelCase	(						self    ):
													for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
																			_snake_case						:					Optional[int]    =   AutoConfig.from_pretrained(lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )

																			_snake_case						:					Any    =   TFAutoModelForTableQuestionAnswering.from_pretrained(lowercase_    )
																			_snake_case    ,_snake_case						:					List[Any]    =   TFAutoModelForTableQuestionAnswering.from_pretrained(
																			    lowercase_				,    output_loading_info=lowercase_    )
																			self.assertIsNotNone(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Optional[Any]    =   TFAutoModelWithLMHead.from_pretrained(lowercase_    )
													self.assertIsInstance(lowercase_				,    lowercase_    )
													self.assertEqual(model.num_parameters()				,    14_410    )
													self.assertEqual(model.num_parameters(only_trainable=lowercase_    )				,    14_410    )



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Dict    =   TFAutoModelWithLMHead.from_pretrained(lowercase_    )
													self.assertIsInstance(lowercase_				,    lowercase_    )
													self.assertEqual(model.num_parameters()				,    14_410    )
													self.assertEqual(model.num_parameters(only_trainable=lowercase_    )				,    14_410    )



							def 	UpperCamelCase	(						self    ):
													# For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
													_snake_case						:					List[Any]    =   TFAutoModel.from_pretrained("sgugger/funnel-random-tiny"    )
													self.assertIsInstance(lowercase_				,    lowercase_    )

													_snake_case						:					Optional[int]    =   copy.deepcopy(model.config    )
													_snake_case						:					str    =   ["FunnelBaseModel"]
													_snake_case						:					Any    =   TFAutoModel.from_config(lowercase_    )
													self.assertIsInstance(lowercase_				,    lowercase_    )

													with tempfile.TemporaryDirectory() as tmp_dir:
																			model.save_pretrained(lowercase_    )
																			_snake_case						:					Optional[Any]    =   TFAutoModel.from_pretrained(lowercase_    )
																			self.assertIsInstance(lowercase_				,    lowercase_    )



							def 	UpperCamelCase	(						self    ):
													try:
																			AutoConfig.register("new-model"				,    lowercase_    )

																			_snake_case						:					Optional[Any]    =   [
																			    TFAutoModel,
																			    TFAutoModelForCausalLM,
																			    TFAutoModelForMaskedLM,
																			    TFAutoModelForPreTraining,
																			    TFAutoModelForQuestionAnswering,
																			    TFAutoModelForSequenceClassification,
																			    TFAutoModelForTokenClassification,
																			]

																			for auto_class in auto_classes:
																									with self.subTest(auto_class.__name__    ):
																															# Wrong config class will raise an error
																															with self.assertRaises(lowercase_    ):
																																					auto_class.register(lowercase_				,    lowercase_    )
																															auto_class.register(lowercase_				,    lowercase_    )
																															# Trying to register something existing in the Transformers library will raise an error
																															with self.assertRaises(lowercase_    ):
																																					auto_class.register(lowercase_				,    lowercase_    )

																															# Now that the config is registered, it can be used as any other config with the auto-API
																															_snake_case						:					str    =   BertModelTester(self    ).get_config()
																															_snake_case						:					int    =   NewModelConfig(**tiny_config.to_dict()    )
																															_snake_case						:					str    =   auto_class.from_config(lowercase_    )
																															self.assertIsInstance(lowercase_				,    lowercase_    )

																															with tempfile.TemporaryDirectory() as tmp_dir:
																																					model.save_pretrained(lowercase_    )
																																					_snake_case						:					Optional[Any]    =   auto_class.from_pretrained(lowercase_    )
																																					self.assertIsInstance(lowercase_				,    lowercase_    )

													finally:
																			if "new-model" in CONFIG_MAPPING._extra_content:
																									del CONFIG_MAPPING._extra_content["new-model"]
																			for mapping in (
																			    TF_MODEL_MAPPING,
																			    TF_MODEL_FOR_PRETRAINING_MAPPING,
																			    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
																			    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
																			    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
																			    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
																			    TF_MODEL_FOR_MASKED_LM_MAPPING,
																			):
																									if NewModelConfig in mapping._extra_content:
																															del mapping._extra_content[NewModelConfig]



							def 	UpperCamelCase	(						self    ):
													with self.assertRaisesRegex(
													    lowercase_				,    "bert-base is not a local folder and is not a valid model identifier"    ):
																			_snake_case						:					Tuple    =   TFAutoModel.from_pretrained("bert-base"    )



							def 	UpperCamelCase	(						self    ):
													with self.assertRaisesRegex(
													    lowercase_				,    R"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"    ):
																			_snake_case						:					List[Any]    =   TFAutoModel.from_pretrained(lowercase_				,    revision="aaaaaa"    )



							def 	UpperCamelCase	(						self    ):
													with self.assertRaisesRegex(
													    lowercase_				,    "hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"				,    ):
																			_snake_case						:					Union[str, Any]    =   TFAutoModel.from_pretrained("hf-internal-testing/config-no-model"    )



							def 	UpperCamelCase	(						self    ):
													with self.assertRaisesRegex(lowercase_				,    "Use `from_pt=True` to load this model"    ):
																			_snake_case						:					Union[str, Any]    =   TFAutoModel.from_pretrained("hf-internal-testing/tiny-bert-pt-only"    )



							def 	UpperCamelCase	(						self    ):
													# Make sure we have cached the model.
													_snake_case						:					Union[str, Any]    =   TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert"    )
													with RequestCounter() as counter:
																			_snake_case						:					Optional[Any]    =   TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert"    )
																			self.assertEqual(counter.get_request_count				,    0    )
																			self.assertEqual(counter.head_request_count				,    1    )
																			self.assertEqual(counter.other_request_count				,    0    )

													# With a sharded checkpoint
													_snake_case						:					Dict    =   TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded"    )
													with RequestCounter() as counter:
																			_snake_case						:					Tuple    =   TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded"    )
																			self.assertEqual(counter.get_request_count				,    0    )
																			self.assertEqual(counter.head_request_count				,    1    )
																			self.assertEqual(counter.other_request_count				,    0    )
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,
    flip_channel_order,
    get_resize_output_image_size,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, is_vision_available, logging


if is_vision_available():
				import PIL

if is_torch_available():
				import torch


__SCREAMING_SNAKE_CASE						:   Optional[int]         =			logging.get_logger(__name__)



class        lowercase_   (					__snake_case		):
							_lowerCamelCase						=		['pixel_values']



							def __init__(						self				,    lowercase_ = True				,    lowercase_ = None				,    lowercase_ = PILImageResampling.BILINEAR				,    lowercase_ = True				,    lowercase_ = 1 / 255				,    lowercase_ = True				,    lowercase_ = None				,    lowercase_ = True				,    **lowercase_				,    ):
													super().__init__(**lowercase_    )
													_snake_case						:					Dict    =   size if size is not None else {"shortest_edge": 224}
													_snake_case						:					Optional[int]    =   get_size_dict(lowercase_				,    default_to_square=lowercase_    )
													_snake_case						:					Any    =   crop_size if crop_size is not None else {"height": 256, "width": 256}
													_snake_case						:					str    =   get_size_dict(lowercase_				,    param_name="crop_size"    )

													_snake_case						:					List[Any]    =   do_resize
													_snake_case						:					Tuple    =   size
													_snake_case						:					Union[str, Any]    =   resample
													_snake_case						:					str    =   do_rescale
													_snake_case						:					Dict    =   rescale_factor
													_snake_case						:					int    =   do_center_crop
													_snake_case						:					int    =   crop_size
													_snake_case						:					List[Any]    =   do_flip_channel_order



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_ = PIL.Image.BILINEAR				,    lowercase_ = None				,    **lowercase_				,    ):
													_snake_case						:					Optional[Any]    =   get_size_dict(lowercase_				,    default_to_square=lowercase_    )
													if "shortest_edge" not in size:
																			raise ValueError(f"""The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}"""    )
													_snake_case						:					int    =   get_resize_output_image_size(lowercase_				,    size=size["shortest_edge"]				,    default_to_square=lowercase_    )
													return resize(lowercase_				,    size=lowercase_				,    resample=lowercase_				,    data_format=lowercase_				,    **lowercase_    )



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_ = None				,    **lowercase_				,    ):
													_snake_case						:					List[str]    =   get_size_dict(lowercase_    )
													if "height" not in size or "width" not in size:
																			raise ValueError(f"""The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}"""    )
													return center_crop(lowercase_				,    size=(size["height"], size["width"])				,    data_format=lowercase_				,    **lowercase_    )



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_				,    lowercase_ = None				,    **lowercase_				,    ):
													return rescale(lowercase_				,    scale=lowercase_				,    data_format=lowercase_				,    **lowercase_    )



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_ = None    ):
													return flip_channel_order(lowercase_				,    data_format=lowercase_    )



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_ = None				,    lowercase_ = None				,    lowercase_ = None				,    lowercase_ = None				,    lowercase_ = None				,    lowercase_ = None				,    lowercase_ = None				,    lowercase_ = None				,    lowercase_ = None				,    lowercase_ = ChannelDimension.FIRST				,    **lowercase_				,    ):
													_snake_case						:					Dict    =   do_resize if do_resize is not None else self.do_resize
													_snake_case						:					Union[str, Any]    =   resample if resample is not None else self.resample
													_snake_case						:					Dict    =   do_rescale if do_rescale is not None else self.do_rescale
													_snake_case						:					str    =   rescale_factor if rescale_factor is not None else self.rescale_factor
													_snake_case						:					Optional[int]    =   do_center_crop if do_center_crop is not None else self.do_center_crop
													_snake_case						:					List[str]    =   (
													    do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
													)

													_snake_case						:					Optional[int]    =   size if size is not None else self.size
													_snake_case						:					Optional[int]    =   get_size_dict(lowercase_				,    default_to_square=lowercase_    )
													_snake_case						:					Optional[int]    =   crop_size if crop_size is not None else self.crop_size
													_snake_case						:					Union[str, Any]    =   get_size_dict(lowercase_				,    param_name="crop_size"    )

													_snake_case						:					Tuple    =   make_list_of_images(lowercase_    )

													if not valid_images(lowercase_    ):
																			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_rescale and rescale_factor is None:
																			raise ValueError("Rescale factor must be specified if do_rescale is True."    )

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

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

													if do_resize:
																			_snake_case						:					int    =   [self.resize(image=lowercase_				,    size=lowercase_				,    resample=lowercase_    ) for image in images]

													if do_center_crop:
																			_snake_case						:					List[Any]    =   [self.center_crop(image=lowercase_				,    size=lowercase_    ) for image in images]

													if do_rescale:
																			_snake_case						:					Optional[int]    =   [self.rescale(image=lowercase_				,    scale=lowercase_    ) for image in images]

													# the pretrained checkpoints assume images are BGR, not RGB
													if do_flip_channel_order:
																			_snake_case						:					int    =   [self.flip_channel_order(image=lowercase_    ) for image in images]

													_snake_case						:					str    =   [to_channel_dimension_format(lowercase_				,    lowercase_    ) for image in images]

													_snake_case						:					List[Any]    =   {"pixel_values": images}
													return BatchFeature(data=lowercase_				,    tensor_type=lowercase_    )



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_ = None    ):
													_snake_case						:					int    =   outputs.logits

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

																			if is_torch_tensor(lowercase_    ):
																									_snake_case						:					Union[str, Any]    =   target_sizes.numpy()

																			_snake_case						:					int    =   []

																			for idx in range(len(lowercase_    )    ):
																									_snake_case						:					Optional[Any]    =   torch.nn.functional.interpolate(
																									    logits[idx].unsqueeze(dim=0    )				,    size=target_sizes[idx]				,    mode="bilinear"				,    align_corners=lowercase_    )
																									_snake_case						:					List[str]    =   resized_logits[0].argmax(dim=0    )
																									semantic_segmentation.append(lowercase_    )
													else:
																			_snake_case						:					List[Any]    =   logits.argmax(dim=1    )
																			_snake_case						:					Optional[Any]    =   [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]    )]

													return semantic_segmentation