





import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

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


if is_torch_available():
     import torch

     from transformers import (
         FlaubertForMultipleChoice,
         FlaubertForQuestionAnswering,
         FlaubertForQuestionAnsweringSimple,
         FlaubertForSequenceClassification,
         FlaubertForTokenClassification,
         FlaubertModel,
         FlaubertWithLMHeadModel,
     )
     from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST

class     __SCREAMING_SNAKE_CASE   (       _lowerCAmelCase    ):






       def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase=13					,	__lowerCAmelCase=7					,	__lowerCAmelCase=True					,	__lowerCAmelCase=True					,	__lowerCAmelCase=True					,	__lowerCAmelCase=True					,	__lowerCAmelCase=True					,	__lowerCAmelCase=False					,	__lowerCAmelCase=False					,	__lowerCAmelCase=False					,	__lowerCAmelCase=2					,	__lowerCAmelCase=99					,	__lowerCAmelCase=0					,	__lowerCAmelCase=32					,	__lowerCAmelCase=5					,	__lowerCAmelCase=4					,	__lowerCAmelCase=0.1					,	__lowerCAmelCase=0.1					,	__lowerCAmelCase=512					,	__lowerCAmelCase=12					,	__lowerCAmelCase=2					,	__lowerCAmelCase=0.02					,	__lowerCAmelCase=3					,	__lowerCAmelCase=4					,	__lowerCAmelCase="last"					,	__lowerCAmelCase=None					,	__lowerCAmelCase=None					,	):
              UpperCamelCase__			       =     parent
              UpperCamelCase__			       =     batch_size
              UpperCamelCase__			       =     seq_length
              UpperCamelCase__			       =     is_training
              UpperCamelCase__			       =     use_input_lengths
              UpperCamelCase__			       =     use_token_type_ids
              UpperCamelCase__			       =     use_labels
              UpperCamelCase__			       =     gelu_activation
              UpperCamelCase__			       =     sinusoidal_embeddings
              UpperCamelCase__			       =     causal
              UpperCamelCase__			       =     asm
              UpperCamelCase__			       =     n_langs
              UpperCamelCase__			       =     vocab_size
              UpperCamelCase__			       =     n_special
              UpperCamelCase__			       =     hidden_size
              UpperCamelCase__			       =     num_hidden_layers
              UpperCamelCase__			       =     num_attention_heads
              UpperCamelCase__			       =     hidden_dropout_prob
              UpperCamelCase__			       =     attention_probs_dropout_prob
              UpperCamelCase__			       =     max_position_embeddings
              UpperCamelCase__			       =     type_vocab_size
              UpperCamelCase__			       =     type_sequence_label_size
              UpperCamelCase__			       =     initializer_range
              UpperCamelCase__			       =     num_labels
              UpperCamelCase__			       =     num_choices
              UpperCamelCase__			       =     summary_type
              UpperCamelCase__			       =     use_proj
              UpperCamelCase__			       =     scope






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     ids_tensor([self.batch_size, self.seq_length]					,	self.vocab_size					)
              UpperCamelCase__			       =     random_attention_mask([self.batch_size, self.seq_length]					)

              UpperCamelCase__			       =     None
              if self.use_input_lengths:
                     UpperCamelCase__			       =     (
                         ids_tensor([self.batch_size]					,	vocab_size=2					) + self.seq_length - 2
                     )  # small variation of seq_length

              UpperCamelCase__			       =     None
              if self.use_token_type_ids:
                     UpperCamelCase__			       =     ids_tensor([self.batch_size, self.seq_length]					,	self.n_langs					)

              UpperCamelCase__			       =     None
              UpperCamelCase__			       =     None
              UpperCamelCase__			       =     None
              if self.use_labels:
                     UpperCamelCase__			       =     ids_tensor([self.batch_size]					,	self.type_sequence_label_size					)
                     UpperCamelCase__			       =     ids_tensor([self.batch_size, self.seq_length]					,	self.num_labels					)
                     UpperCamelCase__			       =     ids_tensor([self.batch_size]					,	2					).float()
                     UpperCamelCase__			       =     ids_tensor([self.batch_size]					,	self.num_choices					)

              UpperCamelCase__			       =     self.get_config()

              return (
                  config,
                  input_ids,
                  token_type_ids,
                  input_lengths,
                  sequence_labels,
                  token_labels,
                  is_impossible_labels,
                  choice_labels,
                  input_mask,
              )






       def 			_lowerCamelCase						(  self					):
              return FlaubertConfig(
                  vocab_size=self.vocab_size					,	n_special=self.n_special					,	emb_dim=self.hidden_size					,	n_layers=self.num_hidden_layers					,	n_heads=self.num_attention_heads					,	dropout=self.hidden_dropout_prob					,	attention_dropout=self.attention_probs_dropout_prob					,	gelu_activation=self.gelu_activation					,	sinusoidal_embeddings=self.sinusoidal_embeddings					,	asm=self.asm					,	causal=self.causal					,	n_langs=self.n_langs					,	max_position_embeddings=self.max_position_embeddings					,	initializer_range=self.initializer_range					,	summary_type=self.summary_type					,	use_proj=self.use_proj					,	)






       def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	):
              UpperCamelCase__			       =     FlaubertModel(config=_lowercase					)
              model.to(_lowercase					)
              model.eval()
              UpperCamelCase__			       =     model(_lowercase					,	lengths=_lowercase					,	langs=_lowercase					)
              UpperCamelCase__			       =     model(_lowercase					,	langs=_lowercase					)
              UpperCamelCase__			       =     model(_lowercase					)
              self.parent.assertEqual(result.last_hidden_state.shape					,	(self.batch_size, self.seq_length, self.hidden_size)					)






       def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	):
              UpperCamelCase__			       =     FlaubertWithLMHeadModel(_lowercase					)
              model.to(_lowercase					)
              model.eval()

              UpperCamelCase__			       =     model(_lowercase					,	token_type_ids=_lowercase					,	labels=_lowercase					)
              self.parent.assertEqual(result.loss.shape					,	()					)
              self.parent.assertEqual(result.logits.shape					,	(self.batch_size, self.seq_length, self.vocab_size)					)






       def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	):
              UpperCamelCase__			       =     FlaubertForQuestionAnsweringSimple(_lowercase					)
              model.to(_lowercase					)
              model.eval()

              UpperCamelCase__			       =     model(_lowercase					)

              UpperCamelCase__			       =     model(_lowercase					,	start_positions=_lowercase					,	end_positions=_lowercase					)
              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 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	):
              UpperCamelCase__			       =     FlaubertForQuestionAnswering(_lowercase					)
              model.to(_lowercase					)
              model.eval()

              UpperCamelCase__			       =     model(_lowercase					)

              UpperCamelCase__			       =     model(
                  _lowercase					,	start_positions=_lowercase					,	end_positions=_lowercase					,	cls_index=_lowercase					,	is_impossible=_lowercase					,	p_mask=_lowercase					,	)

              UpperCamelCase__			       =     model(
                  _lowercase					,	start_positions=_lowercase					,	end_positions=_lowercase					,	cls_index=_lowercase					,	is_impossible=_lowercase					,	)

              ((UpperCamelCase__    )       ,           )			       =     result_with_labels.to_tuple()

              UpperCamelCase__			       =     model(_lowercase					,	start_positions=_lowercase					,	end_positions=_lowercase					)

              ((UpperCamelCase__    )       ,           )			       =     result_with_labels.to_tuple()

              self.parent.assertEqual(result_with_labels.loss.shape					,	()					)
              self.parent.assertEqual(result.start_top_log_probs.shape					,	(self.batch_size, model.config.start_n_top)					)
              self.parent.assertEqual(result.start_top_index.shape					,	(self.batch_size, model.config.start_n_top)					)
              self.parent.assertEqual(
                  result.end_top_log_probs.shape					,	(self.batch_size, model.config.start_n_top * model.config.end_n_top)					)
              self.parent.assertEqual(
                  result.end_top_index.shape					,	(self.batch_size, model.config.start_n_top * model.config.end_n_top)					)
              self.parent.assertEqual(result.cls_logits.shape					,	(self.batch_size,)					)






       def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	):
              UpperCamelCase__			       =     FlaubertForSequenceClassification(_lowercase					)
              model.to(_lowercase					)
              model.eval()

              UpperCamelCase__			       =     model(_lowercase					)
              UpperCamelCase__			       =     model(_lowercase					,	labels=_lowercase					)

              self.parent.assertEqual(result.loss.shape					,	()					)
              self.parent.assertEqual(result.logits.shape					,	(self.batch_size, self.type_sequence_label_size)					)






       def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	):
              UpperCamelCase__			       =     self.num_labels
              UpperCamelCase__			       =     FlaubertForTokenClassification(_lowercase					)
              model.to(_lowercase					)
              model.eval()

              UpperCamelCase__			       =     model(_lowercase					,	attention_mask=_lowercase					,	labels=_lowercase					)
              self.parent.assertEqual(result.logits.shape					,	(self.batch_size, self.seq_length, self.num_labels)					)






       def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	):
              UpperCamelCase__			       =     self.num_choices
              UpperCamelCase__			       =     FlaubertForMultipleChoice(config=_lowercase					)
              model.to(_lowercase					)
              model.eval()
              UpperCamelCase__			       =     input_ids.unsqueeze(1					).expand(-1					,	self.num_choices					,	-1					).contiguous()
              UpperCamelCase__			       =     token_type_ids.unsqueeze(1					).expand(-1					,	self.num_choices					,	-1					).contiguous()
              UpperCamelCase__			       =     input_mask.unsqueeze(1					).expand(-1					,	self.num_choices					,	-1					).contiguous()
              UpperCamelCase__			       =     model(
                  _lowercase					,	attention_mask=_lowercase					,	token_type_ids=_lowercase					,	labels=_lowercase					,	)
              self.parent.assertEqual(result.logits.shape					,	(self.batch_size, self.num_choices)					)







       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     self.prepare_config_and_inputs()
              (
                  (
                  UpperCamelCase__    
              )       ,       (
                  UpperCamelCase__    
              )       ,       (
                  UpperCamelCase__    
              )       ,       (
                  UpperCamelCase__    
              )       ,       (
                  UpperCamelCase__    
              )       ,       (
                  UpperCamelCase__    
              )       ,       (
                  UpperCamelCase__    
              )       ,       (
                  UpperCamelCase__    
              )       ,       (
                  UpperCamelCase__    
              )       ,           
              )			       =     config_and_inputs
              UpperCamelCase__			       =     {
                  """input_ids""": input_ids,
                  """token_type_ids""": token_type_ids,
                  """lengths""": input_lengths,
                  """attention_mask""": input_mask,
              }
              return config, inputs_dict






