



import inspect
import unittest

from transformers import SegformerConfig, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
							import torch

							from transformers import (
							    MODEL_MAPPING,
							    SegformerForImageClassification,
							    SegformerForSemanticSegmentation,
							    SegformerModel,
							)
							from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image

							from transformers import SegformerImageProcessor


class      __lowerCAmelCase     (		UpperCAmelCase_):
					"""simple docstring"""






					def     _a							(     self					:       Optional[int]					):






								"""simple docstring"""


								A__       =    self.config_class(**self.inputs_dict					)
								self.parent.assertTrue(hasattr(_snake_case				,					'hidden_sizes'					)					)
								self.parent.assertTrue(hasattr(_snake_case				,					'num_attention_heads'					)					)
								self.parent.assertTrue(hasattr(_snake_case				,					'num_encoder_blocks'					)					)







class      __lowerCAmelCase     :
					"""simple docstring"""






					def __init__(     self					:       List[Any]				,					_snake_case					:       Union[str, Any]				,					_snake_case					:       Union[str, Any]=13				,					_snake_case					:       Union[str, Any]=64				,					_snake_case					:       int=3				,					_snake_case					:       Dict=4				,					_snake_case					:       Dict=[2, 2, 2, 2]				,					_snake_case					:       Dict=[8, 4, 2, 1]				,					_snake_case					:       Any=[16, 32, 64, 1_28]				,					_snake_case					:       int=[1, 4, 8, 16]				,					_snake_case					:       Any=[1, 2, 4, 8]				,					_snake_case					:       Union[str, Any]=True				,					_snake_case					:       List[Any]=True				,					_snake_case					:       Any="gelu"				,					_snake_case					:       str=0.1				,					_snake_case					:       Any=0.1				,					_snake_case					:       Optional[int]=0.02				,					_snake_case					:       str=3				,					_snake_case					:       str=None				,					):






								"""simple docstring"""


								A__       =    parent
								A__       =    batch_size
								A__       =    image_size
								A__       =    num_channels
								A__       =    num_encoder_blocks
								A__       =    sr_ratios
								A__       =    depths
								A__       =    hidden_sizes
								A__       =    downsampling_rates
								A__       =    num_attention_heads
								A__       =    is_training
								A__       =    use_labels
								A__       =    hidden_act
								A__       =    hidden_dropout_prob
								A__       =    attention_probs_dropout_prob
								A__       =    initializer_range
								A__       =    num_labels
								A__       =    scope



					def     _a							(     self					:       Optional[Any]					):






								"""simple docstring"""


								A__       =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

								A__       =    None
								if self.use_labels:
											A__       =    ids_tensor([self.batch_size, self.image_size, self.image_size]				,					self.num_labels					)

								A__       =    self.get_config()
								return config, pixel_values, labels



					def     _a							(     self					:       List[Any]					):






								"""simple docstring"""


								return SegformerConfig(
								    image_size=self.image_size				,					num_channels=self.num_channels				,					num_encoder_blocks=self.num_encoder_blocks				,					depths=self.depths				,					hidden_sizes=self.hidden_sizes				,					num_attention_heads=self.num_attention_heads				,					hidden_act=self.hidden_act				,					hidden_dropout_prob=self.hidden_dropout_prob				,					attention_probs_dropout_prob=self.attention_probs_dropout_prob				,					initializer_range=self.initializer_range				,					)



					def     _a							(     self					:       Union[str, Any]				,					_snake_case					:       List[Any]				,					_snake_case					:       Any				,					_snake_case					:       Union[str, Any]					):






								"""simple docstring"""


								A__       =    SegformerModel(config=_snake_case					)
								model.to(_snake_case					)
								model.eval()
								A__       =    model(_snake_case					)
								A__       =    A__       =    self.image_size // (self.downsampling_rates[-1] * 2)
								self.parent.assertEqual(
								    result.last_hidden_state.shape				,					(self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)					)



					def     _a							(     self					:       int				,					_snake_case					:       Optional[int]				,					_snake_case					:       Optional[Any]				,					_snake_case					:       Optional[int]					):






								"""simple docstring"""


								A__       =    self.num_labels
								A__       =    SegformerForSemanticSegmentation(_snake_case					)
								model.to(_snake_case					)
								model.eval()
								A__       =    model(_snake_case					)
								self.parent.assertEqual(
								    result.logits.shape				,					(self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)					)
								A__       =    model(_snake_case				,					labels=_snake_case					)
								self.parent.assertEqual(
								    result.logits.shape				,					(self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)					)
								self.parent.assertGreater(result.loss				,					0.0					)



					def     _a							(     self					:       int				,					_snake_case					:       Union[str, Any]				,					_snake_case					:       Tuple				,					_snake_case					:       Any					):






								"""simple docstring"""


								A__       =    1
								A__       =    SegformerForSemanticSegmentation(config=_snake_case					)
								model.to(_snake_case					)
								model.eval()
								A__       =    torch.randint(0				,					1				,					(self.batch_size, self.image_size, self.image_size)					).to(_snake_case					)
								A__       =    model(_snake_case				,					labels=_snake_case					)
								self.parent.assertGreater(result.loss				,					0.0					)



					def     _a							(     self					:       Optional[int]					):






								"""simple docstring"""


								A__       =    self.prepare_config_and_inputs()
								A__     , A__     , A__       =    config_and_inputs
								A__       =    {'pixel_values': pixel_values}
								return config, inputs_dict







