
"""simple docstring"""



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

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
 from fairscale.optim import OSS


UpperCAmelCase__          = logging.get_logger(__name__)

UpperCAmelCase__          = {
    'linear': get_linear_schedule_with_warmup,
    'cosine': get_cosine_schedule_with_warmup,
    'cosine_w_restarts': get_cosine_with_hard_restarts_schedule_with_warmup,
    'polynomial': get_polynomial_decay_schedule_with_warmup,
    'constant': get_constant_schedule,
    'constant_w_warmup': get_constant_schedule_with_warmup,
}

class       lowerCAmelCase__				(						__A	):







     def __init__(							self			:			List[Any]    ,      _lowerCamelCase			:			Dict=None    ,      _lowerCamelCase			:			Optional[int]=None    ,      *_lowerCamelCase			:			int    ,      **_lowerCamelCase			:			Dict			):
       super().__init__(*__UpperCAmelCase    ,      **__UpperCAmelCase			)

       if config is None:
         assert isinstance(self.model    ,      __UpperCAmelCase			), (
             "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
             f''' {self.model.__class__}'''
         )
         _snake_case					     =       self.model.config
       else:
         _snake_case					     =       config

       _snake_case					     =       data_args
       _snake_case					     =       self.config.tgt_vocab_size if isinstance(self.config    ,      __UpperCAmelCase			) else self.config.vocab_size

       if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
         assert self.config.pad_token_id is not None, (
             "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
             " calculation or doing label smoothing."
         )

       if self.config.pad_token_id is None and self.config.eos_token_id is not None:
         logger.warning(
             f'''The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for'''
             ''' padding..'''			)

       if self.args.label_smoothing == 0:
         _snake_case					     =       torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id			)
       else:
         # dynamically import label_smoothed_nll_loss
         from utils import label_smoothed_nll_loss

         _snake_case					     =       label_smoothed_nll_loss







     def      lowercase	(							self			:			Optional[int]    ,      _lowerCamelCase			:			Tuple			):
       if self.optimizer is None:
         _snake_case					     =       ['''bias''', '''LayerNorm.weight''']
         _snake_case					     =       [
             {
                 '''params''': [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay			)],
                 '''weight_decay''': self.args.weight_decay,
             },
             {
                 '''params''': [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay			)],
                 '''weight_decay''': 0.0,
             },
         ]
         _snake_case					     =       Adafactor if self.args.adafactor else AdamW
         if self.args.adafactor:
           _snake_case					     =       Adafactor
           _snake_case					     =       {'''scale_parameter''': False, '''relative_step''': False}
         else:
           _snake_case					     =       AdamW
           _snake_case					     =       {
               '''betas''': (self.args.adam_betaa, self.args.adam_betaa),
               '''eps''': self.args.adam_epsilon,
           }
         _snake_case					     =       self.args.learning_rate
         if self.sharded_ddp:
           _snake_case					     =       OSS(
               params=__UpperCAmelCase    ,      optim=__UpperCAmelCase    ,      **__UpperCAmelCase    ,      )
         else:
           _snake_case					     =       optimizer_cls(__UpperCAmelCase    ,      **__UpperCAmelCase			)

       if self.lr_scheduler is None:
         _snake_case					     =       self._get_lr_scheduler(__UpperCAmelCase			)
       else:  # ignoring --lr_scheduler
         logger.warning('''scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'''			)







     def      lowercase	(							self			:			List[str]    ,      _lowerCamelCase			:			Optional[int]			):
       _snake_case					     =       arg_to_scheduler[self.args.lr_scheduler]
       if self.args.lr_scheduler == "constant":
         _snake_case					     =       schedule_func(self.optimizer			)
       elif self.args.lr_scheduler == "constant_w_warmup":
         _snake_case					     =       schedule_func(self.optimizer    ,      num_warmup_steps=self.args.warmup_steps			)
       else:
         _snake_case					     =       schedule_func(
             self.optimizer    ,      num_warmup_steps=self.args.warmup_steps    ,      num_training_steps=__UpperCAmelCase			)
       return scheduler







     def      lowercase	(							self			:			Any			):
       if isinstance(self.train_dataset    ,      torch.utils.data.IterableDataset			):
         return None
       elif is_torch_tpu_available():
         return get_tpu_sampler(self.train_dataset			)
       else:
         if self.args.sortish_sampler:
           self.train_dataset.make_sortish_sampler(
               self.args.per_device_train_batch_size    ,      distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED)    ,      )

         return (
             RandomSampler(self.train_dataset			)
             if self.args.local_rank == -1
             else DistributedSampler(self.train_dataset			)
         )







     def      lowercase	(							self			:			Dict    ,      _lowerCamelCase			:			List[Any]    ,      _lowerCamelCase			:			Optional[Any]    ,      _lowerCamelCase			:			List[str]			):
       if self.args.label_smoothing == 0:
         if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
           # force training to ignore pad token
           _snake_case					     =       model(**__UpperCAmelCase    ,      use_cache=__UpperCAmelCase			)[0]
           _snake_case					     =       self.loss_fn(logits.view(-1    ,      logits.shape[-1]			)    ,      labels.view(-1			)			)
         else:
           # compute usual loss via models
           _snake_case     ,					_snake_case					     =       model(**__UpperCAmelCase    ,      labels=__UpperCAmelCase    ,      use_cache=__UpperCAmelCase			)[:2]
       else:
         # compute label smoothed loss
         _snake_case					     =       model(**__UpperCAmelCase    ,      use_cache=__UpperCAmelCase			)[0]
         _snake_case					     =       torch.nn.functional.log_softmax(__UpperCAmelCase    ,      dim=-1			)
         _snake_case     ,					_snake_case					     =       self.loss_fn(__UpperCAmelCase    ,      __UpperCAmelCase    ,      self.args.label_smoothing    ,      ignore_index=self.config.pad_token_id			)
       return loss, logits







     def      lowercase	(							self			:			Optional[int]    ,      _lowerCamelCase			:			Any    ,      _lowerCamelCase			:			Optional[Any]			):
       _snake_case					     =       inputs.pop('''labels'''			)
       _snake_case     ,					_snake_case					     =       self._compute_loss(__UpperCAmelCase    ,      __UpperCAmelCase    ,      __UpperCAmelCase			)
       return loss







     def      lowercase	(							self			:			Optional[Any]    ,      _lowerCamelCase			:			Optional[int]    ,      _lowerCamelCase			:			Tuple    ,      _lowerCamelCase			:			Optional[Any]    ,      _lowerCamelCase			:			Dict = None    ,      ):
       _snake_case					     =       self._prepare_inputs(__UpperCAmelCase			)

       _snake_case					     =       {
           '''max_length''': self.data_args.val_max_target_length
           if self.data_args is not None
           else self.config.max_length,
           '''num_beams''': self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
       }

       if self.args.predict_with_generate and not self.args.prediction_loss_only:
         _snake_case					     =       self.model.generate(
             inputs['''input_ids''']    ,      attention_mask=inputs['''attention_mask''']    ,      **__UpperCAmelCase    ,      )
         # in case the batch is shorter than max length, the output should be padded
         if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
           _snake_case					     =       self._pad_tensors_to_max_len(__UpperCAmelCase    ,      gen_kwargs['''max_length''']			)

       _snake_case					     =       inputs.pop('''labels'''			)
       with torch.no_grad():
         # compute loss on predict data
         _snake_case     ,					_snake_case					     =       self._compute_loss(__UpperCAmelCase    ,      __UpperCAmelCase    ,      __UpperCAmelCase			)

       _snake_case					     =       loss.mean().detach()
       if self.args.prediction_loss_only:
         return (loss, None, None)

       _snake_case					     =       generated_tokens if self.args.predict_with_generate else logits

       if labels.shape[-1] < gen_kwargs["max_length"]:
         _snake_case					     =       self._pad_tensors_to_max_len(__UpperCAmelCase    ,      gen_kwargs['''max_length''']			)

       return (loss, logits, labels)







     def      lowercase	(							self			:			int    ,      _lowerCamelCase			:			Dict    ,      _lowerCamelCase			:			int			):
       # If PAD token is not defined at least EOS token has to be defined
       _snake_case					     =       self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

       if pad_token_id is None:
         raise ValueError(
             '''Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be'''
             f''' padded to `max_length`={max_length}'''			)

       _snake_case					     =       pad_token_id * torch.ones(
           (tensor.shape[0], max_length)    ,      dtype=tensor.dtype    ,      device=tensor.device			)
       _snake_case					     =       tensor
       return padded_tensor