@require_torch
class     __SCREAMING_SNAKE_CASE   (       _lowerCAmelCase				,						_lowerCAmelCase				,						unittest.TestCase    ):
       snake_case       : List[Any]         =    (
           (
               FlaubertModel,
               FlaubertWithLMHeadModel,
               FlaubertForQuestionAnswering,
               FlaubertForQuestionAnsweringSimple,
               FlaubertForSequenceClassification,
               FlaubertForTokenClassification,
               FlaubertForMultipleChoice,
           )
           if is_torch_available()
           else ()
       )
       snake_case       : str         =    (
           {
               "feature-extraction": FlaubertModel,
               "fill-mask": FlaubertWithLMHeadModel,
               "question-answering": FlaubertForQuestionAnsweringSimple,
               "text-classification": FlaubertForSequenceClassification,
               "token-classification": FlaubertForTokenClassification,
               "zero-shot": FlaubertForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )






       def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					):
              if (
                  pipeline_test_casse_name == "QAPipelineTests"
                  and tokenizer_name is not None
                  and not tokenizer_name.endswith("""Fast"""					)
              ):
                     # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
                     # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
                     # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
                     return True

              return False






       def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase=False					):
              UpperCamelCase__			       =     super()._prepare_for_class(_lowercase					,	_lowercase					,	return_labels=_lowercase					)

              if return_labels:
                     if model_class.__name__ == "FlaubertForQuestionAnswering":
                            UpperCamelCase__			       =     torch.zeros(
                                self.model_tester.batch_size					,	dtype=torch.long					,	device=_lowercase					)
                            UpperCamelCase__			       =     torch.zeros(
                                self.model_tester.batch_size					,	dtype=torch.long					,	device=_lowercase					)

              return inputs_dict






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     FlaubertModelTester(self					)
              UpperCamelCase__			       =     ConfigTester(self					,	config_class=_lowercase					,	emb_dim=37					)






       def 			_lowerCamelCase						(  self					):
              self.config_tester.run_common_tests()






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_flaubert_model(*_lowercase					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_flaubert_lm_head(*_lowercase					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_flaubert_simple_qa(*_lowercase					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_flaubert_qa(*_lowercase					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_flaubert_sequence_classif(*_lowercase					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_flaubert_token_classif(*_lowercase					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_flaubert_multiple_choice(*_lowercase					)






       @slow
       def 			_lowerCamelCase						(  self					):
              for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                     UpperCamelCase__			       =     FlaubertModel.from_pretrained(_lowercase					)
                     self.assertIsNotNone(_lowercase					)







       @slow
       @require_torch_gpu
       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__       ,       UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs_for_common()
              for model_class in self.all_model_classes:
                     # FlauBertForMultipleChoice behaves incorrectly in JIT environments.
                     if model_class == FlaubertForMultipleChoice:
                            return

                     UpperCamelCase__			       =     True
                     UpperCamelCase__			       =     model_class(config=_lowercase					)

                     UpperCamelCase__			       =     self._prepare_for_class(_lowercase					,	_lowercase					)
                     UpperCamelCase__			       =     torch.jit.trace(
                         _lowercase					,	(inputs_dict["""input_ids"""].to("""cpu"""					), inputs_dict["""attention_mask"""].to("""cpu"""					))					)

                     with tempfile.TemporaryDirectory() as tmp:
                            torch.jit.save(_lowercase					,	os.path.join(_lowercase					,	"""traced_model.pt"""					)					)
                            UpperCamelCase__			       =     torch.jit.load(os.path.join(_lowercase					,	"""traced_model.pt"""					)					,	map_location=_lowercase					)
                            loaded(inputs_dict["""input_ids"""].to(_lowercase					)					,	inputs_dict["""attention_mask"""].to(_lowercase					)					)






@require_torch
class     __SCREAMING_SNAKE_CASE   (       unittest.TestCase    ):






       @slow
       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     FlaubertModel.from_pretrained("""flaubert/flaubert_base_cased"""					)
              UpperCamelCase__			       =     torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]					)
              with torch.no_grad():
                     UpperCamelCase__			       =     model(_lowercase					)[0]
              UpperCamelCase__			       =     torch.Size((1, 11, 768)					)
              self.assertEqual(output.shape					,	_lowercase					)
              UpperCamelCase__			       =     torch.tensor(
                  [[[-2.6251, -1.4298, -0.0227], [-2.8510, -1.6387, 0.2258], [-2.8114, -1.1832, -0.3066]]]					)

              self.assertTrue(torch.allclose(output[:, :3, :3]					,	_lowercase					,	atol=1E-4					)					)











from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


UpperCamelCase__     =	[
    "openmmlab/upernet-convnext-tiny",
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
UpperCamelCase__     =	"UperNetConfig"

class     __SCREAMING_SNAKE_CASE   (       nn.Module    ):






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase = 0					,	__lowerCAmelCase = False					,	__lowerCAmelCase = 1					,	):
														super().__init__()
														UpperCamelCase__			       =     nn.Convad(
														    in_channels=__lowerCAmelCase					,	out_channels=__lowerCAmelCase					,	kernel_size=__lowerCAmelCase					,	padding=__lowerCAmelCase					,	bias=__lowerCAmelCase					,	dilation=__lowerCAmelCase					,	)
														UpperCamelCase__			       =     nn.BatchNormad(__lowerCAmelCase					)
														UpperCamelCase__			       =     nn.ReLU()







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     self.conv(__lowerCAmelCase					)
														UpperCamelCase__			       =     self.batch_norm(__lowerCAmelCase					)
														UpperCamelCase__			       =     self.activation(__lowerCAmelCase					)

														return output






class     __SCREAMING_SNAKE_CASE   (       nn.Module    ):






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					):
														super().__init__()
														UpperCamelCase__			       =     [
														    nn.AdaptiveAvgPoolad(__lowerCAmelCase					),
														    UperNetConvModule(__lowerCAmelCase					,	__lowerCAmelCase					,	kernel_size=1					),
														]
														for i, layer in enumerate(self.layers					):
																					self.add_module(str(__lowerCAmelCase					)					,	__lowerCAmelCase					)







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     input
														for layer in self.layers:
																					UpperCamelCase__			       =     layer(__lowerCAmelCase					)
														return hidden_state






class     __SCREAMING_SNAKE_CASE   (       nn.Module    ):






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					):
														super().__init__()
														UpperCamelCase__			       =     pool_scales
														UpperCamelCase__			       =     align_corners
														UpperCamelCase__			       =     in_channels
														UpperCamelCase__			       =     channels
														UpperCamelCase__			       =     []
														for i, pool_scale in enumerate(__lowerCAmelCase					):
																					UpperCamelCase__			       =     UperNetPyramidPoolingBlock(pool_scale=__lowerCAmelCase					,	in_channels=__lowerCAmelCase					,	channels=__lowerCAmelCase					)
																					self.blocks.append(__lowerCAmelCase					)
																					self.add_module(str(__lowerCAmelCase					)					,	__lowerCAmelCase					)







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     []
														for ppm in self.blocks:
																					UpperCamelCase__			       =     ppm(__lowerCAmelCase					)
																					UpperCamelCase__			       =     nn.functional.interpolate(
																					    __lowerCAmelCase					,	size=x.size()[2:]					,	mode="""bilinear"""					,	align_corners=self.align_corners					)
																					ppm_outs.append(__lowerCAmelCase					)
														return ppm_outs






class     __SCREAMING_SNAKE_CASE   (       nn.Module    ):






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					):
														super().__init__()

														UpperCamelCase__			       =     config
														UpperCamelCase__			       =     config.pool_scales  # e.g. (1, 2, 3, 6)
														UpperCamelCase__			       =     in_channels
														UpperCamelCase__			       =     config.hidden_size
														UpperCamelCase__			       =     False
														UpperCamelCase__			       =     nn.Convad(self.channels					,	config.num_labels					,	kernel_size=1					)

														# PSP Module
														UpperCamelCase__			       =     UperNetPyramidPoolingModule(
														    self.pool_scales					,	self.in_channels[-1]					,	self.channels					,	align_corners=self.align_corners					,	)
														UpperCamelCase__			       =     UperNetConvModule(
														    self.in_channels[-1] + len(self.pool_scales					) * self.channels					,	self.channels					,	kernel_size=3					,	padding=1					,	)
														# FPN Module
														UpperCamelCase__			       =     nn.ModuleList()
														UpperCamelCase__			       =     nn.ModuleList()
														for in_channels in self.in_channels[:-1]:  # skip the top layer
																					UpperCamelCase__			       =     UperNetConvModule(__lowerCAmelCase					,	self.channels					,	kernel_size=1					)
																					UpperCamelCase__			       =     UperNetConvModule(self.channels					,	self.channels					,	kernel_size=3					,	padding=1					)
																					self.lateral_convs.append(__lowerCAmelCase					)
																					self.fpn_convs.append(__lowerCAmelCase					)

														UpperCamelCase__			       =     UperNetConvModule(
														    len(self.in_channels					) * self.channels					,	self.channels					,	kernel_size=3					,	padding=1					,	)






							def 			_lowerCamelCase						(  self					):
														self.apply(self._init_weights					)






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														if isinstance(__lowerCAmelCase					,	nn.Convad					):
																					module.weight.data.normal_(mean=0.0					,	std=self.config.initializer_range					)
																					if module.bias is not None:
																												module.bias.data.zero_()






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     inputs[-1]
														UpperCamelCase__			       =     [x]
														psp_outs.extend(self.psp_modules(__lowerCAmelCase					)					)
														UpperCamelCase__			       =     torch.cat(__lowerCAmelCase					,	dim=1					)
														UpperCamelCase__			       =     self.bottleneck(__lowerCAmelCase					)

														return output







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														# build laterals
														UpperCamelCase__			       =     [lateral_conv(encoder_hidden_states[i]					) for i, lateral_conv in enumerate(self.lateral_convs					)]

														laterals.append(self.psp_forward(__lowerCAmelCase					)					)

														# build top-down path
														UpperCamelCase__			       =     len(__lowerCAmelCase					)
														for i in range(used_backbone_levels - 1					,	0					,	-1					):
																					UpperCamelCase__			       =     laterals[i - 1].shape[2:]
																					UpperCamelCase__			       =     laterals[i - 1] + nn.functional.interpolate(
																					    laterals[i]					,	size=__lowerCAmelCase					,	mode="""bilinear"""					,	align_corners=self.align_corners					)

														# build outputs
														UpperCamelCase__			       =     [self.fpn_convs[i](laterals[i]					) for i in range(used_backbone_levels - 1					)]
														# append psp feature
														fpn_outs.append(laterals[-1]					)

														for i in range(used_backbone_levels - 1					,	0					,	-1					):
																					UpperCamelCase__			       =     nn.functional.interpolate(
																					    fpn_outs[i]					,	size=fpn_outs[0].shape[2:]					,	mode="""bilinear"""					,	align_corners=self.align_corners					)
														UpperCamelCase__			       =     torch.cat(__lowerCAmelCase					,	dim=1					)
														UpperCamelCase__			       =     self.fpn_bottleneck(__lowerCAmelCase					)
														UpperCamelCase__			       =     self.classifier(__lowerCAmelCase					)

														return output






class     __SCREAMING_SNAKE_CASE   (       nn.Module    ):






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase = 2					,	__lowerCAmelCase = 3					,	__lowerCAmelCase = 1					):
														super().__init__()

														UpperCamelCase__			       =     config
														UpperCamelCase__			       =     config.auxiliary_in_channels
														UpperCamelCase__			       =     config.auxiliary_channels
														UpperCamelCase__			       =     config.auxiliary_num_convs
														UpperCamelCase__			       =     config.auxiliary_concat_input
														UpperCamelCase__			       =     in_index

														UpperCamelCase__			       =     (kernel_size // 2) * dilation
														UpperCamelCase__			       =     []
														convs.append(
														    UperNetConvModule(
														        self.in_channels					,	self.channels					,	kernel_size=__lowerCAmelCase					,	padding=__lowerCAmelCase					,	dilation=__lowerCAmelCase					)					)
														for i in range(self.num_convs - 1					):
																					convs.append(
																					    UperNetConvModule(
																					        self.channels					,	self.channels					,	kernel_size=__lowerCAmelCase					,	padding=__lowerCAmelCase					,	dilation=__lowerCAmelCase					)					)
														if self.num_convs == 0:
																					UpperCamelCase__			       =     nn.Identity()
														else:
																					UpperCamelCase__			       =     nn.Sequential(*__lowerCAmelCase					)
														if self.concat_input:
																					UpperCamelCase__			       =     UperNetConvModule(
																					    self.in_channels + self.channels					,	self.channels					,	kernel_size=__lowerCAmelCase					,	padding=kernel_size // 2					)

														UpperCamelCase__			       =     nn.Convad(self.channels					,	config.num_labels					,	kernel_size=1					)






							def 			_lowerCamelCase						(  self					):
														self.apply(self._init_weights					)






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														if isinstance(__lowerCAmelCase					,	nn.Convad					):
																					module.weight.data.normal_(mean=0.0					,	std=self.config.initializer_range					)
																					if module.bias is not None:
																												module.bias.data.zero_()







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														# just take the relevant feature maps
														UpperCamelCase__			       =     encoder_hidden_states[self.in_index]
														UpperCamelCase__			       =     self.convs(__lowerCAmelCase					)
														if self.concat_input:
																					UpperCamelCase__			       =     self.conv_cat(torch.cat([hidden_states, output]					,	dim=1					)					)
														UpperCamelCase__			       =     self.classifier(__lowerCAmelCase					)
														return output






class     __SCREAMING_SNAKE_CASE   (       _a    ):
							snake_case       : Any         =    UperNetConfig
							snake_case       : List[Any]         =    """pixel_values"""
							snake_case       : Optional[Any]         =    True






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														if isinstance(__lowerCAmelCase					,	__lowerCAmelCase					):
																					module.backbone.init_weights()
																					module.decode_head.init_weights()
																					module.auxiliary_head.init_weights()






							def 			_lowerCamelCase						(  self					):
														self.backbone.init_weights()
														self.decode_head.init_weights()
														self.auxiliary_head.init_weights()







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase=False					):
														if isinstance(__lowerCAmelCase					,	__lowerCAmelCase					):
																					UpperCamelCase__			       =     value






UpperCamelCase__     =	r"\n    Parameters:\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n        config ([`UperNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

UpperCamelCase__     =	r"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using\n            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.\n        output_attentions (`bool`, *optional*):\n            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See\n            `attentions` under returned tensors for more detail.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers of the backbone. See `hidden_states` under\n            returned tensors for more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"

@add_start_docstrings(
    """UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes."""				,						_a				,						)
class     __SCREAMING_SNAKE_CASE   (       _a    ):






							def __init__(  self					,	__lowerCAmelCase					):
														super().__init__(__lowerCAmelCase					)

														UpperCamelCase__			       =     AutoBackbone.from_config(config.backbone_config					)

														# Semantic segmentation head(s)
														UpperCamelCase__			       =     UperNetHead(__lowerCAmelCase					,	in_channels=self.backbone.channels					)
														UpperCamelCase__			       =     UperNetFCNHead(__lowerCAmelCase					) if config.use_auxiliary_head else None

														# Initialize weights and apply final processing
														self.post_init()







							@add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format("""batch_size, sequence_length"""					)					)
							@replace_return_docstrings(output_type=__lowerCAmelCase					,	config_class=_CONFIG_FOR_DOC					)
							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase = None					,	__lowerCAmelCase = None					,	__lowerCAmelCase = None					,	__lowerCAmelCase = None					,	__lowerCAmelCase = None					,	):
														UpperCamelCase__			       =     return_dict if return_dict is not None else self.config.use_return_dict
														UpperCamelCase__			       =     (
														    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
														)
														UpperCamelCase__			       =     output_attentions if output_attentions is not None else self.config.output_attentions

														UpperCamelCase__			       =     self.backbone.forward_with_filtered_kwargs(
														    __lowerCAmelCase					,	output_hidden_states=__lowerCAmelCase					,	output_attentions=__lowerCAmelCase					)
														UpperCamelCase__			       =     outputs.feature_maps

														UpperCamelCase__			       =     self.decode_head(__lowerCAmelCase					)
														UpperCamelCase__			       =     nn.functional.interpolate(__lowerCAmelCase					,	size=pixel_values.shape[2:]					,	mode="""bilinear"""					,	align_corners=__lowerCAmelCase					)

														UpperCamelCase__			       =     None
														if self.auxiliary_head is not None:
																					UpperCamelCase__			       =     self.auxiliary_head(__lowerCAmelCase					)
																					UpperCamelCase__			       =     nn.functional.interpolate(
																					    __lowerCAmelCase					,	size=pixel_values.shape[2:]					,	mode="""bilinear"""					,	align_corners=__lowerCAmelCase					)

														UpperCamelCase__			       =     None
														if labels is not None:
																					if self.config.num_labels == 1:
																												raise ValueError("""The number of labels should be greater than one"""					)
																					else:
																												# compute weighted loss
																												UpperCamelCase__			       =     CrossEntropyLoss(ignore_index=self.config.loss_ignore_index					)
																												UpperCamelCase__			       =     loss_fct(__lowerCAmelCase					,	__lowerCAmelCase					)
																												UpperCamelCase__			       =     loss_fct(__lowerCAmelCase					,	__lowerCAmelCase					)
																												UpperCamelCase__			       =     main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

														if not return_dict:
																					if output_hidden_states:
																												UpperCamelCase__			       =     (logits,) + outputs[1:]
																					else:
																												UpperCamelCase__			       =     (logits,) + outputs[2:]
																					return ((loss,) + output) if loss is not None else output

														return SemanticSegmenterOutput(
														    loss=__lowerCAmelCase					,	logits=__lowerCAmelCase					,	hidden_states=outputs.hidden_states					,	attentions=outputs.attentions					,	)