@require_torch
class      __lowerCAmelCase     (		UpperCAmelCase_							,     UpperCAmelCase_							,     unittest.TestCase):
					"""simple docstring"""






					A__				:		Dict			       =    (
					    (
					        SegformerModel,
					        SegformerForSemanticSegmentation,
					        SegformerForImageClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					A__				:		str			       =    (
					    {
					        "feature-extraction": SegformerModel,
					        "image-classification": SegformerForImageClassification,
					        "image-segmentation": SegformerForSemanticSegmentation,
					    }
					    if is_torch_available()
					    else {}
					)

					A__				:		Any			       =    True
					A__				:		List[Any]			       =    False
					A__				:		Any			       =    False
					A__				:		Union[str, Any]			       =    False



					def     _a							(     self					:       str					):






								"""simple docstring"""


								A__       =    SegformerModelTester(self					)
								A__       =    SegformerConfigTester(self				,					config_class=_snake_case					)



					def     _a							(     self					:       int					):






								"""simple docstring"""


								self.config_tester.run_common_tests()



					def     _a							(     self					:       Optional[Any]					):






								"""simple docstring"""


								A__       =    self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case					)



					def     _a							(     self					:       Dict					):






								"""simple docstring"""


								A__       =    self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_binary_image_segmentation(*_snake_case					)



					def     _a							(     self					:       int					):






								"""simple docstring"""


								A__       =    self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_segmentation(*_snake_case					)



					@unittest.skip('SegFormer does not use inputs_embeds'					)
					def     _a							(     self					:       Dict					):






								"""simple docstring"""


								pass



					@unittest.skip('SegFormer does not have get_input_embeddings method and get_output_embeddings methods'					)
					def     _a							(     self					:       Dict					):






								"""simple docstring"""


								pass



					def     _a							(     self					:       str					):






								"""simple docstring"""


								A__     , A__       =    self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											A__       =    model_class(_snake_case					)
											A__       =    inspect.signature(model.forward					)
											# signature.parameters is an OrderedDict => so arg_names order is deterministic
											A__       =    [*signature.parameters.keys()]

											A__       =    ['pixel_values']
											self.assertListEqual(arg_names[:1]				,					_snake_case					)



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






								"""simple docstring"""


								A__     , A__       =    self.model_tester.prepare_config_and_inputs_for_common()
								A__       =    True

								for model_class in self.all_model_classes:
											A__       =    True
											A__       =    False
											A__       =    True
											A__       =    model_class(_snake_case					)
											model.to(_snake_case					)
											model.eval()
											with torch.no_grad():
														A__       =    model(**self._prepare_for_class(_snake_case				,					_snake_case					)					)
											A__       =    outputs.attentions

											A__       =    sum(self.model_tester.depths					)
											self.assertEqual(len(_snake_case					)				,					_snake_case					)

											# check that output_attentions also work using config
											del inputs_dict["output_attentions"]
											A__       =    True
											A__       =    model_class(_snake_case					)
											model.to(_snake_case					)
											model.eval()
											with torch.no_grad():
														A__       =    model(**self._prepare_for_class(_snake_case				,					_snake_case					)					)
											A__       =    outputs.attentions

											self.assertEqual(len(_snake_case					)				,					_snake_case					)

											# verify the first attentions (first block, first layer)
											A__       =    (self.model_tester.image_size // 4) ** 2
											A__       =    (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
											self.assertListEqual(
											    list(attentions[0].shape[-3:]					)				,					[self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]				,					)

											# verify the last attentions (last block, last layer)
											A__       =    (self.model_tester.image_size // 32) ** 2
											A__       =    (self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
											self.assertListEqual(
											    list(attentions[-1].shape[-3:]					)				,					[self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]				,					)
											A__       =    len(_snake_case					)

											# Check attention is always last and order is fine
											A__       =    True
											A__       =    True
											A__       =    model_class(_snake_case					)
											model.to(_snake_case					)
											model.eval()
											with torch.no_grad():
														A__       =    model(**self._prepare_for_class(_snake_case				,					_snake_case					)					)

											self.assertEqual(out_len + 1				,					len(_snake_case					)					)

											A__       =    outputs.attentions

											self.assertEqual(len(_snake_case					)				,					_snake_case					)
											# verify the first attentions (first block, first layer)
											A__       =    (self.model_tester.image_size // 4) ** 2
											A__       =    (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
											self.assertListEqual(
											    list(self_attentions[0].shape[-3:]					)				,					[self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]				,					)



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






								"""simple docstring"""


								def check_hidden_states_output(_snake_case					:       Union[str, Any]				,					_snake_case					:       Union[str, Any]				,					_snake_case					:       str					):
											A__       =    model_class(_snake_case					)
											model.to(_snake_case					)
											model.eval()

											with torch.no_grad():
														A__       =    model(**self._prepare_for_class(_snake_case				,					_snake_case					)					)

											A__       =    outputs.hidden_states

											A__       =    self.model_tester.num_encoder_blocks
											self.assertEqual(len(_snake_case					)				,					_snake_case					)

											# verify the first hidden states (first block)
											self.assertListEqual(
											    list(hidden_states[0].shape[-3:]					)				,					[
											        self.model_tester.hidden_sizes[0],
											        self.model_tester.image_size // 4,
											        self.model_tester.image_size // 4,
											    ]				,					)

								A__     , A__       =    self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											A__       =    True
											check_hidden_states_output(_snake_case				,					_snake_case				,					_snake_case					)

											# check that output_hidden_states also work using config
											del inputs_dict["output_hidden_states"]
											A__       =    True

											check_hidden_states_output(_snake_case				,					_snake_case				,					_snake_case					)



					def     _a							(     self					:       Tuple					):






								"""simple docstring"""


								if not self.model_tester.is_training:
											return

								A__     , A__       =    self.model_tester.prepare_config_and_inputs_for_common()
								A__       =    True

								for model_class in self.all_model_classes:
											if model_class in get_values(_snake_case					):
														continue

											A__       =    model_class(_snake_case					)
											model.to(_snake_case					)
											model.train()
											A__       =    self._prepare_for_class(_snake_case				,					_snake_case				,					return_labels=_snake_case					)
											A__       =    model(**_snake_case					).loss
											loss.backward()



					@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'					)
					def     _a							(     self					:       Optional[Any]					):






								"""simple docstring"""


								pass



					@slow
					def     _a							(     self					:       Dict					):






								"""simple docstring"""


								for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											A__       =    SegformerModel.from_pretrained(_snake_case					)
											self.assertIsNotNone(_snake_case					)







def A						(       )						->      Optional[Any]:
			A__       =    Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'					)
			return image






@require_torch
class      __lowerCAmelCase     (		unittest.TestCase):
					"""simple docstring"""






					@slow
					def     _a							(     self					:       Dict					):






								"""simple docstring"""


								A__       =    SegformerImageProcessor(
								    image_scale=(5_12, 5_12)				,					keep_ratio=_snake_case				,					align=_snake_case				,					do_random_crop=_snake_case					)
								A__       =    SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512'					).to(
								    _snake_case					)

								A__       =    prepare_img()
								A__       =    image_processor(images=_snake_case				,					return_tensors='pt'					)
								A__       =    encoded_inputs.pixel_values.to(_snake_case					)

								with torch.no_grad():
											A__       =    model(_snake_case					)

								A__       =    torch.Size((1, model.config.num_labels, 1_28, 1_28)					)
								self.assertEqual(outputs.logits.shape				,					_snake_case					)

								A__       =    torch.tensor(
								    [
								        [[-4.6310, -5.5232, -6.2356], [-5.1921, -6.1444, -6.5996], [-5.4424, -6.2790, -6.7574]],
								        [[-12.1391, -13.3122, -13.9554], [-12.8732, -13.9352, -14.3563], [-12.9438, -13.8226, -14.2513]],
								        [[-12.5134, -13.4686, -14.4915], [-12.8669, -14.4343, -14.7758], [-13.2523, -14.5819, -15.0694]],
								    ]					).to(_snake_case					)
								self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]				,					_snake_case				,					atol=1E-4					)					)



					@slow
					def     _a							(     self					:       List[Any]					):






								"""simple docstring"""


								A__       =    SegformerImageProcessor(
								    image_scale=(5_12, 5_12)				,					keep_ratio=_snake_case				,					align=_snake_case				,					do_random_crop=_snake_case					)
								A__       =    SegformerForSemanticSegmentation.from_pretrained(
								    'nvidia/segformer-b1-finetuned-cityscapes-1024-1024'					).to(_snake_case					)

								A__       =    prepare_img()
								A__       =    image_processor(images=_snake_case				,					return_tensors='pt'					)
								A__       =    encoded_inputs.pixel_values.to(_snake_case					)

								with torch.no_grad():
											A__       =    model(_snake_case					)

								A__       =    torch.Size((1, model.config.num_labels, 1_28, 1_28)					)
								self.assertEqual(outputs.logits.shape				,					_snake_case					)

								A__       =    torch.tensor(
								    [
								        [[-13.5748, -13.9111, -12.6500], [-14.3500, -15.3683, -14.2328], [-14.7532, -16.0424, -15.6087]],
								        [[-17.1651, -15.8725, -12.9653], [-17.2580, -17.3718, -14.8223], [-16.6058, -16.8783, -16.7452]],
								        [[-3.6456, -3.0209, -1.4203], [-3.0797, -3.1959, -2.0000], [-1.8757, -1.9217, -1.6997]],
								    ]					).to(_snake_case					)
								self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]				,					_snake_case				,					atol=1E-1					)					)



					@slow
					def     _a							(     self					:       List[Any]					):






								"""simple docstring"""


								A__       =    SegformerImageProcessor(
								    image_scale=(5_12, 5_12)				,					keep_ratio=_snake_case				,					align=_snake_case				,					do_random_crop=_snake_case					)
								A__       =    SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512'					).to(
								    _snake_case					)

								A__       =    prepare_img()
								A__       =    image_processor(images=_snake_case				,					return_tensors='pt'					)
								A__       =    encoded_inputs.pixel_values.to(_snake_case					)

								with torch.no_grad():
											A__       =    model(_snake_case					)

								A__       =    outputs.logits.detach().cpu()

								A__       =    image_processor.post_process_semantic_segmentation(outputs=_snake_case				,					target_sizes=[(5_00, 3_00)]					)
								A__       =    torch.Size((5_00, 3_00)					)
								self.assertEqual(segmentation[0].shape				,					_snake_case					)

								A__       =    image_processor.post_process_semantic_segmentation(outputs=_snake_case					)
								A__       =    torch.Size((1_28, 1_28)					)
								self.assertEqual(segmentation[0].shape				,					_snake_case					)









import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
							import optax
							from flax.training.common_utils import onehot

							from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
							from transformers.models.ta.modeling_flax_ta import shift_tokens_right


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






					@slow
					def     _a							(     self					:       int					):






								"""simple docstring"""


								A__       =    FlaxMTaForConditionalGeneration.from_pretrained('google/mt5-small'					)
								A__       =    AutoTokenizer.from_pretrained('google/mt5-small'					)

								A__       =    tokenizer('Hello there'				,					return_tensors='np'					).input_ids
								A__       =    tokenizer('Hi I am'				,					return_tensors='np'					).input_ids

								A__       =    shift_tokens_right(_snake_case				,					model.config.pad_token_id				,					model.config.decoder_start_token_id					)

								A__       =    model(_snake_case				,					decoder_input_ids=_snake_case					).logits
								A__       =    optax.softmax_cross_entropy(_snake_case				,					onehot(_snake_case				,					logits.shape[-1]					)					).mean()

								A__       =    -(labels.shape[-1] * loss.item())

								A__       =    -84.9127
								self.assertTrue(abs(mtf_score - EXPECTED_SCORE					) < 1E-4					)