"""simple docstring"""




import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import 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 transformers import UperNetForSemanticSegmentation
				from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
				from PIL import Image

				from transformers import AutoImageProcessor
class 			UpperCamelCase__:





						def __init__(				self						,__UpperCAmelCase						,__UpperCAmelCase=13						,__UpperCAmelCase=32						,__UpperCAmelCase=3						,__UpperCAmelCase=4						,__UpperCAmelCase=[10, 20, 30, 40]						,__UpperCAmelCase=[2, 2, 3, 2]						,__UpperCAmelCase=True						,__UpperCAmelCase=True						,__UpperCAmelCase=37						,__UpperCAmelCase="gelu"						,__UpperCAmelCase=10						,__UpperCAmelCase=0.0_2						,__UpperCAmelCase=["stage2", "stage3", "stage4"]						,__UpperCAmelCase=3						,__UpperCAmelCase=None						,)    -> Optional[int]:
													A__		   =							parent
													A__		   =							batch_size
													A__		   =							image_size
													A__		   =							num_channels
													A__		   =							num_stages
													A__		   =							hidden_sizes
													A__		   =							depths
													A__		   =							is_training
													A__		   =							use_labels
													A__		   =							intermediate_size
													A__		   =							hidden_act
													A__		   =							type_sequence_label_size
													A__		   =							initializer_range
													A__		   =							out_features
													A__		   =							num_labels
													A__		   =							scope
													A__		   =							num_stages





						def  snake_case__		(				self   )    -> List[Any]:
													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.type_sequence_label_size   )

													A__		   =							self.get_config()

													return config, pixel_values, labels





						def  snake_case__		(				self   )    -> str:
													return ConvNextConfig(
													    num_channels=self.num_channels						,num_stages=self.num_stages						,hidden_sizes=self.hidden_sizes						,depths=self.depths						,is_training=self.is_training						,intermediate_size=self.intermediate_size						,hidden_act=self.hidden_act						,out_features=self.out_features						,)





						def  snake_case__		(				self   )    -> Tuple:
													return UperNetConfig(
													    backbone_config=self.get_backbone_config()						,hidden_size=5_12						,pool_scales=[1, 2, 3, 6]						,use_auxiliary_head=__UpperCAmelCase						,auxiliary_loss_weight=0.4						,auxiliary_in_channels=40						,auxiliary_channels=2_56						,auxiliary_num_convs=1						,auxiliary_concat_input=__UpperCAmelCase						,loss_ignore_index=2_55						,num_labels=self.num_labels						,)





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )    -> Dict:
													A__		   =							UperNetForSemanticSegmentation(config=__UpperCAmelCase   )
													model.to(__UpperCAmelCase   )
													model.eval()
													A__		   =							model(__UpperCAmelCase   )
													self.parent.assertEqual(
													    result.logits.shape						,(self.batch_size, self.num_labels, self.image_size, self.image_size)   )





						def  snake_case__		(				self   )    -> str:
													A__		   =							self.prepare_config_and_inputs()
													(
													    (
													    A__
													)   ,     (
													    A__
													)   ,     (
													    A__
													)   ,     
													)		   =							config_and_inputs
													A__		   =							{'pixel_values': pixel_values}
													return config, inputs_dict






@require_torch
class 			UpperCamelCase__(	__A						,  __A						,  unittest.TestCase       ):
						lowerCAmelCase__     :  int  							=						(UperNetForSemanticSegmentation,) if is_torch_available() else ()
						lowerCAmelCase__     :  int  							=						{'image-segmentation': UperNetForSemanticSegmentation} if is_torch_available() else {}
						lowerCAmelCase__     :  Optional[int]  							=						False
						lowerCAmelCase__     :  List[Any]  							=						False
						lowerCAmelCase__     :  Tuple  							=						False
						lowerCAmelCase__     :  Optional[Any]  							=						False
						lowerCAmelCase__     :  Union[str, Any]  							=						False
						lowerCAmelCase__     :  Dict  							=						False





						def  snake_case__		(				self   )    -> Union[str, Any]:
													A__		   =							UperNetModelTester(self   )
													A__		   =							ConfigTester(self						,config_class=__UpperCAmelCase						,has_text_modality=__UpperCAmelCase						,hidden_size=37   )





						def  snake_case__		(				self   )    -> List[Any]:
													self.create_and_test_config_common_properties()
													self.config_tester.create_and_test_config_to_json_string()
													self.config_tester.create_and_test_config_to_json_file()
													self.config_tester.create_and_test_config_from_and_save_pretrained()
													self.config_tester.create_and_test_config_with_num_labels()
													self.config_tester.check_config_can_be_init_without_params()
													self.config_tester.check_config_arguments_init()





						def  snake_case__		(				self   )    -> int:
													return





						def  snake_case__		(				self   )    -> List[Any]:
													A__   ,     A__		   =							self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				A__		   =							model_class(__UpperCAmelCase   )
																				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]						,__UpperCAmelCase   )





						def  snake_case__		(				self   )    -> Tuple:
													A__		   =							self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_semantic_segmentation(*__UpperCAmelCase   )





						@unittest.skip(reason='UperNet does not use inputs_embeds'   )
						def  snake_case__		(				self   )    -> Optional[int]:
													pass





						@unittest.skip(reason='UperNet does not support input and output embeddings'   )
						def  snake_case__		(				self   )    -> Tuple:
													pass





						@unittest.skip(reason='UperNet does not have a base model'   )
						def  snake_case__		(				self   )    -> List[Any]:
													pass





						@unittest.skip(reason='UperNet does not have a base model'   )
						def  snake_case__		(				self   )    -> Dict:
													pass





						@require_torch_multi_gpu
						@unittest.skip(reason='UperNet has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'   )
						def  snake_case__		(				self   )    -> Any:
													pass





						@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'   )
						def  snake_case__		(				self   )    -> Dict:
													pass





						def  snake_case__		(				self   )    -> Optional[int]:
													def check_hidden_states_output(__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   ):
																				A__		   =							model_class(__UpperCAmelCase   )
																				model.to(__UpperCAmelCase   )
																				model.eval()

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

																				A__		   =							outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																				A__		   =							self.model_tester.num_stages
																				self.assertEqual(len(__UpperCAmelCase   )						,expected_num_stages + 1   )

																				# ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
																				self.assertListEqual(
																				    list(hidden_states[0].shape[-2:]   )						,[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(__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )

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

																				check_hidden_states_output(__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )





						def  snake_case__		(				self   )    -> Optional[Any]:
													A__   ,     A__		   =							self.model_tester.prepare_config_and_inputs_for_common()

													A__		   =							_config_zero_init(__UpperCAmelCase   )
													A__		   =							_config_zero_init(configs_no_init.backbone_config   )
													for model_class in self.all_model_classes:
																				A__		   =							model_class(config=__UpperCAmelCase   )
																				for name, param in model.named_parameters():
																											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'''						,)





						@unittest.skip(reason='UperNet does not have tied weights'   )
						def  snake_case__		(				self   )    -> str:
													pass





						@slow
						def  snake_case__		(				self   )    -> List[str]:
													for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				A__		   =							UperNetForSemanticSegmentation.from_pretrained(__UpperCAmelCase   )
																				self.assertIsNotNone(__UpperCAmelCase   )






def      UpperCAmelCase    (		):






							"""simple docstring"""





							A__		   =							hf_hub_download(
							    repo_id='hf-internal-testing/fixtures_ade20k'				,			repo_type='dataset'				,			filename='ADE_val_00000001.jpg' )
							A__		   =							Image.open(UpperCamelCase__ ).convert('RGB' )
							return image


@require_torch
@require_vision
@slow
class 			UpperCamelCase__(	unittest.TestCase       ):





						def  snake_case__		(				self   )    -> Dict:
													A__		   =							AutoImageProcessor.from_pretrained('openmmlab/upernet-swin-tiny'   )
													A__		   =							UperNetForSemanticSegmentation.from_pretrained('openmmlab/upernet-swin-tiny'   ).to(__UpperCAmelCase   )

													A__		   =							prepare_img()
													A__		   =							processor(images=__UpperCAmelCase						,return_tensors='pt'   ).to(__UpperCAmelCase   )

													with torch.no_grad():
																				A__		   =							model(**__UpperCAmelCase   )

													A__		   =							torch.Size((1, model.config.num_labels, 5_12, 5_12)   )
													self.assertEqual(outputs.logits.shape						,__UpperCAmelCase   )

													A__		   =							torch.tensor(
													    [[-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.4_7_9_7, -7.4_7_9_7, -7.3_0_6_8]]   ).to(__UpperCAmelCase   )
													self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]						,__UpperCAmelCase						,atol=1e-4   )   )





						def  snake_case__		(				self   )    -> str:
													A__		   =							AutoImageProcessor.from_pretrained('openmmlab/upernet-convnext-tiny'   )
													A__		   =							UperNetForSemanticSegmentation.from_pretrained('openmmlab/upernet-convnext-tiny'   ).to(__UpperCAmelCase   )

													A__		   =							prepare_img()
													A__		   =							processor(images=__UpperCAmelCase						,return_tensors='pt'   ).to(__UpperCAmelCase   )

													with torch.no_grad():
																				A__		   =							model(**__UpperCAmelCase   )

													A__		   =							torch.Size((1, model.config.num_labels, 5_12, 5_12)   )
													self.assertEqual(outputs.logits.shape						,__UpperCAmelCase   )

													A__		   =							torch.tensor(
													    [[-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.7_7_4_6, -8.7_7_4_6, -8.6_1_3_0]]   ).to(__UpperCAmelCase   )
													self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]						,__UpperCAmelCase						,atol=1e-4   )   )





